{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.10","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":56206,"databundleVersionId":6049520,"sourceType":"competition"}],"dockerImageVersionId":30512,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Detectron2\n\n[Detectron2](https://detectron2.readthedocs.io/en/latest/index.html) is a popular open-source software library developed by Facebook AI Research (FAIR) for building computer vision models. It serves as a powerful framework for object detection, instance segmentation, and keypoint detection tasks. Detectron2 is built on top of PyTorch, geared towards a more convenient way to build modular, flexible pipelines for specific Computer Vision Tasks such as object detection, instance segmentation. \n\nDetectron2 has a collection of trained models for these tasks in their [model zoo](https://github.com/facebookresearch/detectron2/blob/main/MODEL_ZOO.md). We can also use detectron2 to train pre-implemented state-of-the-art models from scratch for new datasets, as we do in this notebook. \n\nRead the [documentation](https://detectron2.readthedocs.io/en/latest/index.html). ","metadata":{}},{"cell_type":"markdown","source":"# 1 Install detectron2","metadata":{}},{"cell_type":"markdown","source":"## 1.1 Recommended Way (is not working on kaggle)","metadata":{}},{"cell_type":"code","source":"# !python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'","metadata":{"_kg_hide-input":false,"_kg_hide-output":false,"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:56:45.568651Z","iopub.execute_input":"2025-01-13T02:56:45.569276Z","iopub.status.idle":"2025-01-13T02:56:45.573770Z","shell.execute_reply.started":"2025-01-13T02:56:45.569245Z","shell.execute_reply":"2025-01-13T02:56:45.572930Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 1.2 Fast Way\nIgnore the warnings.","metadata":{}},{"cell_type":"code","source":"%%capture\nimport sys, os, distutils.core\n# Note: This is a faster way to install detectron2 in Colab, but it does not include all functionalities (e.g. compiled operators).\n# See https://detectron2.readthedocs.io/tutorials/install.html for full installation instructions\n!git clone 'https://github.com/facebookresearch/detectron2'\ndist = distutils.core.run_setup(\"./detectron2/setup.py\")\n!python -m pip install {' '.join([f\"'{x}'\" for x in dist.install_requires])}\nsys.path.insert(0, os.path.abspath('./detectron2'))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:56:45.575081Z","iopub.execute_input":"2025-01-13T02:56:45.575307Z","iopub.status.idle":"2025-01-13T02:57:05.750358Z","shell.execute_reply.started":"2025-01-13T02:56:45.575287Z","shell.execute_reply":"2025-01-13T02:57:05.749131Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 2 Notebook Config","metadata":{}},{"cell_type":"markdown","source":"## 2.1 Decisions","metadata":{}},{"cell_type":"code","source":"from datetime import datetime\n\n# if False, model is set to `PRETRAINED_PATH` model\nis_train = True\n\n# if True, evaluate on validation dataset\nis_evaluate = False\n\n# if True, run inference on test dataset\nis_inference = True\n\n# if True and `is_train` == True, `PRETRAINED_PATH` model is trained further\nis_resume_training = False\n\n# Perform augmentation\nis_augment = False\n\nSEED = int(datetime.now().timestamp())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:05.752049Z","iopub.execute_input":"2025-01-13T02:57:05.752991Z","iopub.status.idle":"2025-01-13T02:57:05.758591Z","shell.execute_reply.started":"2025-01-13T02:57:05.752962Z","shell.execute_reply":"2025-01-13T02:57:05.757638Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2.2 Paths","metadata":{}},{"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\")\n\n# Training output directory\nOUTPUT_DIR = Path(\"./output\")\nOUTPUT_MODEL = OUTPUT_DIR/\"model_final.pth\"\n\n# Path to your pretrained model weights\nPRETRAINED_PATH = Path(\"\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:05.761127Z","iopub.execute_input":"2025-01-13T02:57:05.761457Z","iopub.status.idle":"2025-01-13T02:57:05.776371Z","shell.execute_reply.started":"2025-01-13T02:57:05.761428Z","shell.execute_reply":"2025-01-13T02:57:05.775519Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Model path based on Decisions\nMODEL_PATH = OUTPUT_MODEL if is_train else PRETRAINED_PATH","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:05.777663Z","iopub.execute_input":"2025-01-13T02:57:05.777999Z","iopub.status.idle":"2025-01-13T02:57:05.786360Z","shell.execute_reply.started":"2025-01-13T02:57:05.777970Z","shell.execute_reply":"2025-01-13T02:57:05.785623Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2.3 imports","metadata":{}},{"cell_type":"code","source":"# detectron2\nfrom detectron2.utils.memory import retry_if_cuda_oom\nfrom detectron2.utils.logger import setup_logger\nfrom detectron2.checkpoint import DetectionCheckpointer\nfrom detectron2.modeling import build_model\nfrom detectron2.evaluation import COCOEvaluator, inference_on_dataset\nimport detectron2.data.transforms as T\nfrom detectron2.data import detection_utils as utils\nfrom detectron2.data import DatasetCatalog, MetadataCatalog, build_detection_test_loader, build_detection_train_loader, DatasetMapper\nfrom detectron2.utils.visualizer import Visualizer\nfrom detectron2.structures import BoxMode\nfrom detectron2.engine import DefaultPredictor, DefaultTrainer\nfrom detectron2.config import get_cfg\nfrom detectron2 import model_zoo\n\nimport pandas as pd\nimport numpy as np\nfrom tqdm.notebook import tqdm  # progress bar\nimport matplotlib.pyplot as plt\nimport json\nimport cv2\nimport copy\nfrom typing import Optional\n\nfrom IPython.display import FileLink\n\n# torch\nimport torch\n\nimport gc\n\nimport warnings\n# Ignore \"future\" warnings and Data-Frame-Slicing warnings.\nwarnings.filterwarnings('ignore')\n\nsetup_logger()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:05.787414Z","iopub.execute_input":"2025-01-13T02:57:05.787691Z","iopub.status.idle":"2025-01-13T02:57:06.682202Z","shell.execute_reply.started":"2025-01-13T02:57:05.787665Z","shell.execute_reply":"2025-01-13T02:57:06.680959Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 3 COCO Annotations Data","metadata":{}},{"cell_type":"markdown","source":"## 3.1 Load","metadata":{}},{"cell_type":"code","source":"with TRAIN_COCO_PATH.open() as f:\n    train_dict = json.load(f)\n\nwith TEST_METADATA_PATH.open() as f:\n    test_dict = json.load(f)\n\nprint(\"#### LABELS AND METADATA LOADED ####\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:06.683594Z","iopub.execute_input":"2025-01-13T02:57:06.684259Z","iopub.status.idle":"2025-01-13T02:57:12.227864Z","shell.execute_reply.started":"2025-01-13T02:57:06.684222Z","shell.execute_reply":"2025-01-13T02:57:12.226950Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3.2 Observe","metadata":{}},{"cell_type":"code","source":"def organize_coco_data(data_dict: dict) -> tuple[list[str], list[dict], list[dict]]:\n    thing_classes: list[str] = []\n\n    # Map Category Names to IDs\n    for cat in data_dict['categories']:\n        thing_classes.append(cat['name'])\n\n    # Images\n    images_metadata: list[dict] = data_dict['images']\n\n    # Convert COCO annotations to detectron2 annotations format\n    data_annotations = []\n    for ann in data_dict['annotations']:\n        # coco format -> detectron2 format\n        annot_obj = {\n            # Annotation ID\n            \"id\": ann['id'],\n\n            # Segmentation Polygon (x, y) coords\n            \"gt_masks\": ann['segmentation'],\n\n            # Image ID for this annotation (Which image does this annotation belong to?)\n            \"image_id\": ann['image_id'],\n\n            # Category Label (0: paragraph, 1: text box, 2: image, 3: table)\n            \"category_id\": ann['category_id'],\n\n            \"x_min\": ann['bbox'][0],  # left\n            \"y_min\": ann['bbox'][1],  # top\n            \"x_max\": ann['bbox'][0] + ann['bbox'][2],  # left+width\n            \"y_max\": ann['bbox'][1] + ann['bbox'][3]  # top+height\n        }\n        data_annotations.append(annot_obj)\n\n    return thing_classes, images_metadata, data_annotations","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:12.229215Z","iopub.execute_input":"2025-01-13T02:57:12.229571Z","iopub.status.idle":"2025-01-13T02:57:12.236983Z","shell.execute_reply.started":"2025-01-13T02:57:12.229538Z","shell.execute_reply":"2025-01-13T02:57:12.235994Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"thing_classes, images_metadata, data_annotations = organize_coco_data(\n    train_dict\n)\n\nthing_classes_test, images_metadata_test, _ = organize_coco_data(\n    test_dict\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:12.239711Z","iopub.execute_input":"2025-01-13T02:57:12.240034Z","iopub.status.idle":"2025-01-13T02:57:13.316523Z","shell.execute_reply.started":"2025-01-13T02:57:12.240011Z","shell.execute_reply":"2025-01-13T02:57:13.315548Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"These are the categories we are going to detect.","metadata":{}},{"cell_type":"code","source":"print(thing_classes)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:13.317641Z","iopub.execute_input":"2025-01-13T02:57:13.317977Z","iopub.status.idle":"2025-01-13T02:57:13.322666Z","shell.execute_reply.started":"2025-01-13T02:57:13.317946Z","shell.execute_reply":"2025-01-13T02:57:13.321821Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_metadata = pd.DataFrame(images_metadata)\ntrain_metadata = train_metadata[['id', 'file_name', 'width', 'height']]\ntrain_metadata = train_metadata.rename(columns={\"id\": \"image_id\"})\nprint(\"train_metadata size=\", len(train_metadata))\ntrain_metadata.head(5)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:13.323787Z","iopub.execute_input":"2025-01-13T02:57:13.324127Z","iopub.status.idle":"2025-01-13T02:57:13.387374Z","shell.execute_reply.started":"2025-01-13T02:57:13.324094Z","shell.execute_reply":"2025-01-13T02:57:13.386495Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_annot_df = pd.DataFrame(data_annotations)\nprint(\"train_annot_df size=\", len(train_annot_df))\ntrain_annot_df.head(5)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:13.388487Z","iopub.execute_input":"2025-01-13T02:57:13.388748Z","iopub.status.idle":"2025-01-13T02:57:14.290642Z","shell.execute_reply.started":"2025-01-13T02:57:13.388726Z","shell.execute_reply":"2025-01-13T02:57:14.289739Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Here `gt_masks` are the sequence of `(x, y)` coordinates of vertices of the polygon surrounding the target object. ","metadata":{}},{"cell_type":"code","source":"test_metadata = pd.DataFrame(images_metadata_test)\ntest_metadata = test_metadata[['id', 'file_name', 'width', 'height']]\ntest_metadata = test_metadata.rename(columns={\"id\": \"image_id\"})\nprint(\"test_metadata size=\", len(test_metadata))\ntest_metadata.head(5)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:14.291838Z","iopub.execute_input":"2025-01-13T02:57:14.292089Z","iopub.status.idle":"2025-01-13T02:57:14.326416Z","shell.execute_reply.started":"2025-01-13T02:57:14.292067Z","shell.execute_reply":"2025-01-13T02:57:14.325556Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 4 Preparing Data for Training","metadata":{}},{"cell_type":"markdown","source":" 4.1 Train-Validation Split","metadata":{}},{"cell_type":"code","source":"#TRAIN_SPLIT = 0.95","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:14.328138Z","iopub.execute_input":"2025-01-13T02:57:14.328382Z","iopub.status.idle":"2025-01-13T02:57:14.331873Z","shell.execute_reply.started":"2025-01-13T02:57:14.328361Z","shell.execute_reply":"2025-01-13T02:57:14.330983Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"VAL_SPLIT = 0.05 # 5% validation split, this can be changed","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:14.332885Z","iopub.execute_input":"2025-01-13T02:57:14.333113Z","iopub.status.idle":"2025-01-13T02:57:14.341937Z","shell.execute_reply.started":"2025-01-13T02:57:14.333092Z","shell.execute_reply":"2025-01-13T02:57:14.340967Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"n_dataset = len(train_metadata)\n\n#Calculating the number of training samples based on (validation) split\nn_train = int(n_dataset * (1 - VAL_SPLIT)) # using (1 - VAL_SPLIT) for training set\nprint(\"n_dataset\", n_dataset, \"n_train\", n_train, \"n_val\", n_dataset - n_train)\n\n#SEED is defined for reproducibility\nSEED = 42 # or any other value\nnp.random.seed(SEED)\n\ninds = np.random.permutation(n_dataset)\ntrain_inds, valid_inds = inds[:n_train], inds[n_train:]\n\n#Ensuring the split indices do not overlap and the dataset is shuffled properly\nassert len(set(train_inds).intersection(valid_inds)) == 0, \"Data leakage: Train and Validation sets overlap\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:14.342985Z","iopub.execute_input":"2025-01-13T02:57:14.343220Z","iopub.status.idle":"2025-01-13T02:57:14.357233Z","shell.execute_reply.started":"2025-01-13T02:57:14.343200Z","shell.execute_reply":"2025-01-13T02:57:14.356371Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4.2 Formatting Data for `detectron2`","metadata":{}},{"cell_type":"code","source":"def convert_coco_to_detectron2_format(\n    imgdir: Path,\n    metadata_df: pd.DataFrame,\n    annot_df: Optional[pd.DataFrame] = None,\n    target_indices: Optional[np.ndarray] = None,\n):\n\n    dataset_dicts = []\n    \n    for _, train_meta_row in tqdm(metadata_df.iterrows(), total=len(metadata_df)):\n        # Iterate over each image\n        image_id, filename, width, height = train_meta_row[['image_id', 'file_name', 'width', 'height']]\n        \n        annotations = []\n\n        # If train/validation data, then there will be annotations\n        if annot_df is not None:\n            for _, ann in annot_df.query(\"image_id == @image_id\").iterrows():\n                # Get annotations of current iteration's image\n                class_id = ann[\"category_id\"]\n                gt_masks = ann[\"gt_masks\"]\n\n                #Ensuring gt_masks is a list of polygons (even if there are multiple)\n                if isinstance(gt_masks, list):\n                    if not all(isinstance(mask, list) for mask in gt_masks):\n                        #If gt_masks is a list but not of lists, we wrap it into a list\n                        gt_masks = [gt_masks]\n\n                else:\n                    #If it's not a list, make sure to convert it into one\n                    gt_masks = [gt_masks]\n                    \n                bbox_resized = [\n                    float(ann[\"x_min\"]),\n                    float(ann[\"y_min\"]),\n                    float(ann[\"x_max\"]),\n                    float(ann[\"y_max\"]),\n                ]\n\n                annotation = {\n                    \"bbox\": bbox_resized,\n                    \"bbox_mode\": BoxMode.XYXY_ABS,\n                    \"segmentation\": gt_masks,\n                    \"category_id\": class_id,\n                }\n\n                annotations.append(annotation)\n\n        # coco format -> detectron2 format dict\n        record = {\n            \"file_name\": str(imgdir/filename),\n            \"image_id\": image_id,\n            \"width\": width,\n            \"height\": height,\n            \"annotations\": annotations\n        }\n\n        dataset_dicts.append(record)\n\n    if target_indices is not None:\n        dataset_dicts = [dataset_dicts[i] for i in target_indices]\n\n    return dataset_dicts","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:14.358627Z","iopub.execute_input":"2025-01-13T02:57:14.359233Z","iopub.status.idle":"2025-01-13T02:57:14.368341Z","shell.execute_reply.started":"2025-01-13T02:57:14.359208Z","shell.execute_reply":"2025-01-13T02:57:14.367442Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4.3 Registering and Loading Data for `detectron2`","metadata":{}},{"cell_type":"code","source":"DATA_REGISTER_TRAINING = \"badlad_train\"\nDATA_REGISTER_VALID    = \"badlad_valid\"\nDATA_REGISTER_TEST     = \"badlad_test\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:14.369362Z","iopub.execute_input":"2025-01-13T02:57:14.369601Z","iopub.status.idle":"2025-01-13T02:57:14.381452Z","shell.execute_reply.started":"2025-01-13T02:57:14.369580Z","shell.execute_reply":"2025-01-13T02:57:14.380651Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Define the is_train variable \n#is_train = True # Set this to True or False based on your requirement","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:14.382439Z","iopub.execute_input":"2025-01-13T02:57:14.382745Z","iopub.status.idle":"2025-01-13T02:57:14.391719Z","shell.execute_reply.started":"2025-01-13T02:57:14.382723Z","shell.execute_reply":"2025-01-13T02:57:14.390831Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# TRAINING: Register Training data\nif DATA_REGISTER_TRAINING not in DatasetCatalog.list():\n    print(\"Registering training data...\")\n    DatasetCatalog.register(\n        DATA_REGISTER_TRAINING,\n        lambda: convert_coco_to_detectron2_format(\n            TRAIN_IMG_DIR,\n            train_metadata,\n            train_annot_df,\n            target_indices=train_inds,\n        ),\n    )\n    print(\"Training data registered.\")\n\n    # Set Training data categories\n    MetadataCatalog.get(DATA_REGISTER_TRAINING).set(thing_classes=thing_classes)\n    print(\"Training data categories set.\")\n\ndataset_dicts_train = DatasetCatalog.get(DATA_REGISTER_TRAINING)\nmetadata_dicts_train = MetadataCatalog.get(DATA_REGISTER_TRAINING)\n\nprint(\"dicts training size=\", len(dataset_dicts_train))\nprint(\"################\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T02:57:14.392857Z","iopub.execute_input":"2025-01-13T02:57:14.393090Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"","metadata":{"execution":{"iopub.status.busy":"2025-01-08T16:19:24.85863Z","iopub.execute_input":"2025-01-08T16:19:24.858928Z","iopub.status.idle":"2025-01-08T16:19:24.871599Z","shell.execute_reply.started":"2025-01-08T16:19:24.858901Z","shell.execute_reply":"2025-01-08T16:19:24.870784Z"}}},{"cell_type":"code","source":"# VALIDATION: Register Validation data\nif is_train or is_evaluate:\n    if DATA_REGISTER_VALID not in DatasetCatalog.list():\n        DatasetCatalog.register(\n            DATA_REGISTER_VALID,\n            lambda: convert_coco_to_detectron2_format(\n                TRAIN_IMG_DIR,\n                train_metadata,\n                train_annot_df,\n                target_indices=valid_inds,\n        ),\n    )\n\n    # Set Validation data categories\n    MetadataCatalog.get(DATA_REGISTER_VALID).set(thing_classes=thing_classes)\n\n    dataset_dicts_valid = DatasetCatalog.get(DATA_REGISTER_VALID)\n    metadata_dicts_valid = MetadataCatalog.get(DATA_REGISTER_VALID)\n\n    print(\"dicts valid size=\", len(dataset_dicts_valid))\n    print(\"################\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# INFERENCE: Register Test Inference data\nif DATA_REGISTER_TEST not in DatasetCatalog.list():\n    print(\"Registering test data...\")\n    DatasetCatalog.register(\n        DATA_REGISTER_TEST,\n        lambda: convert_coco_to_detectron2_format(\n            TEST_IMG_DIR,\n            test_metadata,\n        )\n    )\n    print(\"Test data registered.\")\n\n# Set Test data categories\nMetadataCatalog.get(DATA_REGISTER_TEST).set(\n    thing_classes=thing_classes_test\n)\n\n# Retrieve and print test data information\ndataset_dicts_test = DatasetCatalog.get(DATA_REGISTER_TEST)\nmetadata_dicts_test = MetadataCatalog.get(DATA_REGISTER_TEST) \n\nprint(\"dicts test size=\", len(dataset_dicts_test))\nprint(\"################\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"#### DATA REGISTERED ####\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 5 Augmentation","metadata":{}},{"cell_type":"code","source":"try: \n    import utils \nexcept ModuleNotFoundError: \n    print(\"Module 'utils' not found. Installing...\") \n    import pip \n    pip.main(['install', 'utils']) \n    import utils","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def custom_mapper(dataset_dict): \n    dataset_dict = copy.deepcopy(dataset_dict) \n    image = utils.read_image(dataset_dict[\"file_name\"], format=\"BGR\") \n    \n    # Define the transformation list \n    transform_list = [T.RandomBrightness(0.8, 1.2), \n                      T.RandomFlip(prob=0.5, horizontal=False, vertical=True), \n                      T.RandomFlip(prob=0.5, horizontal=True, vertical=False) \n                     ] \n    # Apply transformations \n    image, transforms = T.apply_transform_gens(transform_list, image) \n    \n    # Convert image to float32 before applying sharpening \n    image = image.astype('float32') \n    \n    # Sharpen the image \n    kernel = np.array([[0, -1, 0], \n                       [-1, 5, -1], \n                       [0, -1, 0]]) \n    image = cv2.filter2D(src=image, ddepth=-1, kernel=kernel) \n    \n    # Convert image back to tensor \n    dataset_dict[\"image\"] = torch.as_tensor(image.transpose(2, 0, 1))\n\n    # Processing annotations (if any) and apply the transformation to them\n    if \"annotations\" in dataset_dict: # Ensuring that annotations exist\n        annos = [\n            utils.transform_instance_annotations(obj, transforms, image.shape[:2])\n            for obj in dataset_dict.pop(\"annotations\")\n            if obj.get(\"iscrowd\", 0) == 0 #Ignoring crowd annotations\n        ]\n        instances = utils.annotations_to_instances(annos, image.shape[:2])\n    \n        dataset_dict[\"instances\"] = utils.filter_empty_instances(instances)\n\n    return dataset_dict","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class AugTrainer(DefaultTrainer):\n    @classmethod\n    def build_train_loader(cls, cfg):\n        return build_detection_train_loader(cfg, mapper=custom_mapper)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 6 Hyperparameters","metadata":{}},{"cell_type":"markdown","source":"Detectron2 models need a config file to build the model. This config file has the associated hyperparameters.\n\nYou can play with this: https://detectron2.readthedocs.io/en/latest/modules/config.html#yaml-config-references","metadata":{}},{"cell_type":"code","source":"if is_train:\n    cfg = get_cfg()\n\n    # config_name = \"COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml\"\n    config_name = \"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml\"\n\n    cfg.merge_from_file(model_zoo.get_config_file(config_name))\n\n    cfg.DATASETS.TRAIN = (DATA_REGISTER_TRAINING,)\n    cfg.DATASETS.TEST = (DATA_REGISTER_VALID,)\n\n    # to evaluate during training, you have to implement `build_evaluator()` method of the trainer.\n    # https://github.com/facebookresearch/detectron2/blob/94113be6e12db36b8c7601e13747587f19ec92fe/detectron2/engine/defaults.py#L561\n    # cfg.TEST.EVAL_PERIOD = 500  #Uncomment to evaluate after every 500 iterations\n\n    cfg.DATALOADER.NUM_WORKERS = 4   # 2-> Increasing for faster data loading if possible\n\n    # cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(config_name)\n    if (is_resume_training):\n        print(\"#### SETTING PRETRAINED WEIGHTS TO RESUME TRAINING ####\")\n        cfg.MODEL.WEIGHTS = str(PRETRAINED_PATH)\n    else:\n        print(\"#### TRAINING MODEL FROM SCRATCH ####\")\n\n    cfg.SOLVER.AMP.ENABLED = True\n    cfg.SOLVER.IMS_PER_BATCH = 16  # 8-> Increasing batch size if GPU memory allows \n    cfg.SOLVER.BASE_LR = 0.0005    # 0.001-> Lowering learning rate for finer training\n\n    cfg.SOLVER.WARMUP_ITERS = 500    # 5-> Increasing iterations for longer training\n\n    # Maximum number of iterations\n    cfg.SOLVER.MAX_ITER = 10000     # 500-> Increasing iterations for longer training\n\n    # cfg.SOLVER.STEPS = (500, 1000) # must be less than MAX_ITER\n\n    cfg.SOLVER.GAMMA = 0.1  # 0.05-> Lowering decay rate for more fine-tuning\n\n    # Small value == Frequent save need a lot of storage.\n    cfg.SOLVER.CHECKPOINT_PERIOD = 500\n    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128\n    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4  # Adjusted based on the number of classes\n\n    # Create Output Directory\n    cfg.OUTPUT_DIR = str(OUTPUT_DIR)\n    print(\"creating cfg.OUTPUT_DIR -> \", cfg.OUTPUT_DIR)\n    OUTPUT_DIR.mkdir(exist_ok=True)\n\n    # Setting input image size: Adjust to higher resolution for better accuracy\n    cfg.INPUT.MIN_SIZE_TRAIN = 800    # Increasing the minimum image size for better resolution\n    cfg.INPUT.MAX_SIZE_TRAIN = 1333   # Larger max image size to capture more detail\n\n    #Applying additional data augmentation strategies (in the mapper)\n    cfg.INPUT.FORMAT = \"BGR\"\n    cfg.INPUT.CROP.ENABLED = \"True\"\n    cfg.INPUT.CROP.TYPE = \"absolute\"\n    cfg.INPUT.CROP.SIZE = [640, 640]\n    \n    # Regularization and optimization (optional for better accuracy)\n    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.05","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 7 Training","metadata":{}},{"cell_type":"markdown","source":"## 7.1 Training the model","metadata":{}},{"cell_type":"code","source":"if is_train:\n    trainer = DefaultTrainer(cfg) if not is_augment else AugTrainer(cfg)\n        \n    trainer.resume_or_load(resume=is_resume_training)\n\n    trainer.train()\n    \n    print(\"#### TRAINING COMPLETE ####\")\n    _ = trainer.model.train(False)  # turn off training\n\n    # Saving the trained model checkpoint (optional)\n    model_checkpoint_path = str(OUTPUT_MODEL)\n    torch.save(trainer.model.state_dict(), model_checkpoint_path)\n\n    # Optionally, generate a link to download the trained model\n    FileLink(model_checkpoint_path)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 7.2 Visualizing Training Metrics","metadata":{}},{"cell_type":"code","source":"if is_train:\n    # Load metrics\n    metrics_df = pd.read_json(\n        OUTPUT_DIR/\"metrics.json\", orient=\"records\", lines=True\n    )\n    mdf = metrics_df.sort_values(\"iteration\")\n    \n    print(mdf.head(10).T)\n\n    # Plot loss\n    fig, ax = plt.subplots(figsize=(10, 6))\n\n    mdf1 = mdf[~mdf[\"total_loss\"].isna()]   # Removing NaN values\n    ax.plot(mdf1[\"iteration\"], mdf1[\"total_loss\"], c=\"C0\", label=\"train\", linewidth=2)\n\n    if \"validation_loss\" in mdf.columns:\n        mdf2 = mdf[~mdf[\"validation_loss\"].isna()]\n        ax.plot(mdf2[\"iteration\"], mdf2[\"validation_loss\"], c=\"C1\", label=\"validation\", linewidth=2)\n\n    ax.set_xlabel('Iterations')\n    ax.set_ylabel('Loss')\n    ax.set_title(\"Training and Validation Loss curve\")\n    ax.legend()\n    ax.grid(True)\n    plt.show()\n\n    # Plot Accuracy\n    fig, ax = plt.subplots(figsize=(10, 6))\n\n    mdf1 = mdf[~mdf[\"fast_rcnn/cls_accuracy\"].isna()]\n    ax.plot(mdf1[\"iteration\"], mdf1[\"fast_rcnn/cls_accuracy\"], c=\"C0\", label=\"train\", linewidth=2)\n\n    ax.set_xlabel('Iterations')\n    ax.set_ylabel('Accuracy')\n    ax.set_title(\"Classification curve\")\n    ax.legend()\n    ax.grid(True)\n    plt.show()\n\n    # Plot Bounding Box regressor loss\n    fig, ax = plt.subplots(figsize=(10, 6))\n\n    mdf1 = mdf[~mdf[\"loss_box_reg\"].isna()]\n    ax.plot(mdf1[\"iteration\"], mdf1[\"loss_box_reg\"], c=\"C0\", label=\"train\", linewidth=2)\n\n    ax.set_xlabel('Iterations')\n    ax.set_ylabel('Loss')\n    ax.set_title(\"Bounding Box Regressor Loss curve\")\n    ax.legend()\n    ax.grid(True)\n    plt.show()\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 8 Evaluation","metadata":{}},{"cell_type":"markdown","source":"Can evaluate trained model on validation dataset to obtain different metric scores. ","metadata":{}},{"cell_type":"code","source":"if is_evaluate:\n    print(\"### EVALUATING ON VALIDATION DATA ####\")\n    # trained model weights\n    cfg.MODEL.WEIGHTS = str(MODEL_PATH)\n    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5   # set a custom testing threshold\n\n    cfg.SOLVER.IMS_PER_BATCH = 64\n\n    evaluator = COCOEvaluator(\n        DATA_REGISTER_VALID, \n        cfg, \n        False, \n        output_dir=cfg.OUTPUT_DIR,\n        use_fast_impl=True\n    )\n\n    val_loader = build_detection_test_loader(cfg, DATA_REGISTER_VALID)\n\n    results = inference_on_dataset(\n        trainer.model, \n        val_loader, \n        evaluator=evaluator\n    )\n\n    # Output the evaluation results\n    print(\"Evaluation completed. Results:\")\n    print(results)\n\n    # Optionally saving the evaluation results to a file\n    with open(cfg.OUTPUT_DIR / \"evaluation_results.json\", \"w\") as f:\n        json.dump(results, f)\n\n    print(f\"Mean Average Precision (mAP) Score: {results['bbox']}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 9 Inference","metadata":{}},{"cell_type":"markdown","source":"## 9.1 Setting Up Inference Model","metadata":{}},{"cell_type":"code","source":"inf_cfg = get_cfg()\n\nconfig_name = \"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml\"\n\ninf_cfg.merge_from_file(model_zoo.get_config_file(config_name))\ninf_cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128\ninf_cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4\ninf_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5\ninf_cfg.MODEL.DEVICE = \"cuda\"\n\ninf_cfg.DATALOADER.NUM_WORKERS = 2  # lower this if CUDA overflow occurs\ninf_cfg.MODEL.WEIGHTS = str(MODEL_PATH)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BATCH = 8  # lower this if CUDA overflow occurs\ntest_loader = build_detection_test_loader(inf_cfg, DATA_REGISTER_TEST, batch_size=BATCH)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"How confident should the model be for you to accept the predicted mask?","metadata":{}},{"cell_type":"code","source":"ACCEPTANCE_THRESHOLD = 0.6  # for all categories","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(f\"#### MODEL: {inf_cfg.MODEL.WEIGHTS} FOR INFERENCE ####\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 9.2 Sample Inference","metadata":{}},{"cell_type":"code","source":"predictor = DefaultPredictor(inf_cfg)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fig, ax = plt.subplots(2, 2, figsize=(20, 20))\nindices = [ax[0][0], ax[1][0], ax[0][1], ax[1][1]]\n\n# Show some qualitative results by predicting on test set images\nNUM_TEST_SAMPLES = min(4, len(dataset_dicts_test))  # Ensuring we don't exceed available samples\nsamples = np.random.choice(dataset_dicts_test, NUM_TEST_SAMPLES)\n\nfor i, sample in enumerate(samples):\n    img = cv2.imread(sample[\"file_name\"])\n    outputs = predictor(img)\n    visualizer = Visualizer(img, metadata=metadata_dicts_test, scale=0.5,)\n    visualizer = visualizer.draw_instance_predictions(outputs[\"instances\"].to(\"cpu\"))\n    display_img = visualizer.get_image()[:, :, ::-1]\n    \n    indices[i].grid(False)\n    indices[i].imshow(display_img)\n\n#Adjusting layout for better spacing of the subplots\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 9.3 Test Data Inference and Submission","metadata":{}},{"cell_type":"markdown","source":"### 9.3.1 Building Inference Model","metadata":{}},{"cell_type":"code","source":"def rebuild_model():\n    model = build_model(inf_cfg)\n    _ = DetectionCheckpointer(model).load(inf_cfg.MODEL.WEIGHTS)\n    return model","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = rebuild_model()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 9.3.2 CUDA Problems","metadata":{}},{"cell_type":"code","source":"!export LRU_CACHE_CAPACITY=1\n!export 'PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512'","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Occassionally CUDA memory overflow occurs. Trying to save as much VRAM as we can. ","metadata":{}},{"cell_type":"code","source":"vars_to_del = [\"trainer\", \"predictor\", \"outputs\"]\n\nfor v in vars_to_del:\n    if v in globals():\n        print(f\"Deleting {v} from globals\")\n        del globals()[v]\n    elif v in locals():\n        print(f\"Deleting {v} from locals\")\n        del locals()[v]\n\n#import gc\n#gc.collect()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 9.3.3 Inference Utils","metadata":{}},{"cell_type":"code","source":"def rle_encode(mask):\n    pixels = mask.T.flatten()\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    return ' '.join(str(x) for x in rle)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"@retry_if_cuda_oom\ndef get_masks(prediction):\n    # get masks for each category\n    take = prediction.scores >= ACCEPTANCE_THRESHOLD\n    pred_masks = (prediction.pred_masks[take] != 0)\n    pred_classes = prediction.pred_classes[take]\n  \n    rles = []\n    for cat in range(len(thing_classes)):\n        pred_mask = pred_masks[pred_classes == cat]\n        \n        # pred_mask = retry_if_cuda_oom(torch.any)(pred_mask, dim=0)\n        pred_mask = torch.any(pred_mask, dim=0)\n        rles.append(rle_encode(pred_mask.short().to(\"cpu\").numpy()))\n\n    return rles","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def run_inference(data):\n    results = []\n    with torch.no_grad():\n        outputs = model(data)\n        if torch.cuda.is_available():\n            torch.cuda.synchronize()\n\n        for idx, output in enumerate(outputs):\n            output = output[\"instances\"]\n\n            rles = get_masks(output)\n\n            result = [\n                f\"{data[idx]['image_id']}_{cat},{rles[cat]}\\n\"\n                for cat in range(len(thing_classes))\n            ]\n\n            results.extend(result)\n\n        del outputs, output\n\n    #gc.collect()\n\n    return results","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 9.3.4 Running Inference on Test Data and Creating Submission File\nThis is super slow. ","metadata":{}},{"cell_type":"code","source":"print(\"#### RUNNING INFERENCE ON TEST DATA ####\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"torch.cuda.empty_cache()\ngc.collect()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if is_inference:\n    model.eval()\n    submission_file = open(\"submission.csv\", \"w\")\n    submission_file.write(\"Id,Predicted\\n\")\n\n    results: list[str] = []\n    \n    for i, data in enumerate(tqdm(test_loader)):\n        res = run_inference(data)\n        results.extend(res)\n        \n        if i % (500 // BATCH) == 0:\n            print(f\"Inference on batch {i}/{len(test_loader)} done\")\n            submission_file.writelines(results)\n            results = []\n\n        # Clearing memory in each iteration\n        torch.cuda.empty_cache()\n        gc.collect()\n\n    submission_file.writelines(results)\n    submission_file.close()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if Path(\"submission.csv\").exists:\n    display(FileLink(\"submission.csv\"))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!rm -r detectron2/","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}