{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","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"}],"dockerImageVersionId":30787,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -qU torchsummary\n!pip install -qU torchgeometry\n!pip install -qU segmentation-models-pytorch","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-11-23T07:08:43.695628Z","iopub.execute_input":"2024-11-23T07:08:43.696252Z","iopub.status.idle":"2024-11-23T07:09:15.156432Z","shell.execute_reply.started":"2024-11-23T07:08:43.696201Z","shell.execute_reply":"2024-11-23T07:09:15.155241Z"}},"outputs":[],"execution_count":null},{"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\nimport segmentation_models_pytorch as smp\nimport wandb\nimport segmentation_models_pytorch as smp","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T07:10:45.665364Z","iopub.execute_input":"2024-11-23T07:10:45.665762Z","iopub.status.idle":"2024-11-23T07:10:45.671435Z","shell.execute_reply.started":"2024-11-23T07:10:45.665721Z","shell.execute_reply":"2024-11-23T07:10:45.670591Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DatasetCustom(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        lower_red1 = np.array([0, 100, 20])\n        upper_red1 = np.array([10, 255, 255])\n        lower_red2 = np.array([160,100,20])\n        upper_red2 = np.array([179,255,255])\n        \n        lower_mask_red = cv2.inRange(image, lower_red1, upper_red1)\n        upper_mask_red = cv2.inRange(image, lower_red2, upper_red2)\n        \n        red_mask = lower_mask_red + upper_mask_red\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)  #  BGR\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convert to 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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T07:10:07.36239Z","iopub.execute_input":"2024-11-23T07:10:07.362661Z","iopub.status.idle":"2024-11-23T07:10:07.371186Z","shell.execute_reply.started":"2024-11-23T07:10:07.362634Z","shell.execute_reply":"2024-11-23T07:10:07.370374Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"images_path = \"/kaggle/input/bkai-igh-neopolyp/train/train/\"\nimage_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":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T07:10:07.372294Z","iopub.execute_input":"2024-11-23T07:10:07.372598Z","iopub.status.idle":"2024-11-23T07:10:08.434004Z","shell.execute_reply.started":"2024-11-23T07:10:07.372573Z","shell.execute_reply":"2024-11-23T07:10:08.433146Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T07:10:08.435872Z","iopub.execute_input":"2024-11-23T07:10:08.436168Z","iopub.status.idle":"2024-11-23T07:10:09.312514Z","shell.execute_reply.started":"2024-11-23T07:10:08.436141Z","shell.execute_reply":"2024-11-23T07:10:09.311614Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dataset = DatasetCustom(img_dir= TRAIN_DIR,\n                             label_dir= TRAIN_MASK_DIR,\n                             resize= (256,256),\n                             transform = None)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T07:10:09.313544Z","iopub.execute_input":"2024-11-23T07:10:09.313895Z","iopub.status.idle":"2024-11-23T07:10:09.318754Z","shell.execute_reply.started":"2024-11-23T07:10:09.313856Z","shell.execute_reply":"2024-11-23T07:10:09.318104Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T07:10:15.128771Z","iopub.execute_input":"2024-11-23T07:10:15.129604Z","iopub.status.idle":"2024-11-23T07:10:45.663694Z","shell.execute_reply.started":"2024-11-23T07:10:15.129562Z","shell.execute_reply":"2024-11-23T07:10:45.66275Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = smp.UnetPlusPlus(\n    encoder_name=\"resnet34\",        \n    encoder_weights=\"imagenet\",     \n    in_channels=3,                  \n    classes=3     \n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T07:10:51.054436Z","iopub.execute_input":"2024-11-23T07:10:51.055303Z","iopub.status.idle":"2024-11-23T07:10:52.1988Z","shell.execute_reply.started":"2024-11-23T07:10:51.055267Z","shell.execute_reply":"2024-11-23T07:10:52.198045Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CustomDataset(Dataset):\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        assert image.shape[:2] == label.shape[:2]\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)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T07:10:54.868321Z","iopub.execute_input":"2024-11-23T07:10:54.868673Z","iopub.status.idle":"2024-11-23T07:10:54.875191Z","shell.execute_reply.started":"2024-11-23T07:10:54.86864Z","shell.execute_reply":"2024-11-23T07:10:54.874092Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_transformation = A.Compose([\n    A.HorizontalFlip(p=0.4),\n    A.VerticalFlip(p=0.4),\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_transformation = A.Compose([\n    A.Normalize(mean=(0.485, 0.456, 0.406),std=(0.229, 0.224, 0.225)),\n    ToTensorV2(),\n])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T07:11:01.619255Z","iopub.execute_input":"2024-11-23T07:11:01.619609Z","iopub.status.idle":"2024-11-23T07:11:01.627428Z","shell.execute_reply.started":"2024-11-23T07:11:01.619576Z","shell.execute_reply":"2024-11-23T07:11:01.626512Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_size = int(0.8 * len(images_data))\nval_size = len(images_data) - train_size\ntrain_dataset = CustomDataset(images_data[:train_size], labels_data[:train_size], transform=train_transformation)\nval_dataset = CustomDataset(images_data[train_size:], labels_data[train_size:], transform=val_transformation)\n\ntrain_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T07:11:17.461853Z","iopub.execute_input":"2024-11-23T07:11:17.462764Z","iopub.status.idle":"2024-11-23T07:11:17.467793Z","shell.execute_reply.started":"2024-11-23T07:11:17.462727Z","shell.execute_reply":"2024-11-23T07:11:17.466898Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"learning_rate = 0.0001\noptimizer = optim.Adam(model.parameters(), lr=learning_rate)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T07:11:24.542293Z","iopub.execute_input":"2024-11-23T07:11:24.54295Z","iopub.status.idle":"2024-11-23T07:11:24.548303Z","shell.execute_reply.started":"2024-11-23T07:11:24.542913Z","shell.execute_reply":"2024-11-23T07:11:24.547343Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T07:11:28.558728Z","iopub.execute_input":"2024-11-23T07:11:28.559393Z","iopub.status.idle":"2024-11-23T07:11:28.564627Z","shell.execute_reply.started":"2024-11-23T07:11:28.559355Z","shell.execute_reply":"2024-11-23T07:11:28.563662Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"wandb.login(\n    key = \"0914f403b30196f64eb9744f9310b0029205c241\",\n)\nwandb.init(\n    project = \"PolypSegment\"\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T07:23:53.613621Z","iopub.execute_input":"2024-11-23T07:23:53.613983Z","iopub.status.idle":"2024-11-23T07:23:56.647164Z","shell.execute_reply.started":"2024-11-23T07:23:53.613949Z","shell.execute_reply":"2024-11-23T07:23:56.646331Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from tqdm import tqdm\nimport time\n\nnum_epochs = 250\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else \"cpu\")\nmodel.to(device)\ncriterion = nn.CrossEntropyLoss()\nbest_val_loss = 999\n\nepoch_bar = tqdm(total=num_epochs, desc='Total Progress')\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'model.pth'\n        torch.save(checkpoint, save_path)\n        \n    epoch_bar.update(1)\n    wandb.log({'Val_loss': val_loss/len(val_loader),'Train_loss': train_loss/len(train_loader)})\nepoch_bar.close()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T07:27:26.100387Z","iopub.execute_input":"2024-11-23T07:27:26.100727Z","iopub.status.idle":"2024-11-23T07:27:28.201349Z","shell.execute_reply.started":"2024-11-23T07:27:26.100696Z","shell.execute_reply":"2024-11-23T07:27:28.200254Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# checkpoint = torch.load('/kaggle/input/unet-model/model.pth')\n# model.load_state_dict(checkpoint['model'])\n# device = torch.device('cuda' if torch.cuda.is_available() else \"cpu\")\n# model.to(device)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !mkdir prediction","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.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, (256, 256))\n    transformed = val_transformation(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":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}