{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":30892,"databundleVersionId":2715462,"sourceType":"competition"},{"sourceId":6974612,"sourceType":"datasetVersion","datasetId":4007626}],"dockerImageVersionId":30580,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install segmentation-models-pytorch","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:26:51.556915Z","iopub.execute_input":"2023-11-16T06:26:51.557177Z","iopub.status.idle":"2023-11-16T06:27:11.641856Z","shell.execute_reply.started":"2023-11-16T06:26:51.557151Z","shell.execute_reply":"2023-11-16T06:27:11.640766Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport pandas as pd\nimport numpy as np\nimport cv2\nfrom torchvision.io import read_image\nimport matplotlib.pyplot as plt\nfrom torch.utils.data import Dataset, random_split, DataLoader\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nfrom torchvision.transforms import ToTensor\nfrom PIL import Image\nimport os\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torchvision \nfrom torchvision import transforms\nfrom torchinfo import summary\nimport timm","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:27:11.644007Z","iopub.execute_input":"2023-11-16T06:27:11.644308Z","iopub.status.idle":"2023-11-16T06:27:18.871966Z","shell.execute_reply.started":"2023-11-16T06:27:11.644277Z","shell.execute_reply":"2023-11-16T06:27:18.871004Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!nvidia-smi","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:27:18.873220Z","iopub.execute_input":"2023-11-16T06:27:18.873890Z","iopub.status.idle":"2023-11-16T06:27:19.921795Z","shell.execute_reply.started":"2023-11-16T06:27:18.873860Z","shell.execute_reply":"2023-11-16T06:27:19.920721Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomImageDataset(Dataset):\n    def __init__(self, img_dir, label_dir, resize=None, transform=None):\n        self.img_dir = img_dir\n        self.label_dir = label_dir\n        self.resize = resize\n        self.transform = transform\n        self.images = os.listdir(self.img_dir)\n\n    def __len__(self):\n        return len(self.images)\n    \n    def read_mask(self, mask_path):\n        image = cv2.imread(mask_path)\n        image = cv2.resize(image, self.resize)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)\n\n        lower1 = np.array([0, 100, 20])\n        upper1 = np.array([10, 255, 255])\n\n        lower2 = np.array([160,100,20])\n        upper2 = np.array([179,255,255])\n        lower_mask = cv2.inRange(image, lower1, upper1)\n        upper_mask = cv2.inRange(image, lower2, upper2)\n        \n        red_mask = lower_mask + upper_mask;\n        red_mask[red_mask != 0] = 1\n\n        green_mask = cv2.inRange(image, (36, 25, 25), (70, 255, 255))\n        green_mask[green_mask != 0] = 2\n\n        full_mask = cv2.bitwise_or(red_mask, green_mask)\n        full_mask = np.expand_dims(full_mask, axis=-1) \n        full_mask = full_mask.astype(np.uint8)\n        \n        return full_mask\n\n    def __getitem__(self, idx):\n        img_path = os.path.join(self.img_dir, self.images[idx])\n        label_path = os.path.join(self.label_dir, self.images[idx])\n        image = cv2.imread(img_path)  # Đọc ảnh dưới dạng BGR\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convert sang RGB\n        label = self.read_mask(label_path)  \n        image = cv2.resize(image, self.resize)\n        if self.transform:\n            image = self.transform(image)\n            \n        return image, label\n\n    def show_image(self, idx):\n        img_path = os.path.join(self.img_dir, self.images[idx])\n        label_path = os.path.join(self.label_dir, self.images[idx])\n        image = plt.imread(img_path)\n        label = plt.imread(label_path)\n        fig, axs = plt.subplots(1, 2, figsize=(10, 5))\n        axs[0].imshow(image)\n        axs[0].set_title('Image')\n        axs[1].imshow(label)\n        axs[1].set_title('Label')\n        plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:27:19.925055Z","iopub.execute_input":"2023-11-16T06:27:19.925514Z","iopub.status.idle":"2023-11-16T06:27:19.941371Z","shell.execute_reply.started":"2023-11-16T06:27:19.925472Z","shell.execute_reply":"2023-11-16T06:27:19.940584Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image_path = []\nTRAIN_DIR = '/kaggle/input/bkai-igh-neopolyp/train/train'\nfor root, dirs, files in os.walk(TRAIN_DIR):\n    for file in files:\n        path = os.path.join(root,file)\n        image_path.append(path)\n        \nlen(image_path)","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:27:19.942375Z","iopub.execute_input":"2023-11-16T06:27:19.942699Z","iopub.status.idle":"2023-11-16T06:27:20.558019Z","shell.execute_reply.started":"2023-11-16T06:27:19.942666Z","shell.execute_reply":"2023-11-16T06:27:20.557089Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"mask_path = []\nTRAIN_MASK_DIR = '/kaggle/input/bkai-igh-neopolyp/train_gt/train_gt'\nfor root, dirs, files in os.walk(TRAIN_MASK_DIR):\n    for file in files:\n        path = os.path.join(root,file)\n        mask_path.append(path)\n        \nlen(mask_path)","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:27:20.559105Z","iopub.execute_input":"2023-11-16T06:27:20.559415Z","iopub.status.idle":"2023-11-16T06:27:21.098308Z","shell.execute_reply.started":"2023-11-16T06:27:20.559388Z","shell.execute_reply":"2023-11-16T06:27:21.097348Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset = CustomImageDataset(img_dir= TRAIN_DIR,\n                             label_dir= TRAIN_MASK_DIR,\n                             resize= (256,256),\n                             transform = None)","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:27:21.099301Z","iopub.execute_input":"2023-11-16T06:27:21.099596Z","iopub.status.idle":"2023-11-16T06:27:21.104812Z","shell.execute_reply.started":"2023-11-16T06:27:21.099572Z","shell.execute_reply":"2023-11-16T06:27:21.103926Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import segmentation_models_pytorch as smp\n\nmodel = smp.UnetPlusPlus(\n    encoder_name=\"resnet34\",        \n    encoder_weights=\"imagenet\",     \n    in_channels=3,                  \n    classes=3     \n)","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:27:21.105983Z","iopub.execute_input":"2023-11-16T06:27:21.106260Z","iopub.status.idle":"2023-11-16T06:27:23.394046Z","shell.execute_reply.started":"2023-11-16T06:27:21.106237Z","shell.execute_reply":"2023-11-16T06:27:23.393255Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"batch_size = 8\nimages_data = []\nlabels_data = []\nfor x,y in dataset:\n    images_data.append(x)\n    labels_data.append(y)","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:27:23.395527Z","iopub.execute_input":"2023-11-16T06:27:23.395975Z","iopub.status.idle":"2023-11-16T06:28:02.164952Z","shell.execute_reply.started":"2023-11-16T06:27:23.395937Z","shell.execute_reply":"2023-11-16T06:28:02.163853Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomDataset(CustomImageDataset):\n    def __init__(self, data, targets, transform=None):\n        self.data = data\n        self.targets = targets\n        self.transform = transform\n\n    def __getitem__(self, index):\n        image = self.data[index]\n        label = self.targets[index]\n        if self.transform:\n            transformed = self.transform(image=image, mask=label)\n            image = transformed['image'].float()\n            label = transformed['mask'].float()\n            label = label.permute(2, 0, 1)\n        return image, label\n    \n    def __len__(self):\n        return len(self.data)\n\n    \ntrain_transform = A.Compose([\n    A.HorizontalFlip(p=0.5),\n    A.VerticalFlip(p=0.5),\n    A.RandomGamma (gamma_limit=(70, 130), eps=None, always_apply=False, p=0.2),\n    A.RGBShift(p=0.3, r_shift_limit=10, g_shift_limit=10, b_shift_limit=10),\n    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n    ToTensorV2(),\n])\n\nval_transform = A.Compose([\n    A.Normalize(mean=(0.485, 0.456, 0.406),std=(0.229, 0.224, 0.225)),\n    ToTensorV2(),\n])\n\ntrain_size = int(0.9 * len(images_data))\nval_size = len(images_data) - train_size\ntrain_dataset = CustomDataset(images_data[:train_size], labels_data[:train_size], transform=train_transform)\nval_dataset = CustomDataset(images_data[train_size:], labels_data[train_size:], transform=val_transform)\n\ntrain_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=True)\nprint(len(train_dataset))\nprint(len(val_dataset))","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:28:02.169099Z","iopub.execute_input":"2023-11-16T06:28:02.169921Z","iopub.status.idle":"2023-11-16T06:28:02.184289Z","shell.execute_reply.started":"2023-11-16T06:28:02.169889Z","shell.execute_reply":"2023-11-16T06:28:02.183392Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image,label = train_dataset[2]\n\nlabel_array = label.permute(1, 2, 0).numpy()\nimage_array = image.permute(1, 2, 0).numpy()\n\nfig, axs = plt.subplots(1, 2, figsize=(10, 5))\n\naxs[0].imshow(image_array)\naxs[0].set_title('Image')\naxs[0].axis('off')  \n\naxs[1].imshow(label_array)\naxs[1].set_title('Label')\naxs[1].axis('off')  \n\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:28:02.185422Z","iopub.execute_input":"2023-11-16T06:28:02.185782Z","iopub.status.idle":"2023-11-16T06:28:02.568815Z","shell.execute_reply.started":"2023-11-16T06:28:02.185755Z","shell.execute_reply":"2023-11-16T06:28:02.567952Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"learning_rate = 0.0001\noptimizer = optim.Adam(model.parameters(), lr=learning_rate)","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:28:02.569774Z","iopub.execute_input":"2023-11-16T06:28:02.570027Z","iopub.status.idle":"2023-11-16T06:28:02.576208Z","shell.execute_reply.started":"2023-11-16T06:28:02.570003Z","shell.execute_reply":"2023-11-16T06:28:02.575304Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"color_dict= {0: (0, 0, 0),\n             1: (255, 0, 0),\n             2: (0, 255, 0)}\ndef mask_to_rgb(mask, color_dict):\n    output = np.zeros((mask.shape[0], mask.shape[1], 3))\n\n    for k in color_dict.keys():\n        output[mask==k] = color_dict[k]\n\n    return np.uint8(output)    ","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:28:02.577354Z","iopub.execute_input":"2023-11-16T06:28:02.577644Z","iopub.status.idle":"2023-11-16T06:28:02.585535Z","shell.execute_reply.started":"2023-11-16T06:28:02.577620Z","shell.execute_reply":"2023-11-16T06:28:02.584753Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install wandb\nimport wandb\n!wandb login '7c74de70e64d4ae9ebc72e3d26a7847ac9337b3d'","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:28:02.586657Z","iopub.execute_input":"2023-11-16T06:28:02.586921Z","iopub.status.idle":"2023-11-16T06:28:18.062693Z","shell.execute_reply.started":"2023-11-16T06:28:02.586898Z","shell.execute_reply":"2023-11-16T06:28:18.061699Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"wandb.init(\n    project = 'Unet_polyp-Segmentation'\n)","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:28:18.064075Z","iopub.execute_input":"2023-11-16T06:28:18.064427Z","iopub.status.idle":"2023-11-16T06:28:51.290136Z","shell.execute_reply.started":"2023-11-16T06:28:18.064395Z","shell.execute_reply":"2023-11-16T06:28:51.289136Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"num_epochs = 300\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else \"cpu\")\nmodel.to(device)\ncriterion = nn.CrossEntropyLoss()\nbest_val_loss = 999\n\nfor epoch in range(num_epochs):\n    model.train()\n    train_loss = 0\n    for images, labels in train_loader:\n        images = images.to(device)\n        labels = labels.to(device)\n\n        labels = labels.squeeze(dim=1).long()\n        outputs = model(images)\n    \n        loss = criterion(outputs, labels)\n        loss.backward()\n        optimizer.step()\n        optimizer.zero_grad()\n        train_loss += loss.item()\n    model.eval()\n    with torch.no_grad():\n        val_loss = 0\n        for images, labels in val_loader:\n            images = images.to(device)\n            labels = labels.to(device)\n            labels = labels.squeeze(dim=1).long()\n            \n            outputs = model(images)\n\n            val_loss += criterion(outputs.float(),labels.long()).item()\n\n    print(f\"Epoch [{epoch+1}/{num_epochs}], Loss: {val_loss/len(val_loader):.10f}\")\n    if val_loss < best_val_loss:\n        best_val_loss = val_loss\n        checkpoint = { \n            'epoch': epoch,\n            'model': model.state_dict(),\n            'optimizer': optimizer.state_dict(),\n            'loss': val_loss,\n        }\n        save_path = f'colorization_model.pth'\n        torch.save(checkpoint, save_path)\n        print('Save new model')\n    wandb.log({'Val_loss': val_loss/len(val_loader),\n               'Train_loss': train_loss/len(train_loader)\n              })\n\n    label = labels[0].cpu().numpy()\n    label = mask_to_rgb(label,color_dict)\n    outputs[0] = outputs[0].softmax(dim=0)\n    output = outputs[0].cpu().numpy()\n    output = np.argmax(output, axis=0)\n    output = mask_to_rgb(output,color_dict)\n    fig, axs = plt.subplots(1, 2, figsize=(10, 5))\n    axs[0].imshow(label)\n    axs[0].set_title('Label')\n    axs[1].imshow(output)\n    axs[1].set_title('Output')\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:30:28.434272Z","iopub.execute_input":"2023-11-16T06:30:28.435181Z","iopub.status.idle":"2023-11-16T06:31:25.755954Z","shell.execute_reply.started":"2023-11-16T06:30:28.435144Z","shell.execute_reply":"2023-11-16T06:31:25.754655Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"checkpoint = torch.load('/kaggle/working/colorization_model.pth')\nmodel.load_state_dict(checkpoint['model'])\ndevice = torch.device('cuda' if torch.cuda.is_available() else \"cpu\")\nmodel.to(device)","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:31:32.445242Z","iopub.execute_input":"2023-11-16T06:31:32.446188Z","iopub.status.idle":"2023-11-16T06:31:32.771720Z","shell.execute_reply.started":"2023-11-16T06:31:32.446145Z","shell.execute_reply":"2023-11-16T06:31:32.770820Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!mkdir prediction","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:31:33.375225Z","iopub.execute_input":"2023-11-16T06:31:33.376111Z","iopub.status.idle":"2023-11-16T06:31:34.406909Z","shell.execute_reply.started":"2023-11-16T06:31:33.376074Z","shell.execute_reply":"2023-11-16T06:31:34.405476Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainsize = 256\nmodel.eval()\nfor i in os.listdir(\"/kaggle/input/bkai-igh-neopolyp/test/test\"):\n    img_path = os.path.join(\"/kaggle/input/bkai-igh-neopolyp/test/test\", i)\n    ori_img = cv2.imread(img_path)\n    ori_img = cv2.cvtColor(ori_img, cv2.COLOR_BGR2RGB)\n    ori_w = ori_img.shape[0]\n    ori_h = ori_img.shape[1]\n    img = cv2.resize(ori_img, (trainsize, trainsize))\n    transformed = val_transform(image=img)\n    input_img = transformed[\"image\"]\n    input_img = input_img.unsqueeze(0).to(device)\n    with torch.no_grad():\n        output_mask = model.forward(input_img).squeeze(0).cpu().numpy().transpose(1,2,0)\n    mask = cv2.resize(output_mask, (ori_h, ori_w))\n    mask = np.argmax(mask, axis=2)\n    mask_rgb = mask_to_rgb(mask, color_dict)\n    mask_rgb = cv2.cvtColor(mask_rgb, cv2.COLOR_RGB2BGR)\n    cv2.imwrite(\"prediction/{}\".format(i), mask_rgb) ","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:31:34.489578Z","iopub.execute_input":"2023-11-16T06:31:34.489876Z","iopub.status.idle":"2023-11-16T06:32:00.516948Z","shell.execute_reply.started":"2023-11-16T06:31:34.489845Z","shell.execute_reply":"2023-11-16T06:32:00.515939Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport cv2\nimport os\n\ndef rle_to_string(runs):\n    return ' '.join(str(x) for x in runs)\n\ndef rle_encode_one_mask(mask):\n    pixels = mask.flatten()\n    pixels[pixels > 225] = 255\n    pixels[pixels <= 225] = 0\n    use_padding = False\n    if pixels[0] or pixels[-1]:\n        use_padding = True\n        pixel_padded = np.zeros([len(pixels) + 2], dtype=pixels.dtype)\n        pixel_padded[1:-1] = pixels\n        pixels = pixel_padded\n    rle = np.where(pixels[1:] != pixels[:-1])[0] + 2\n    if use_padding:\n        rle = rle - 1\n    rle[1::2] = rle[1::2] - rle[:-1:2]\n    \n    return rle_to_string(rle)\n\ndef rle2mask(mask_rle, shape=(3,3)):\n    s = mask_rle.split()\n    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n    starts -= 1\n    ends = starts + lengths\n    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)\n    for lo, hi in zip(starts, ends):\n        img[lo:hi] = 1\n    return img.reshape(shape).T\n\ndef mask2string(dir):\n    strings = []\n    ids = []\n    ws, hs = [[] for i in range(2)]\n    for image_id in os.listdir(dir):\n        id = image_id.split('.')[0]\n        path = os.path.join(dir, image_id)\n        print(path)\n        img = cv2.imread(path)[:,:,::-1]\n        h, w = img.shape[0], img.shape[1]\n        for channel in range(2):\n            ws.append(w)\n            hs.append(h)\n            ids.append(f'{id}_{channel}')\n            string = rle_encode_one_mask(img[:,:,channel])\n            strings.append(string)\n    r = {\n        'ids': ids,\n        'strings': strings,\n    }\n    return r\n\n\nMASK_DIR_PATH = '/kaggle/working/prediction'\ndir = MASK_DIR_PATH\nres = mask2string(dir)\ndf = pd.DataFrame(columns=['Id', 'Expected'])\ndf['Id'] = res['ids']\ndf['Expected'] = res['strings']\n\ndf.to_csv(r'output.csv', index=False)\nprint('Done')","metadata":{"execution":{"iopub.status.busy":"2023-11-16T06:32:00.519740Z","iopub.execute_input":"2023-11-16T06:32:00.520025Z","iopub.status.idle":"2023-11-16T06:32:03.808736Z","shell.execute_reply.started":"2023-11-16T06:32:00.519998Z","shell.execute_reply":"2023-11-16T06:32:03.807199Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}