{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np \nimport pandas as pd \nimport os","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:12:46.827890Z","iopub.execute_input":"2023-07-30T15:12:46.828252Z","iopub.status.idle":"2023-07-30T15:12:46.834238Z","shell.execute_reply.started":"2023-07-30T15:12:46.828223Z","shell.execute_reply":"2023-07-30T15:12:46.833255Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nos.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"max_split_size_mb:512\"","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:12:47.066904Z","iopub.execute_input":"2023-07-30T15:12:47.067267Z","iopub.status.idle":"2023-07-30T15:12:47.072308Z","shell.execute_reply.started":"2023-07-30T15:12:47.067236Z","shell.execute_reply":"2023-07-30T15:12:47.070928Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"label_map = np.array([\n    (0, 0, 0),  # background\n    (128, 0, 0),  # aeroplane\n    (0, 128, 0),  # bicycle\n    (128, 128, 0),  # bird\n    (0, 0, 128),  # boat\n    (128, 0, 128),  # bottle\n    (0, 128, 128),  # bus\n    (128, 128, 128),  # car\n    (64, 0, 0),  # cat\n    (192, 0, 0),  # chair\n    (64, 128, 0),  # cow\n    (192, 128, 0),  # dining table\n    (64, 0, 128),  # dog\n    (192, 0, 128),  # horse\n    (64, 128, 128),  # motorbike\n    (192, 128, 128),  # person\n    (0, 64, 0),  # potted plant\n    (128, 64, 0),  # sheep\n    (0, 192, 0),  # sofa\n    (128, 192, 0),  # train\n    (0, 64, 128),  # tv/monitor\n])","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:12:47.586690Z","iopub.execute_input":"2023-07-30T15:12:47.587594Z","iopub.status.idle":"2023-07-30T15:12:47.596359Z","shell.execute_reply.started":"2023-07-30T15:12:47.587550Z","shell.execute_reply":"2023-07-30T15:12:47.594958Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def draw_segmentation_map(outputs):\n    labels = torch.argmax(outputs.squeeze(), dim=0).numpy()\n  \n    # Create 3 Numpy arrays containing zeros.\n    # Later each pixel will be filled with respective red, green, and blue pixels\n    # depending on the predicted class.\n  \n    red_map   = np.zeros_like(labels).astype(np.uint8)\n    green_map = np.zeros_like(labels).astype(np.uint8)\n    blue_map  = np.zeros_like(labels).astype(np.uint8)\n  \n    for label_num in range(0, len(label_map)):\n        index = labels == label_num\n         \n        R, G, B = label_map[label_num]\n  \n        red_map[index]   = R\n        green_map[index] = G\n        blue_map[index]  = B\n  \n    segmentation_map = np.stack([red_map, green_map, blue_map], axis=2)\n    return segmentation_map\n\n\ndef image_overlay(image, segmented_image):\n    alpha = 1  # transparency for the original image\n    beta  = 0.8  # transparency for the segmentation map\n    gamma = 0  # scalar added to each sum\n  \n    image = np.array(image)\n    segmented_image = cv2.cvtColor(segmented_image, cv2.COLOR_RGB2BGR)\n    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)\n     \n    cv2.addWeighted(image, alpha, segmented_image, beta, gamma, image)\n  \n    return image","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:12:47.637342Z","iopub.execute_input":"2023-07-30T15:12:47.637623Z","iopub.status.idle":"2023-07-30T15:12:47.647136Z","shell.execute_reply.started":"2023-07-30T15:12:47.637598Z","shell.execute_reply":"2023-07-30T15:12:47.646131Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install cython","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:12:47.648982Z","iopub.execute_input":"2023-07-30T15:12:47.649605Z","iopub.status.idle":"2023-07-30T15:12:59.719550Z","shell.execute_reply.started":"2023-07-30T15:12:47.649572Z","shell.execute_reply":"2023-07-30T15:12:59.718398Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:12:59.722194Z","iopub.execute_input":"2023-07-30T15:12:59.722564Z","iopub.status.idle":"2023-07-30T15:13:20.034743Z","shell.execute_reply.started":"2023-07-30T15:12:59.722526Z","shell.execute_reply":"2023-07-30T15:13:20.033547Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import json\nfrom pycocotools.coco import COCO\nimport torch\nimport torchvision\nimport cv2\nimport torchvision.transforms as transforms\nfrom PIL import Image","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:20.036469Z","iopub.execute_input":"2023-07-30T15:13:20.037107Z","iopub.status.idle":"2023-07-30T15:13:20.467122Z","shell.execute_reply.started":"2023-07-30T15:13:20.037065Z","shell.execute_reply":"2023-07-30T15:13:20.466105Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_loc=\"/kaggle/input/dlsprint2/badlad/images/train\"\nlabel_loc = \"/kaggle/input/dlsprint2/badlad/labels/yolov8_format/train\"\ntrain_label_loc=\"/kaggle/input/dlsprint2/badlad/labels/coco_format/train/badlad-train-coco.json\"\ntest_loc=\"/kaggle/input/dlsprint2/badlad/images/test\"","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:20.469834Z","iopub.execute_input":"2023-07-30T15:13:20.470308Z","iopub.status.idle":"2023-07-30T15:13:20.476213Z","shell.execute_reply.started":"2023-07-30T15:13:20.470271Z","shell.execute_reply":"2023-07-30T15:13:20.474374Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def channel_shift(img, img1):\n    img2 = np.zeros_like(img)\n    img2[:,:,0] = img1\n    img2[:,:,1] = img1\n    img2[:,:,2] = img1\n    return img2","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:20.477502Z","iopub.execute_input":"2023-07-30T15:13:20.478380Z","iopub.status.idle":"2023-07-30T15:13:20.487184Z","shell.execute_reply.started":"2023-07-30T15:13:20.478346Z","shell.execute_reply":"2023-07-30T15:13:20.486185Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def channel_shift(img, img1):\n    img2 = np.zeros_like(img)\n    img2[:,:,0] = img1\n    img2[:,:,1] = img1\n    img2[:,:,2] = img1\n    return img2","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:20.488538Z","iopub.execute_input":"2023-07-30T15:13:20.488938Z","iopub.status.idle":"2023-07-30T15:13:20.498926Z","shell.execute_reply.started":"2023-07-30T15:13:20.488901Z","shell.execute_reply":"2023-07-30T15:13:20.497934Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def make_mask_form_ann(img, lbl):\n    masks = []\n    height, width = 800, 1333\n    img = cv2.resize(img, (width, height))\n    h, w, c = img.shape\n    for i in range(4):\n        masks.append(channel_shift(img, empty_mask(img)))\n    f = open(lbl)\n    coords = f.readlines()\n    for co in coords:\n        if len(co)==0:\n            continue\n        cls = int(co[0])\n        pts = co[2:].split(' ')\n        points = np.array(get_pts(pts, h, w))\n        cv2.fillPoly(masks[cls], pts=[points], color=(255, 255, 255))\n    mx = []\n    for mask in masks:\n        gs = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)\n        gs = np.array(gs//255, dtype=np.float32)\n        mx.append(gs)\n    return mx","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:20.500421Z","iopub.execute_input":"2023-07-30T15:13:20.500848Z","iopub.status.idle":"2023-07-30T15:13:20.511600Z","shell.execute_reply.started":"2023-07-30T15:13:20.500813Z","shell.execute_reply":"2023-07-30T15:13:20.510695Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_pts(pts, h, w):\n    points = []\n    for i in range(0, len(pts), 2):\n        x, y = int(float(pts[i])*w), int(float(pts[i+1])*h)\n        points.append([x, y])\n    return points","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:20.513386Z","iopub.execute_input":"2023-07-30T15:13:20.513861Z","iopub.status.idle":"2023-07-30T15:13:20.523019Z","shell.execute_reply.started":"2023-07-30T15:13:20.513829Z","shell.execute_reply":"2023-07-30T15:13:20.522005Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def empty_mask(img):\n    mask = np.zeros(img.shape[:2], dtype=\"uint8\")\n    return np.array(mask, dtype=bool)","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:20.524472Z","iopub.execute_input":"2023-07-30T15:13:20.525006Z","iopub.status.idle":"2023-07-30T15:13:20.534866Z","shell.execute_reply.started":"2023-07-30T15:13:20.524973Z","shell.execute_reply":"2023-07-30T15:13:20.533847Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from pathlib import Path\n\nTRAIN_IMG_DIR = Path(\"/kaggle/input/dlsprint2/badlad/images/train\")\n\nTRAIN_COCO_PATH = Path(\"/kaggle/input/dlsprint2/badlad/labels/coco_format/train/badlad-train-coco.json\")\n\nTEST_IMG_DIR = Path(\"/kaggle/input/dlsprint2/badlad/images/test\")\n\nTEST_METADATA_PATH = Path(\"/kaggle/input/dlsprint2/badlad/badlad-test-metadata.json\")","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:20.541252Z","iopub.execute_input":"2023-07-30T15:13:20.541522Z","iopub.status.idle":"2023-07-30T15:13:20.548200Z","shell.execute_reply.started":"2023-07-30T15:13:20.541492Z","shell.execute_reply":"2023-07-30T15:13:20.547207Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import json\nwith TRAIN_COCO_PATH.open() as f:\n    train_dict = json.load(f)\n# Usage\ntrains_ds = train_dict['annotations']\nimages = train_dict['images']\n# prepared_dataset = prepare_maskformer_dataset(annotations, images)","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:20.550473Z","iopub.execute_input":"2023-07-30T15:13:20.550773Z","iopub.status.idle":"2023-07-30T15:13:27.074684Z","shell.execute_reply.started":"2023-07-30T15:13:20.550748Z","shell.execute_reply":"2023-07-30T15:13:27.073639Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def transform_annotations(annotations):\n    transformed_annotations = {}\n\n    for record in annotations:\n        image_id = record['image_id']\n        category_id = record['category_id']\n        segmentation = record['segmentation']\n\n        if image_id in transformed_annotations:\n            transformed_annotations[image_id]['segmentations'].append(segmentation)\n            transformed_annotations[image_id]['category_ids'].append(category_id)\n        else:\n            transformed_annotations[image_id] = {\n                'segmentations': [segmentation],\n                'category_ids': [category_id]\n            }\n\n    transformed_annotations_list = [{'image_id': key, 'segmentations': value['segmentations'], 'category_ids': value['category_ids']} for key, value in transformed_annotations.items()]\n    \n    return transformed_annotations_list\nnew_annotations = transform_annotations(trains_ds)","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:27.076138Z","iopub.execute_input":"2023-07-30T15:13:27.076602Z","iopub.status.idle":"2023-07-30T15:13:27.416479Z","shell.execute_reply.started":"2023-07-30T15:13:27.076563Z","shell.execute_reply":"2023-07-30T15:13:27.415473Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import random\nfrom sklearn.model_selection import train_test_split\n\ndef filter_annotations_by_category(annotations, category_id):\n    filtered_annotations_table = [ann for ann in annotations if category_id in ann['category_ids']]\n    filtered_annotations_without_table = [ann for ann in annotations if category_id not in ann['category_ids']]\n\n    # Reduce the size of filtered_annotations_table and filtered_annotations_without_table separately\n    random.seed(42)  # For reproducibility\n    filtered_annotations_without_table = random.sample(filtered_annotations_without_table, int(0.3 * len(filtered_annotations_without_table)))\n\n    return filtered_annotations_table, filtered_annotations_without_table\n\n# Assuming new_annotations is a list of annotations in the format provided by transform_annotations()\ncategory_id_to_keep = 3\nfiltered_annotations_table, filtered_annotations_without_table = filter_annotations_by_category(new_annotations, category_id_to_keep)\n\n# Split filtered_annotations_table into train_table and valid_table\ntrain_table, valid_table = train_test_split(filtered_annotations_table, test_size=0.2, random_state=42)\n\n# Split filtered_annotations_without_table into train_without_table and valid_without_table\ntrain_without_table, valid_without_table = train_test_split(filtered_annotations_without_table, test_size=0.2, random_state=42)\n\n# Combine train_table and train_without_table to get the final train set\ntrain_annotations = train_table + train_without_table\n\n# Combine valid_table and valid_without_table to get the final validation set\nvalid_annotations = valid_table + valid_without_table\n","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:27.418054Z","iopub.execute_input":"2023-07-30T15:13:27.418425Z","iopub.status.idle":"2023-07-30T15:13:27.962753Z","shell.execute_reply.started":"2023-07-30T15:13:27.418389Z","shell.execute_reply":"2023-07-30T15:13:27.961775Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(len(train_annotations), len(valid_annotations))","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:27.964217Z","iopub.execute_input":"2023-07-30T15:13:27.964551Z","iopub.status.idle":"2023-07-30T15:13:27.970221Z","shell.execute_reply.started":"2023-07-30T15:13:27.964519Z","shell.execute_reply":"2023-07-30T15:13:27.969164Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del train_dict\nimport gc\ngc.collect()","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:27.971769Z","iopub.execute_input":"2023-07-30T15:13:27.972406Z","iopub.status.idle":"2023-07-30T15:13:28.475018Z","shell.execute_reply.started":"2023-07-30T15:13:27.972371Z","shell.execute_reply":"2023-07-30T15:13:28.473937Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nfrom torch.utils.data import Dataset\n\nclass ImageSegmentationDataset(Dataset):\n    \"\"\"Image segmentation dataset.\"\"\"\n\n    def __init__(self, dataset, image_info_dataset, transform_1, transform_2):\n        self.dataset = dataset\n        self.image_info_dataset = image_info_dataset\n        self.transform_1 = transform_1\n        self.transform_2 = transform_2\n        self.image_info_df = pd.DataFrame(image_info_dataset)        \n        self.image_size = (512, 512)\n    \n    def read_file(self, img):\n        img = cv2.resize(img, self.image_size, interpolation=cv2.INTER_NEAREST)\n        return img\n    \n    def __len__(self):\n        return len(self.dataset)\n    \n    def __getitem__(self, index):\n        annotation = self.dataset[index]\n        image_id = annotation['image_id']\n        category_ids = annotation['category_ids']\n        segmentations = annotation['segmentations']\n\n        row_index = self.image_info_df.index[self.image_info_df['id'] == image_id].tolist()[0]\n        file_name = '/kaggle/input/dlsprint2/badlad/images/train/' + self.image_info_df.loc[row_index, 'file_name']\n        \n        image = Image.open(file_name)\n        \n        pixel_mask = []\n        pixel_mask.append(np.ones((image.height, image.width), dtype=np.uint8))\n        for i in range(1,5):\n            pixel_mask.append(np.zeros((image.height, image.width), dtype=np.uint8))\n        \n        \n        # Create the pixel mask for each (category_id, segmentation) pair\n        for category_id, segmentation in zip(category_ids, segmentations):\n            mask = np.zeros((image.height, image.width), dtype=np.uint8)\n            segmentation = np.array(segmentation).reshape(-1, 2).astype(np.int32)\n            cv2.fillPoly(mask, [segmentation], 1) \n            mask_indices = np.where(mask > 0)\n            pixel_mask[category_id+1][mask_indices[0], mask_indices[1]] = 1\n            \n        for i in range(1,5):\n            mask_indices = np.where(pixel_mask[i] > 0)\n            pixel_mask[0][mask_indices] = 0\n        \n        for i in range(0, 5):\n            pixel_mask[i] = self.read_file(pixel_mask[i])\n            pixel_mask[i] = pixel_mask[i][:, :, np.newaxis]\n#         new_mask = np.zeros((*self.image_size, 1))\n#         new_mask[:,:,0] = pixel_mask\n        \n        \n        _mask = np.zeros((*self.image_size, 5), dtype=np.float32)\n         \n        # BACKGROUND\n        _mask[:, :, 0] = np.where(pixel_mask[0][:, :, 0] == 1,   1.0, 0.0) \n         \n        # Paragraph\n        _mask[:, :, 1] = np.where(pixel_mask[1][:, :, 0] == 1, 1.0, 0.0)  \n        \n        # Textbox\n        _mask[:, :, 2] = np.where(pixel_mask[2][:, :, 0] == 1, 1.0, 0.0)\n        \n        \n        # image\n        _mask[:, :, 3] = np.where(pixel_mask[3][:, :, 0] == 1, 1.0, 0.0)\n        \n        \n        # table\n        _mask[:, :, 4] = np.where(pixel_mask[4][:, :, 0] == 1, 1.0, 0.0)\n \n        mask = torch.from_numpy(_mask).permute(2, 0, 1)\n\n#         new_mask = torch.from_numpy(new_mask).permute(2, 0, 1)\n\n        \n#         pixel_mask = torch.from_numpy(pixel_mask).permute(2, 0, 1)\n        \n        original_image = np.array(image.convert('RGB'))\n        \n        if original_image.ndim == 3:\n            transformed = self.transform_1(image=original_image)\n            \n        else:\n            transformed = self.transform_2(image=original_image)\n            \n            \n        image = transformed['image']\n        # convert to C, H, W\n        \n        image = image.transpose(2,0,1)\n\n        return image, mask","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:28.476669Z","iopub.execute_input":"2023-07-30T15:13:28.477160Z","iopub.status.idle":"2023-07-30T15:13:28.498932Z","shell.execute_reply.started":"2023-07-30T15:13:28.477124Z","shell.execute_reply":"2023-07-30T15:13:28.497933Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import albumentations as A\n\nADE_MEAN = (0.485, 0.456, 0.406)\nADE_STD = (0.229, 0.224, 0.225)\n\n# ADE_MEAN = np.array([127.5, 127.5, 127.5])\n# ADE_STD = np.array([127.5, 127.5, 127.5])\n\n# ADE_MEAN = np.array([0, 0, 0])\n# ADE_STD = np.array([1, 1, 1]) / 255\nprint(ADE_MEAN, \"  \", ADE_STD)\n\ntrain_transform_1 = A.Compose([\n    A.Resize(width=512, height=512),  # Add the Resize transformation\n    A.Normalize(mean=ADE_MEAN, std=ADE_STD),\n])\n\ntrain_transform_2 = A.Compose([\n    A.Resize(width=512, height=512),  # Add the Resize transformation\n    A.Normalize(mean=127.5, std=127.5),\n])\n\n\ntest_transform = A.Compose([\n    A.Resize(width=512, height=512),\n    A.Normalize(mean=ADE_MEAN, std=ADE_STD),\n])\n\n\ntrain_dataset = ImageSegmentationDataset(train_annotations, images, transform_1=train_transform_1, transform_2=train_transform_2)\nvalid_dataset = ImageSegmentationDataset(valid_annotations, images, transform_1=train_transform_1, transform_2=train_transform_2)","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:28.500558Z","iopub.execute_input":"2023-07-30T15:13:28.500969Z","iopub.status.idle":"2023-07-30T15:13:29.530002Z","shell.execute_reply.started":"2023-07-30T15:13:28.500936Z","shell.execute_reply":"2023-07-30T15:13:29.528874Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image, segmentation_map = train_dataset[0]\nimage_path = images[0]['file_name']\nimage_path","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:29.533122Z","iopub.execute_input":"2023-07-30T15:13:29.533465Z","iopub.status.idle":"2023-07-30T15:13:30.378358Z","shell.execute_reply.started":"2023-07-30T15:13:29.533437Z","shell.execute_reply":"2023-07-30T15:13:30.377365Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image_path = \"/kaggle/input/dlsprint2/badlad/images/train/\" + image_path\nimage = Image.open(image_path)\nimage\noriginal_image = np.array(image)\nprint(original_image.shape)\nimage = train_transform_2(image = original_image)\nprint(image['image'].shape)","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:30.379668Z","iopub.execute_input":"2023-07-30T15:13:30.380342Z","iopub.status.idle":"2023-07-30T15:13:30.413870Z","shell.execute_reply.started":"2023-07-30T15:13:30.380305Z","shell.execute_reply":"2023-07-30T15:13:30.412809Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"id2label = {0: 'background', 1: 'paragraph', 2:'textbox', 3: 'image', 4: 'table'}","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:30.415545Z","iopub.execute_input":"2023-07-30T15:13:30.416249Z","iopub.status.idle":"2023-07-30T15:13:30.421561Z","shell.execute_reply.started":"2023-07-30T15:13:30.416211Z","shell.execute_reply":"2023-07-30T15:13:30.420356Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\ndef visualize_segmentation_map(segmentation_map):\n    # Define the colormap for visualization (0: black, 1: white)\n    cmap = plt.get_cmap('gray', 2)\n\n    # Create a figure and axis\n    fig, ax = plt.subplots(figsize=(8, 8))\n\n    # Plot the segmentation map using the colormap\n    ax.imshow(segmentation_map, cmap=cmap, vmin=0, vmax=1)\n\n    # Hide axis ticks and labels\n    ax.set_xticks([])\n    ax.set_yticks([])\n\n    # Show the plot\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:30.423067Z","iopub.execute_input":"2023-07-30T15:13:30.423413Z","iopub.status.idle":"2023-07-30T15:13:30.435079Z","shell.execute_reply.started":"2023-07-30T15:13:30.423380Z","shell.execute_reply":"2023-07-30T15:13:30.434031Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# from PIL import Image\n# for i in range(0,10):\n#     image, segmentation_map = valid_dataset[i]\n#     print(image.shape)\n#     print(segmentation_map.shape)\n\n#     unnormalized_image = (image * np.array(ADE_STD)[:, None, None]) + np.array(ADE_MEAN)[:, None, None]\n#     unnormalized_image = (unnormalized_image * 255).astype(np.uint8)\n#     unnormalized_image = np.moveaxis(unnormalized_image, 0, -1)\n    \n#     image = Image.fromarray(unnormalized_image)\n\n#     # Get unique labels from the flattened map\n#     label1 = [int(label) for label in np.unique(segmentation_map[0])]\n\n#     label2 = [int(label) for label in np.unique(segmentation_map[1])]\n#     label3 = [int(label) for label in np.unique(segmentation_map[2])]\n#     label4 = [int(label) for label in np.unique(segmentation_map[3])]\n#     label5 = [int(label) for label in np.unique(segmentation_map[4])]\n#     print(label1)\n#     print(label2)\n#     print(label3)\n#     print(label4)\n#     print(label5)\n#     # Show the image using matplotlib\n#     plt.imshow(image)\n#     plt.axis('off')  # Turn off axis labels\n#     plt.show()\n#     image.show()\n#     visualize_segmentation_map(segmentation_map[0])\n#     visualize_segmentation_map(segmentation_map[1])\n#     visualize_segmentation_map(segmentation_map[2])\n#     visualize_segmentation_map(segmentation_map[3])\n#     visualize_segmentation_map(segmentation_map[4])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torchvision.models.segmentation.deeplabv3 import DeepLabHead\nfrom torchvision import models\n\n\ndef createDeepLabv3(outputchannels=5):\n    \"\"\"DeepLabv3 class with custom head\n\n    Args:\n        outputchannels (int, optional): The number of output channels\n        in your dataset masks. Defaults to 1.\n\n    Returns:\n        model: Returns the DeepLabv3 model with the ResNet101 backbone.\n    \"\"\"\n    \n#     model = models.segmentation.deeplabv3_mobilenet_v3_large(pretrained=True,\n#                                                     progress=True)\n    model = models.segmentation.deeplabv3_resnet50(pretrained=True,\n                                                    progress=True)\n    model.classifier = DeepLabHead(2048, outputchannels)\n    # Set the model in training mode\n    model.train()\n    return model","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:58.589882Z","iopub.execute_input":"2023-07-30T15:13:58.590278Z","iopub.status.idle":"2023-07-30T15:13:58.596841Z","shell.execute_reply.started":"2023-07-30T15:13:58.590246Z","shell.execute_reply":"2023-07-30T15:13:58.595729Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Initialize the model\nmodel = createDeepLabv3(5)\n","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:13:59.257476Z","iopub.execute_input":"2023-07-30T15:13:59.257882Z","iopub.status.idle":"2023-07-30T15:14:00.821858Z","shell.execute_reply.started":"2023-07-30T15:13:59.257846Z","shell.execute_reply":"2023-07-30T15:14:00.820847Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader\nfrom torchvision.models.segmentation import deeplabv3_mobilenet_v3_large\nfrom torchvision.models.segmentation import deeplabv3_resnet50, deeplabv3_resnet101\nfrom torchvision import transforms\nfrom tqdm import tqdm\n\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'\nmodel.to(device)\nmodel.train()\n# Define the training parameters\nnum_classes = 5\nbatch_size = 2\nlr = 0.001\nnum_epochs = 6\nshuffle = True\nnum_workers = 2\n\ntrain_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)\nvalid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)\noptimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4)\ncriterion = nn.CrossEntropyLoss()","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:36:56.506449Z","iopub.execute_input":"2023-07-30T15:36:56.506894Z","iopub.status.idle":"2023-07-30T15:36:56.528218Z","shell.execute_reply.started":"2023-07-30T15:36:56.506859Z","shell.execute_reply":"2023-07-30T15:36:56.526686Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport csv\nimport time\nimport numpy as np\nimport torch\nfrom torch.optim import Adam\nfrom tqdm import tqdm","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:36:56.662553Z","iopub.execute_input":"2023-07-30T15:36:56.663207Z","iopub.status.idle":"2023-07-30T15:36:56.668357Z","shell.execute_reply.started":"2023-07-30T15:36:56.663167Z","shell.execute_reply":"2023-07-30T15:36:56.667343Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def intermediate_metric_calculation(predictions, targets, use_dice=False, smooth=1e-6, dims=(2,3)):\n    # dims corresponding to image height and width: [B, C, H, W].\n    # Intersection: |G ∩ P|. Shape: (batch_size, num_classes)\n    intersection = (predictions * targets).sum(dim=dims) + smooth \n \n    # Summation: |G| + |P|. Shape: (batch_size, num_classes).\n    summation = (predictions.sum(dim=dims) + targets.sum(dim=dims)) + smooth \n  \n    if use_dice:\n        # Dice Shape: (batch_size, num_classes) \n        metric = (2.0 * intersection) / summation\n    else:\n        # Union. Shape: (batch_size, num_classes)\n        union = summation - intersection\n \n        # IoU Shape: (batch_size, num_classes)\n        metric = intersection /  union\n         \n    # Compute the mean over the remaining axes (batch and classes). \n    # Shape: Scalar\n    total = metric.mean()\n    return total","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:36:56.836876Z","iopub.execute_input":"2023-07-30T15:36:56.839412Z","iopub.status.idle":"2023-07-30T15:36:56.848821Z","shell.execute_reply.started":"2023-07-30T15:36:56.839372Z","shell.execute_reply":"2023-07-30T15:36:56.847769Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_model(model, train_loader, valid_loader, criterion, optimizer, num_epochs=3):\n    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n    model.to(device)\n\n    best_valid_loss = float('inf')  # Initialize with a very large value\n\n    for epoch in range(1, num_epochs+1):\n        model.train()\n        train_loss = 0.0\n        valid_loss = 0.0\n        error = 0\n        for inputs, masks in tqdm(train_loader):\n            try:\n                inputs = inputs.to(device)\n                masks = masks.to(device)\n\n                optimizer.zero_grad()\n                outputs = model(inputs)['out']\n                loss = criterion(outputs, masks.argmax(dim=1))\n                loss.backward()\n                optimizer.step()\n                train_loss += loss.item()\n                \n            except:\n                error += 1\n                print(\"train error: \", error)\n\n        model.eval()\n        with torch.no_grad():\n            for inputs, masks in valid_loader:\n                try:\n                    inputs = inputs.to(device)\n                    masks = masks.to(device)\n    \n                    # Get the predictions from the 'out' key of the OrderedDict\n                    outputs = model(inputs)['out']\n                    loss = criterion(outputs, masks.argmax(dim=1))\n                    valid_loss += loss.item()\n                except:\n                    error += 1\n                    print(\"val error: \", error)\n\n        train_loss /= len(train_loader)\n        valid_loss /= len(valid_loader)\n\n        print(f\"Epoch [{epoch + 1}/{num_epochs}] - Train Loss: {train_loss:.4f}, Valid Loss: {valid_loss:.4f}\")\n\n        # Save the model every 3 epochs if the current validation loss is the best so far\n        if epoch % 2 == 0 and valid_loss < best_valid_loss:\n            best_valid_loss = valid_loss\n            torch.save(model.state_dict(), 'best_model_epoch_{}.pth'.format(epoch))\n\n    print(\"Training complete!\")","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:41:16.753236Z","iopub.execute_input":"2023-07-30T15:41:16.753627Z","iopub.status.idle":"2023-07-30T15:41:16.767509Z","shell.execute_reply.started":"2023-07-30T15:41:16.753593Z","shell.execute_reply":"2023-07-30T15:41:16.766447Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import time\n# import csv\n# import numpy as np\n# from tqdm import tqdm\n# import torch\n# import torch.nn.functional as F\n# from torch.optim import Adam\n\n# def intermediate_metric_calculation(predictions, targets, use_dice=False, smooth=1e-6, dims=(2,3)):\n#     # dims corresponding to image height and width: [B, C, H, W].\n#     # Intersection: |G ∩ P|. Shape: (batch_size, num_classes)\n#     intersection = (predictions * targets).sum(dim=dims) + smooth \n \n#     # Summation: |G| + |P|. Shape: (batch_size, num_classes).\n#     summation = (predictions.sum(dim=dims) + targets.sum(dim=dims)) + smooth \n  \n#     if use_dice:\n#         # Dice Shape: (batch_size, num_classes) \n#         metric = (2.0 * intersection) / summation\n#     else:\n#         # Union. Shape: (batch_size, num_classes)\n#         union = summation - intersection\n \n#         # IoU Shape: (batch_size, num_classes)\n#         metric = intersection /  union\n         \n#     # Compute the mean over the remaining axes (batch and classes). \n#     # Shape: Scalar\n#     total = metric.mean()\n#     return total\n\n# def train_model(model, train_loader, valid_loader, num_epochs=3):\n#     since = time.time()\n#     best_model_wts = model.state_dict()\n#     best_loss = 1e10\n#     device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n#     model.to(device)\n#     fieldnames = ['epoch', 'Train_loss', 'Test_loss']\n\n#     with open('log.csv', 'w', newline='') as csvfile:\n#         writer = csv.DictWriter(csvfile, fieldnames=fieldnames)\n#         writer.writeheader()\n\n#     optimizer = Adam(model.parameters())\n#     criterion = torch.nn.CrossEntropyLoss()  # CrossEntropyLoss for multi-class segmentation\n\n#     for epoch in range(1, num_epochs + 1):\n#         print(f'Epoch {epoch}/{num_epochs}')\n#         print('-' * 10)\n#         batchsummary = {a: [0] for a in fieldnames}\n\n#         for phase in ['Train', 'Test']:\n#             model.train() if phase == 'Train' else model.eval()\n#             dataloader = train_loader if phase == 'Train' else valid_loader\n#             for inputs, masks in tqdm(dataloader):\n#                 inputs = inputs.to(device)\n#                 masks = masks.to(device)\n\n#                 with torch.set_grad_enabled(phase == 'Train'):\n#                     try:\n#                         outputs = model(inputs)\n#                         predictions = outputs['out']  # No need for sigmoid activation\n#                         print(predictions.shape)\n#                         # Flatten the predictions and masks for the loss function\n#                         predictions_flat = predictions.permute(0, 2, 3, 1).contiguous().view(-1, 5)\n#                         print(predictions_flat.shape)\n#                         masks_flat = masks.view(-1)\n#                         print(masks_flat.shape)\n\n#                         loss = criterion(predictions_flat, masks_flat)\n\n#                         if phase == 'Train':\n#                             optimizer.zero_grad()\n#                             loss.backward()\n#                             optimizer.step()\n#                             batchsummary[f'{phase}_loss'].append(loss.item())\n#                     except:\n#                         print(\"error\")\n\n#             batchsummary['epoch'] = epoch\n#             epoch_loss = np.mean(batchsummary[f'{phase}_loss'])\n#             batchsummary[f'{phase}_loss'] = epoch_loss\n#             print(f'{phase} Loss: {epoch_loss:.4f}')\n        \n#         with open('log.csv', 'a', newline='') as csvfile:\n#             writer = csv.DictWriter(csvfile, fieldnames=fieldnames)\n#             writer.writerow(batchsummary)\n\n#             if phase == 'Test' and epoch_loss < best_loss:\n#                 best_loss = epoch_loss\n#                 best_model_wts = model.state_dict()\n\n#         if epoch % 3 == 0:\n#             torch.save(best_model_wts, f'best_model_epoch_{epoch}.pth')\n\n#     time_elapsed = time.time() - since\n#     print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')\n#     print(f'Lowest Loss: {best_loss:.4f}')\n\n#     torch.save(best_model_wts, 'best_model.pth')\n#     return model","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:41:17.294808Z","iopub.execute_input":"2023-07-30T15:41:17.295180Z","iopub.status.idle":"2023-07-30T15:41:17.303566Z","shell.execute_reply.started":"2023-07-30T15:41:17.295151Z","shell.execute_reply":"2023-07-30T15:41:17.302489Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_model(model, train_loader, valid_loader, criterion, optimizer, 8)","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:41:17.802915Z","iopub.execute_input":"2023-07-30T15:41:17.803298Z","iopub.status.idle":"2023-07-30T15:41:21.941378Z","shell.execute_reply.started":"2023-07-30T15:41:17.803266Z","shell.execute_reply":"2023-07-30T15:41:21.939762Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import torch\n# import torch.nn as nn\n# from torchvision.models.segmentation import deeplabv3_mobilenet_v3_large\n\n# def create_model_with_modified_classifier(num_classes=2):\n#     # Load the base DeepLabV3 model with the specified number of classes (2 in this case)\n#     base_model = deeplabv3_mobilenet_v3_large(num_classes=num_classes)\n\n#     # Modify the classification layer to have the desired number of classes\n#     base_model.classifier[4] = nn.Conv2d(256, num_classes, kernel_size=(1, 1))\n\n#     return base_model\n\n# # Create a new model with the same architecture as the loaded model but with 2 classes in the classification layer\n# model = create_model_with_modified_classifier(num_classes=2)\n\n# # Load the saved state_dict from the checkpoint file with strict=False\n# state_dict = torch.load('/kaggle/input/deeplabv3-models/best_model_epoch_3.pth')\n\n# # Load the state_dict into the model with strict=False\n# model.load_state_dict(state_dict, strict=False)\n","metadata":{"execution":{"iopub.status.busy":"2023-07-30T15:37:08.309688Z","iopub.status.idle":"2023-07-30T15:37:08.310414Z","shell.execute_reply.started":"2023-07-30T15:37:08.310165Z","shell.execute_reply":"2023-07-30T15:37:08.310189Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.eval()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from PIL import Image\nimport torch\nimport numpy as np\nimport albumentations as A\n\n# Load the image and convert it to RGB\ninput_image = Image.open(\"/kaggle/input/dlsprint2/badlad/images/test/02bf7f8a-976d-4b51-84ce-2a09cc17578e.png\").convert(\"RGB\")\n# model1.eval()\n# Resize the input image to a larger size\ninput_image = input_image.resize((512, 512))\n\n# Define the normalization parameters (replace with actual values if available)\nmean = [0.485, 0.456, 0.406]\nstd = [0.229, 0.224, 0.225]\n\n# Create the normalization transformation using albumentations\ntrain_transform = A.Compose([\n    A.Normalize(mean=mean, std=std),\n])\n\n# Convert the input image to a numpy array\ninput_array = np.array(input_image)\n\n# Apply the normalization transformation\ninput_array = train_transform(image=input_array)[\"image\"]\n\n# Convert the numpy array to a PyTorch tensor and add batch dimension\ninput_tensor = torch.tensor(input_array).permute(2, 0, 1).unsqueeze(0)\n\n# Move the input and model to GPU for speed if available\nif torch.cuda.is_available():\n    input_tensor = input_tensor.to('cuda')\n    model.to('cuda')\n\nprint(input_tensor.shape)\nwith torch.no_grad():\n    output = model(input_tensor)['out'][0]\n    output_predictions = output.argmax(0)\n\nprint(output_predictions.shape)\n# Convert the tensor to a numpy array\noutput_predictions_np = output_predictions.cpu().numpy()\n\n# Optionally, you can visualize the segmentation mask\nimport matplotlib.pyplot as plt\n\nplt.imshow(output_predictions_np, cmap='jet')\nplt.axis('off')\nplt.show()\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torchvision import transforms\nfilename = \"/kaggle/input/dlsprint2/badlad/images/test/02bf7f8a-976d-4b51-84ce-2a09cc17578e.png\"\ninput_image = Image.open(filename)\ninput_image = input_image.convert(\"RGB\")\npreprocess = transforms.Compose([\n    transforms.ToTensor(),\n    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n])\n\ninput_tensor = preprocess(input_image)\ninput_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model\n\n# move the input and model to GPU for speed if available\nif torch.cuda.is_available():\n    input_batch = input_batch.to('cuda')\n    model.to('cuda')\n\nwith torch.no_grad():\n    output = model(input_batch)['out'][0]\noutput_predictions = output.argmax(0)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"input_image","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Get the number of classes\nnum_classes = model.classifier[4].out_channels\nprint(\"Number of classes in the model:\", num_classes)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# create a color pallette, selecting a color for each class\npalette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])\ncolors = torch.as_tensor([i for i in range(21)])[:, None] * palette\ncolors = (colors % 255).numpy().astype(\"uint8\")\n\n# plot the semantic segmentation predictions of 21 classes in each color\nr = Image.fromarray(output_predictions.byte().cpu().numpy()).resize(input_image.size)\nr.putpalette(colors)\n\nimport matplotlib.pyplot as plt\nplt.imshow(r)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pixel_to_class = {\n    0: \"background\",\n    1: \"table\",\n    # Add more mappings as needed\n}\n\n# Create the class label mask\nclass_label_mask = np.vectorize(pixel_to_class.get)(output_predictions_np)\n\n# Get all the unique detected classes in the image\ndetected_classes = np.unique(class_label_mask)\n\n# Print the detected classes\nprint(\"Detected Classes:\")\nfor detected_class in detected_classes:\n    print(detected_class)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset[0]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}