{"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":"#### This notebook was inspired from [here](https://www.kaggle.com/code/ammarnassanalhajali/layout-parser-model-training)","metadata":{"_uuid":"61741213-ab72-4bf7-abc7-d916b4ae2e25","_cell_guid":"40cb089f-e96a-4a5d-b258-6a60ddab9518","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"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":{"_uuid":"73c5592d-cbc5-4b82-b91a-7a65caabb8a7","_cell_guid":"68edabe4-6f96-4db1-8dfb-e33cf3b39236","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"markdown","source":"# 1 Install detectron2","metadata":{"_uuid":"681daf82-4cad-4788-b35b-6e2223c67a1a","_cell_guid":"90eaeabe-0436-45af-9e11-20a7dd7f1acb","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"markdown","source":"## 1.1 Recommended Way (is not working on kaggle)","metadata":{"_uuid":"0b20c7d9-645b-441c-b023-57b35cbe1c82","_cell_guid":"a9b70df9-fffe-4a27-9b1d-88d998a48aca","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"# !python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'","metadata":{"_uuid":"3d48a95b-49e6-4192-b27e-727af0b43d0e","_cell_guid":"05e5faba-91ed-4bb6-96bc-315e301e053c","trusted":true,"collapsed":false,"_kg_hide-input":false,"_kg_hide-output":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:20.602692Z","iopub.execute_input":"2025-10-25T21:22:20.602929Z","iopub.status.idle":"2025-10-25T21:22:20.607173Z","shell.execute_reply.started":"2025-10-25T21:22:20.602905Z","shell.execute_reply":"2025-10-25T21:22:20.606317Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 1.2 Fast Way\nIgnore the warnings.","metadata":{"_uuid":"145ff485-77d1-44c0-8c34-63266e72696a","_cell_guid":"1d42000c-c7a0-44d2-bf65-1d2970e5a2c2","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"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":{"_uuid":"72a17d28-d249-43ec-94ee-6d331d12f29e","_cell_guid":"edb14357-ce62-48bc-83a4-0998806fde4a","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:20.610946Z","iopub.execute_input":"2025-10-25T21:22:20.611232Z","iopub.status.idle":"2025-10-25T21:22:31.784246Z","shell.execute_reply.started":"2025-10-25T21:22:20.611202Z","shell.execute_reply":"2025-10-25T21:22:31.782802Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 2 Notebook Config","metadata":{"_uuid":"9fa0745e-3399-4446-8ca1-3d8f68d8694d","_cell_guid":"b789762c-cd36-46f3-b7e0-33cedfa38d18","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"markdown","source":"## 2.1 Decisions","metadata":{"_uuid":"8f7d9a9f-84b1-46f3-aba5-055370d810da","_cell_guid":"7534f8d7-6674-4a19-b390-b777b61c47ee","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"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 = True\n\nSEED = int(datetime.now().timestamp())","metadata":{"_uuid":"16cb65ed-eecd-4405-8873-5479a1fcbcce","_cell_guid":"0d572225-9c02-4413-a945-010354bf5581","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:31.786383Z","iopub.execute_input":"2025-10-25T21:22:31.788146Z","iopub.status.idle":"2025-10-25T21:22:31.792904Z","shell.execute_reply.started":"2025-10-25T21:22:31.788100Z","shell.execute_reply":"2025-10-25T21:22:31.792033Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2.2 Paths","metadata":{"_uuid":"c03b5713-9dd2-4930-b793-6e5ca8d0fc58","_cell_guid":"4ccb1a20-c713-4b69-8333-355370ffe9e3","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"!ls /kaggle/input\n!ls /kaggle/input/dlsprint2\n!ls /kaggle/input/dlsprint2/badlad\n!ls /kaggle/input/dlsprint2/badlad/labels\n!ls /kaggle/input/dlsprint2/badlad/labels/coco_format/train","metadata":{"_uuid":"74ad507d-fd14-4a45-8b7d-10ec25bf103b","_cell_guid":"861a1c41-0fb2-410e-8cb3-3716ff7d3ffa","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:31.794201Z","iopub.execute_input":"2025-10-25T21:22:31.794534Z","iopub.status.idle":"2025-10-25T21:22:36.800363Z","shell.execute_reply.started":"2025-10-25T21:22:31.794502Z","shell.execute_reply":"2025-10-25T21:22:36.799211Z"}},"outputs":[],"execution_count":null},{"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":{"_uuid":"713109da-2374-488a-91a1-43946a07fc53","_cell_guid":"44c8834b-1f70-44c2-96d7-22a288e61d5a","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:36.802145Z","iopub.execute_input":"2025-10-25T21:22:36.803073Z","iopub.status.idle":"2025-10-25T21:22:36.808619Z","shell.execute_reply.started":"2025-10-25T21:22:36.803031Z","shell.execute_reply":"2025-10-25T21:22:36.807604Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Model path based on Decisions\nMODEL_PATH = OUTPUT_MODEL if is_train else PRETRAINED_PATH","metadata":{"_uuid":"f8dc9fb4-84c8-4160-8646-20b1dbacf300","_cell_guid":"72a6f2d8-51e0-4dd0-aaa0-4ee948294fcb","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:36.809729Z","iopub.execute_input":"2025-10-25T21:22:36.809961Z","iopub.status.idle":"2025-10-25T21:22:36.822798Z","shell.execute_reply.started":"2025-10-25T21:22:36.809941Z","shell.execute_reply":"2025-10-25T21:22:36.822099Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2.3 imports","metadata":{"_uuid":"396ce18f-8e59-4e4a-ab8d-843574616986","_cell_guid":"87489e30-5b4a-4496-80ff-1ff5b110021a","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"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\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\nfrom pathlib import Path\nfrom detectron2.data.datasets import load_coco_json\n\nimport warnings\n# Ignore \"future\" warnings and Data-Frame-Slicing warnings.\nwarnings.filterwarnings('ignore')\n\nsetup_logger()","metadata":{"_uuid":"36e1589a-6196-476f-ac29-3f3170dbd807","_cell_guid":"b9c4cb30-44f0-4406-bd8a-df7d9ae3002b","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:36.823998Z","iopub.execute_input":"2025-10-25T21:22:36.824197Z","iopub.status.idle":"2025-10-25T21:22:37.150248Z","shell.execute_reply.started":"2025-10-25T21:22:36.824180Z","shell.execute_reply":"2025-10-25T21:22:37.149368Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 3 COCO Annotations Data","metadata":{"_uuid":"336306fd-9325-4a09-b010-828c7c199e82","_cell_guid":"c2e5a395-9b66-408d-8e2c-0b9247e3503e","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"markdown","source":"## 3.1 Load","metadata":{"_uuid":"fda5920f-7c52-426a-b430-44b32e2ffa6d","_cell_guid":"c88ab989-a30a-460f-a75c-4774c4acc489","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"ROOT_PATH = Path(\"/kaggle/input/dlsprint2/badlad\")\nTRAIN_PATH = ROOT_PATH  # json files are in data/\n#TRAIN_IMG_DIR = ROOT_PATH / \"train_images\"","metadata":{"_uuid":"65010308-5096-4c5a-bdc1-52cfaeade879","_cell_guid":"8af76877-2aa6-41f8-b129-fc75a94ca722","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:37.151233Z","iopub.execute_input":"2025-10-25T21:22:37.151502Z","iopub.status.idle":"2025-10-25T21:22:37.155688Z","shell.execute_reply.started":"2025-10-25T21:22:37.151466Z","shell.execute_reply":"2025-10-25T21:22:37.154727Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TRAIN_IMG_DIR = ROOT_PATH / \"images\" / \"train\"\nTEST_IMG_DIR = ROOT_PATH / \"images\" / \"test\"\n\nTRAIN_COCO_PATH = ROOT_PATH / \"labels\" / \"coco_format\" / \"train\" / \"badlad-train-coco.json\"\n#If you don’t have a validation JSON, just skip for now or create one later\n#VAL_COCO_PATH = ROOT_PATH / \"labels\" / \"coco_format\" / \"valid\" / \"badlad-valid-coco.json\"\nTEST_METADATA_PATH = Path(\"/kaggle/input/dlsprint2/badlad/badlad-test-metadata.json\")","metadata":{"_uuid":"f2132165-4967-406c-8c55-6ed40e0c28ea","_cell_guid":"937d92dd-53f8-4051-8270-4cf9f773b5a6","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:37.156807Z","iopub.execute_input":"2025-10-25T21:22:37.157160Z","iopub.status.idle":"2025-10-25T21:22:37.166161Z","shell.execute_reply.started":"2025-10-25T21:22:37.157137Z","shell.execute_reply":"2025-10-25T21:22:37.165206Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Define paths to annotation files\nTRAIN_COCO_PATH = TRAIN_PATH / \"/kaggle/input/dlsprint2/badlad/images/train.json\"\nVAL_COCO_PATH = TRAIN_PATH / \"valid.json\"\nTEST_METADATA_PATH = TRAIN_PATH / \"test.json\"\n\n# Define image directories\nTRAIN_IMG_DIR = TRAIN_PATH / \"train\"\nTEST_IMG_DIR = TRAIN_PATH / \"test\"","metadata":{"_uuid":"d9793667-0642-4d2f-9e16-15557e50de52","_cell_guid":"c1af1da4-9fb9-4a3a-8616-8da5cf624bcb","trusted":true,"collapsed":false,"execution":{"iopub.status.busy":"2025-07-30T11:23:37.639862Z","iopub.execute_input":"2025-07-30T11:23:37.6402Z","iopub.status.idle":"2025-07-30T11:23:37.644942Z","shell.execute_reply.started":"2025-07-30T11:23:37.640169Z","shell.execute_reply":"2025-07-30T11:23:37.643971Z"},"jupyter":{"outputs_hidden":false}}},{"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":{"_uuid":"b1b6b32d-05fa-4116-84d1-0ef21c7997ca","_cell_guid":"0a9eac7a-e66a-4ae8-825c-637bb80d1fba","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:37.167146Z","iopub.execute_input":"2025-10-25T21:22:37.167405Z","iopub.status.idle":"2025-10-25T21:22:41.706008Z","shell.execute_reply.started":"2025-10-25T21:22:37.167384Z","shell.execute_reply":"2025-10-25T21:22:41.705119Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3.2 Observe","metadata":{"_uuid":"58b9d575-868b-492f-a56c-1cca94bcb452","_cell_guid":"2d39bc5c-3840-4c2a-b5ab-50f8a6540b53","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"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":{"_uuid":"4f7d70ab-f7d8-463a-985b-28a4806b6384","_cell_guid":"606fba88-e7ed-48fd-aac4-e0a21bcada01","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:41.707188Z","iopub.execute_input":"2025-10-25T21:22:41.707567Z","iopub.status.idle":"2025-10-25T21:22:41.714294Z","shell.execute_reply.started":"2025-10-25T21:22:41.707535Z","shell.execute_reply":"2025-10-25T21:22:41.713389Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#from detectron2.data import DatasetCatalog, MetadataCatalog\n#from detectron2.data.datasets import load_coco_json\n\nclass_names = [\"paragraph\", \"text_box\", \"image\", \"table\"]  # Adjust based on your dataset\n\nDatasetCatalog.register(\"badlad_train\", lambda: load_coco_json(str(TRAIN_COCO_PATH), TRAIN_IMG_DIR, \"badlad_train\"))\nMetadataCatalog.get(\"badlad_train\").set(\n    thing_classes=class_names,\n    evaluator_type=\"coco\",\n    json_file=str(TRAIN_COCO_PATH),\n    image_root=TRAIN_IMG_DIR\n)\n\nDatasetCatalog.register(\"badlad_valid\", lambda: load_coco_json(str(VAL_COCO_PATH), TRAIN_IMG_DIR, \"badlad_valid\"))\nMetadataCatalog.get(\"badlad_valid\").set(\n    thing_classes=class_names,\n    evaluator_type=\"coco\",\n    json_file=str(VAL_COCO_PATH),\n    image_root=TRAIN_IMG_DIR\n)\n\nDatasetCatalog.register(\"badlad_test\", lambda: load_coco_json(str(TEST_METADATA_PATH), TEST_IMG_DIR, \"badlad_test\"))\nMetadataCatalog.get(\"badlad_test\").set(\n    thing_classes=class_names,\n    evaluator_type=\"coco\",\n    json_file=str(TEST_METADATA_PATH),\n    image_root=TEST_IMG_DIR\n)\n\nprint(\"#### COCO DATASETS REGISTERED ####\")","metadata":{"_uuid":"372510a3-fdd1-453e-b424-71cf4fe8b171","_cell_guid":"f12e9ac3-88df-4007-b7f1-92574ca6e2ab","trusted":true,"collapsed":false,"execution":{"iopub.status.busy":"2025-07-30T11:05:39.796849Z","iopub.execute_input":"2025-07-30T11:05:39.797554Z","iopub.status.idle":"2025-07-30T11:05:39.804312Z","shell.execute_reply.started":"2025-07-30T11:05:39.797518Z","shell.execute_reply":"2025-07-30T11:05:39.803341Z"},"jupyter":{"outputs_hidden":false}}},{"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":{"_uuid":"4bdf4d20-2238-454f-85d1-08658b30ab86","_cell_guid":"ea9f4380-577b-497e-8f29-eb15ec990301","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:41.715350Z","iopub.execute_input":"2025-10-25T21:22:41.715931Z","iopub.status.idle":"2025-10-25T21:22:42.794240Z","shell.execute_reply.started":"2025-10-25T21:22:41.715906Z","shell.execute_reply":"2025-10-25T21:22:42.793201Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"These are the categories we are going to detect.","metadata":{"_uuid":"2a466f99-7bc7-436b-bb61-619df114cb78","_cell_guid":"70864998-d5c2-4fd8-994d-6e375d7e790c","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"print(thing_classes)","metadata":{"_uuid":"68b35d4e-baab-45ed-aad9-fc03c7bf7d6b","_cell_guid":"f630c3cb-2295-42d5-b45e-c0390538698d","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:42.795325Z","iopub.execute_input":"2025-10-25T21:22:42.795603Z","iopub.status.idle":"2025-10-25T21:22:42.800185Z","shell.execute_reply.started":"2025-10-25T21:22:42.795581Z","shell.execute_reply":"2025-10-25T21:22:42.799325Z"}},"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":{"_uuid":"33f948ec-576d-4050-9751-03a95bd6fa06","_cell_guid":"bcbc02aa-1f00-48e4-9190-1f08f4dc8ba5","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:42.801099Z","iopub.execute_input":"2025-10-25T21:22:42.801293Z","iopub.status.idle":"2025-10-25T21:22:42.852740Z","shell.execute_reply.started":"2025-10-25T21:22:42.801277Z","shell.execute_reply":"2025-10-25T21:22:42.851862Z"}},"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":{"_uuid":"12d9b2f7-9433-43d3-bd94-af54fa6d7410","_cell_guid":"b3ad8c10-12fd-4c2f-9bf5-b89828ff43bf","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:42.853853Z","iopub.execute_input":"2025-10-25T21:22:42.854080Z","iopub.status.idle":"2025-10-25T21:22:43.720116Z","shell.execute_reply.started":"2025-10-25T21:22:42.854060Z","shell.execute_reply":"2025-10-25T21:22:43.719203Z"}},"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":{"_uuid":"5532993e-4d23-4595-8471-a11618b682bf","_cell_guid":"99332a9d-c016-4e44-8acc-e3eecc0f154c","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"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":{"_uuid":"cae92904-d9c7-43fe-9a53-5f44c40d6fc5","_cell_guid":"601b5fe2-db0b-4f9e-a96f-6d6f9e3a68d4","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:43.721182Z","iopub.execute_input":"2025-10-25T21:22:43.721518Z","iopub.status.idle":"2025-10-25T21:22:43.751402Z","shell.execute_reply.started":"2025-10-25T21:22:43.721490Z","shell.execute_reply":"2025-10-25T21:22:43.750502Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 4 Preparing Data for Training","metadata":{"_uuid":"45ba1be2-d935-4cd2-af52-6981b8a7cddf","_cell_guid":"af841e7a-685b-41c0-bcdb-7655763db59e","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"markdown","source":"## 4.1 Train-Validation Split","metadata":{"_uuid":"0dc60e25-d846-42ac-a874-3efbb4a80db8","_cell_guid":"7fff0da9-90aa-43d6-9aeb-aff41d10c4d7","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"TRAIN_SPLIT = 0.95","metadata":{"_uuid":"629d875b-f1bf-4733-a9d6-55cc71ca1849","_cell_guid":"99006b54-9f12-4afa-94b4-d765ee2560aa","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:43.752719Z","iopub.execute_input":"2025-10-25T21:22:43.752976Z","iopub.status.idle":"2025-10-25T21:22:43.757659Z","shell.execute_reply.started":"2025-10-25T21:22:43.752943Z","shell.execute_reply":"2025-10-25T21:22:43.756771Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"n_dataset = len(train_metadata)\nn_train = int(n_dataset * TRAIN_SPLIT)\nprint(\"n_dataset\", n_dataset, \"n_train\", n_train, \"n_val\", n_dataset-n_train)\n\nnp.random.seed(SEED)\n\ninds = np.random.permutation(n_dataset)\ntrain_inds, valid_inds = inds[:n_train], inds[n_train:]","metadata":{"_uuid":"4d54d479-0bbf-44d5-9462-740c2e7bf9d8","_cell_guid":"b17242eb-67b7-4426-9bfc-c021ff0c05f3","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:43.758714Z","iopub.execute_input":"2025-10-25T21:22:43.759062Z","iopub.status.idle":"2025-10-25T21:22:43.769337Z","shell.execute_reply.started":"2025-10-25T21:22:43.759015Z","shell.execute_reply":"2025-10-25T21:22:43.768603Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4.2 Formatting Data for `detectron2`","metadata":{"_uuid":"0f5d46ae-c0b6-4504-84c9-e847b65e23df","_cell_guid":"db826b3c-071c-45f1-8441-636c14f96fc6","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"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    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.values\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                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":{"_uuid":"08819df1-da88-44b5-804e-4f6b98ca571c","_cell_guid":"1105e191-fabb-4530-91e9-b31883cb8d76","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:43.770479Z","iopub.execute_input":"2025-10-25T21:22:43.770743Z","iopub.status.idle":"2025-10-25T21:22:43.781011Z","shell.execute_reply.started":"2025-10-25T21:22:43.770721Z","shell.execute_reply":"2025-10-25T21:22:43.780160Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4.3 Registering and Loading Data for `detectron2`","metadata":{"_uuid":"878bdaab-09c3-4d24-9d0f-70644961c4b9","_cell_guid":"8992a8c4-4b57-4c05-bf8c-c455bcf4560a","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"DATA_REGISTER_TRAINING = \"badlad_train\"\nDATA_REGISTER_VALID    = \"badlad_valid\"\nDATA_REGISTER_TEST     = \"badlad_test\"","metadata":{"_uuid":"42b0a543-94a4-4c79-979c-73eefd37c13b","_cell_guid":"2f46ae12-e071-41db-8e9b-8e85b274c5a3","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:43.782009Z","iopub.execute_input":"2025-10-25T21:22:43.782255Z","iopub.status.idle":"2025-10-25T21:22:43.794597Z","shell.execute_reply.started":"2025-10-25T21:22:43.782235Z","shell.execute_reply":"2025-10-25T21:22:43.793730Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Register Training data\nif is_train:\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\n    # Set Training data categories\n    MetadataCatalog.get(DATA_REGISTER_TRAINING).set(thing_classes=thing_classes)\n\n    dataset_dicts_train = DatasetCatalog.get(DATA_REGISTER_TRAINING)\n    metadata_dicts_train = MetadataCatalog.get(DATA_REGISTER_TRAINING)\n\n    print(\"dicts training size=\", len(dataset_dicts_train))\n    print(\"################\")","metadata":{"_uuid":"46a56dde-16b5-4106-9e2c-e8bfd433d9d1","_cell_guid":"3ca22006-56e9-4647-b72c-b1875e36be15","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:22:43.796262Z","iopub.execute_input":"2025-10-25T21:22:43.796536Z","iopub.status.idle":"2025-10-25T21:23:58.996283Z","shell.execute_reply.started":"2025-10-25T21:22:43.796516Z","shell.execute_reply":"2025-10-25T21:23:58.995353Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Register Validation data\nif is_train or is_evaluate:\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":{"_uuid":"53c9de21-9a5d-49bf-870c-abb55a79f8cc","_cell_guid":"95b1381d-e351-4ef3-92c5-37f027cd3cf7","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:23:58.997438Z","iopub.execute_input":"2025-10-25T21:23:58.997730Z","iopub.status.idle":"2025-10-25T21:25:14.772333Z","shell.execute_reply.started":"2025-10-25T21:23:58.997707Z","shell.execute_reply":"2025-10-25T21:25:14.771334Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Register Test Inference data\nDatasetCatalog.register(\n    DATA_REGISTER_TEST,\n    lambda: convert_coco_to_detectron2_format(\n        TEST_IMG_DIR,\n        test_metadata,\n    )\n)\n\n# Set Test data categories\nMetadataCatalog.get(DATA_REGISTER_TEST).set(\n    thing_classes=thing_classes_test\n)\n\ndataset_dicts_test = DatasetCatalog.get(DATA_REGISTER_TEST)\ndataset_dicts_test = dataset_dicts_test[:2]  # Limit to 2 samples\n\nmetadata_dicts_test = MetadataCatalog.get(DATA_REGISTER_TEST)","metadata":{"_uuid":"4bd1b142-0536-477f-872b-3fb2b85531a3","_cell_guid":"554bc888-9eae-44f1-83a2-7f408d4df0ca","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:25:14.776648Z","iopub.execute_input":"2025-10-25T21:25:14.776916Z","iopub.status.idle":"2025-10-25T21:25:15.350490Z","shell.execute_reply.started":"2025-10-25T21:25:14.776893Z","shell.execute_reply":"2025-10-25T21:25:15.349709Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"#### DATA REGISTERED ####\")","metadata":{"_uuid":"4c5eb241-010f-45c9-8a3b-8356ba3503c4","_cell_guid":"7d7f7e3d-42cb-40dd-85b7-1a13205e3d07","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:25:15.351420Z","iopub.execute_input":"2025-10-25T21:25:15.351689Z","iopub.status.idle":"2025-10-25T21:25:15.355742Z","shell.execute_reply.started":"2025-10-25T21:25:15.351661Z","shell.execute_reply":"2025-10-25T21:25:15.354891Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 5 Augmentation","metadata":{"_uuid":"19d26691-2cb8-4be2-8a14-34204b23ae0f","_cell_guid":"547f2c79-64b6-45ad-870d-a7b6badfa6a5","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"markdown","source":"# Log Transform Function","metadata":{"_uuid":"9f5a1642-fa9d-4534-9cd3-a9fd15a83e79","_cell_guid":"074c5eac-5b99-4da2-9c99-69bb44cd3a19","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"def log_transform_image(image: np.ndarray) -> np.ndarray:\n    # Convert image to float to prevent log(0) errors\n    image = image.astype(np.float32)\n    \n    # Calculate log transform\n    c = 255 / np.log(1 + np.max(image))  # scaling constant\n    log_image = c * np.log(1 + image)\n    \n    # Convert back to uint8 format\n    return np.uint8(log_image)","metadata":{"_uuid":"6834c236-b3fb-460d-9113-e2482a39e842","_cell_guid":"e937b5fe-6692-4115-96db-31941f73c92a","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:25:15.356812Z","iopub.execute_input":"2025-10-25T21:25:15.357049Z","iopub.status.idle":"2025-10-25T21:25:15.368218Z","shell.execute_reply.started":"2025-10-25T21:25:15.357028Z","shell.execute_reply":"2025-10-25T21:25:15.367583Z"}},"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     # Applying Log Transform here\n    image = log_transform_image(image)\n\n    # Normalizing to 0–1 range\n    image = image.astype(np.float32) / 255.0\n    \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    image, transforms = T.apply_transform_gens(transform_list, image)\n\n    dataset_dict[\"image\"] = torch.as_tensor(\n        image.transpose(2, 0, 1).astype(\"float32\"))\n\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\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":{"_uuid":"cadf45f1-fed2-40dd-b2ea-754df6623173","_cell_guid":"c3ee7ccf-1bb9-407f-bd99-4bdef3ec1a73","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:25:15.369231Z","iopub.execute_input":"2025-10-25T21:25:15.369621Z","iopub.status.idle":"2025-10-25T21:25:15.383572Z","shell.execute_reply.started":"2025-10-25T21:25:15.369598Z","shell.execute_reply":"2025-10-25T21:25:15.382651Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"from detectron2.engine import DefaultTrainer","metadata":{"_uuid":"d4b7b499-14d8-45fc-b812-1fb49d62feb3","_cell_guid":"22a6c681-b25e-46e9-ac0c-d4163797e9ba","trusted":true,"collapsed":false,"execution":{"iopub.status.busy":"2025-07-30T16:14:50.604733Z","iopub.execute_input":"2025-07-30T16:14:50.605581Z","iopub.status.idle":"2025-07-30T16:14:50.609336Z","shell.execute_reply.started":"2025-07-30T16:14:50.60555Z","shell.execute_reply":"2025-07-30T16:14:50.60841Z"},"jupyter":{"outputs_hidden":false}}},{"cell_type":"markdown","source":"class AugTrainer(DefaultTrainer):\n    @classmethod\n    def build_train_loader(cls, cfg):\n        return build_detection_train_loader(cfg, mapper=custom_mapper)","metadata":{"_uuid":"5671894d-91ef-4b9f-bc00-039d65c3a8ba","_cell_guid":"11310e07-8805-4d83-9804-9e2af9b3b873","trusted":true,"collapsed":false,"execution":{"iopub.status.busy":"2025-07-30T16:14:58.443931Z","iopub.execute_input":"2025-07-30T16:14:58.444653Z","iopub.status.idle":"2025-07-30T16:14:58.448795Z","shell.execute_reply.started":"2025-07-30T16:14:58.444619Z","shell.execute_reply":"2025-07-30T16:14:58.44785Z"},"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"# Visual Check for Log Transform\n\n#import matplotlib.pyplot as plt\n#import cv2\n\n# Show a sample image before and after log transform\nimg_path = str(TRAIN_IMG_DIR / train_metadata.iloc[0][\"file_name\"])\noriginal = cv2.imread(img_path)\nlog_img = log_transform_image(original)\n\nplt.figure(figsize=(10, 5))\nplt.subplot(1, 2, 1)\nplt.imshow(cv2.cvtColor(original, cv2.COLOR_BGR2RGB))\nplt.title(\"Original Image\")\n\nplt.subplot(1, 2, 2)\nplt.imshow(cv2.cvtColor(log_img, cv2.COLOR_BGR2RGB))\nplt.title(\"Log Transformed Image\")\nplt.show()","metadata":{"_uuid":"3b870bbc-acad-422f-ba62-645a790cae9b","_cell_guid":"e18e8e33-8150-4c58-b5c6-1f4c5bef2ea4","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:25:15.384689Z","iopub.execute_input":"2025-10-25T21:25:15.385240Z","iopub.status.idle":"2025-10-25T21:25:15.883711Z","shell.execute_reply.started":"2025-10-25T21:25:15.385217Z","shell.execute_reply":"2025-10-25T21:25:15.882778Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 6 Hyperparameters","metadata":{"_uuid":"de574c9e-a382-4bea-abdb-adf4f1dd8107","_cell_guid":"ec1ac112-74f1-48e5-bb0f-fe71109ab738","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"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":{"_uuid":"dd667000-6f1e-4952-aeaf-28f5d74113db","_cell_guid":"fd4f51e8-0a7d-4224-a360-74a2c863fb08","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"markdown","source":"# Seting alloc config for PyTorch (add this before training starts)\nimport os\nos.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"max_split_size_mb:64\"   \n\nimport torch\ngc.collect()\ntorch.cuda.empty_cache()\n\ntorch.backends.cudnn.benchmark = False\n\n# Optional: Cap memory usage per process\ntorch.cuda.set_per_process_memory_fraction(0.8, device=0)","metadata":{"execution":{"iopub.status.busy":"2025-10-25T13:08:48.566706Z","iopub.execute_input":"2025-10-25T13:08:48.567029Z","iopub.status.idle":"2025-10-25T13:08:49.734317Z","shell.execute_reply.started":"2025-10-25T13:08:48.567005Z","shell.execute_reply":"2025-10-25T13:08:49.733317Z"}}},{"cell_type":"markdown","source":"config_name = \"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml\"","metadata":{}},{"cell_type":"code","source":"# Seting alloc config for PyTorch (add this before training starts)\nimport os\nos.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"max_split_size_mb:32\"  # 32 -> 16   \n\nimport torch\n#gc.collect()\n#torch.cuda.empty_cache()\n\ntorch.backends.cudnn.benchmark = False\n\n# Optional: Cap memory usage per process\ntorch.cuda.set_per_process_memory_fraction(0.8, device=0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T21:28:40.318761Z","iopub.execute_input":"2025-10-25T21:28:40.319421Z","iopub.status.idle":"2025-10-25T21:28:40.324242Z","shell.execute_reply.started":"2025-10-25T21:28:40.319386Z","shell.execute_reply":"2025-10-25T21:28:40.323430Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if is_train:\n    cfg = get_cfg()\n\n    config_name = \"COCO-Detection/faster_rcnn_R_50_FPN_1x.yaml\"\n    #config_name = \"COCO-InstanceSegmentation/mask_rcnn_R_50_C4_1x.yaml\"\n   # config_name = \"COCO-InstanceSegmentation/mask_rcnn_R_18_FPN_1x.yaml\"\n\n    # Override backbone to ResNet-18\n    cfg.MODEL.BACKBONE.NAME = \"build_resnet_backbone\"\n    cfg.MODEL.RESNETS.DEPTH = 18\n    cfg.MODEL.WEIGHTS = \"\"  # No pretrained weights\n\n    \n    cfg.merge_from_file(model_zoo.get_config_file(config_name))\n\n    # Gradiant Checkpoint\n    cfg.MODEL.GRADIENT_CHECKPOINTING = True\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\n\n    cfg.DATALOADER.NUM_WORKERS = 0     # 2 -> 0\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   # True -> False\n    cfg.SOLVER.AMP.LOSS_SCALE = \"dynamic\"\n\n    cfg.SOLVER.IMS_PER_BATCH = 1   # 8 -> 2 -> 1\n    cfg.SOLVER.OPTIMIZER = \"AdamW\"\n    cfg.SOLVER.BASE_LR = 0.0001  # Lower LR for AdamW\n\n    #cfg.SOLVER.BASE_LR = 0.001\n\n    # Gradient clipping to reduce memory spikes\n    cfg.SOLVER.CLIP_GRADIENTS.ENABLED = False      # True -> False\n    cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE = \"value\"\n    cfg.SOLVER.CLIP_GRADIENTS.CLIP_VALUE = 0.5\n\n    cfg.SOLVER.WARMUP_ITERS = 5\n\n    # Maximum number of iterations\n    cfg.SOLVER.MAX_ITER = 10000       # 500 -> 10000\n\n    # cfg.SOLVER.STEPS = (500, 1000) # must be less than MAX_ITER\n\n    cfg.SOLVER.GAMMA = 0.05\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 = 1  # 128 -> 64 -> 32 -> 16 -> 8 -> 4 -> 2 ->\n    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4\n\n    # Reducing proposals — useful for memory\n    cfg.MODEL.RPN.POST_NMS_TOPK_TRAIN = 200     # 1000 -> 500 -> 200\n    cfg.MODEL.RPN.POST_NMS_TOPK_TEST = 100       # 500 -> 300 -> 100\n\n    # Limiting detections per image during test time\n    cfg.TEST.DETECTIONS_PER_IMAGE = 10         # 50 -> 25 -> 10\n\n    # Optional image size limit to reduce memory\n    cfg.INPUT.MIN_SIZE_TRAIN = (80,)        # 256 -> 224 -> 128 -> 96 -> 80\n    cfg.INPUT.MAX_SIZE_TRAIN = 80          # 480 -> 320 -> 256 -> 128 -> 96 -> 80\n\n    # Enabling AMP (Mixed Precision) to save GPU memory\n    cfg.SOLVER.AMP.ENABLED = True\n\n    # Seting alloc config for PyTorch (add this before training starts)\n    #import os\n    #os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"max_split_size_mb:32\"   \n\n    #import torch\n    #gc.collect()\n    #torch.cuda.empty_cache()\n\n    torch.backends.cudnn.benchmark = False\n\n    # Optional: Cap memory usage per process\n    #torch.cuda.set_per_process_memory_fraction(0.8, device=0)\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)","metadata":{"_uuid":"e76ba724-bf9f-488f-a946-04e10b29c3b2","_cell_guid":"ba66e9b5-9d6c-495c-b425-0feaaa952682","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:28:43.197002Z","iopub.execute_input":"2025-10-25T21:28:43.197890Z","iopub.status.idle":"2025-10-25T21:28:43.219224Z","shell.execute_reply.started":"2025-10-25T21:28:43.197856Z","shell.execute_reply":"2025-10-25T21:28:43.218420Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"try:\n    is_train\nexcept NameError:\n    is_train = True\n    is_resume_training = False\n    is_augment = True","metadata":{"_uuid":"5f5f3d05-80b0-43e7-acce-eb1a9b115f1f","_cell_guid":"425ca2b2-9cb4-40c0-bebf-8477239e2539","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:28:47.793090Z","iopub.execute_input":"2025-10-25T21:28:47.793422Z","iopub.status.idle":"2025-10-25T21:28:47.797708Z","shell.execute_reply.started":"2025-10-25T21:28:47.793394Z","shell.execute_reply":"2025-10-25T21:28:47.796804Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#import gc\n#gc.collect()\n#torch.cuda.empty_cache()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T21:15:05.321635Z","iopub.execute_input":"2025-10-25T21:15:05.321876Z","iopub.status.idle":"2025-10-25T21:15:05.331757Z","shell.execute_reply.started":"2025-10-25T21:15:05.321856Z","shell.execute_reply":"2025-10-25T21:15:05.330960Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#torch.cuda.reset_peak_memory_stats()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T21:15:05.332691Z","iopub.execute_input":"2025-10-25T21:15:05.332929Z","iopub.status.idle":"2025-10-25T21:15:05.340122Z","shell.execute_reply.started":"2025-10-25T21:15:05.332906Z","shell.execute_reply":"2025-10-25T21:15:05.339502Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 7 Training","metadata":{"_uuid":"981a13ca-4ff5-4be6-9851-cdd939069597","_cell_guid":"84db7e17-74f9-4b6b-9de2-381f600dc5be","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"markdown","source":"## 7.1 Training the model","metadata":{"_uuid":"a0a243c3-ad46-49d6-b577-04047ed4c227","_cell_guid":"c1051d3c-38a1-47c1-8ee1-0ae73178cd3c","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"markdown","source":"from detectron2.engine import DefaultTrainer\n\nclass AugTrainer(DefaultTrainer):\n    @classmethod\n    def build_model(cls, cfg):\n        model = build_model(cfg)\n        model.backbone.bottom_up.use_checkpoint = True  # ✅ This is the correct place\n        print(\"Gradient checkpointing enabled:\", model.backbone.bottom_up.use_checkpoint)\n        return model","metadata":{"_uuid":"378e2303-f039-4228-8af6-fdb47668ab47","_cell_guid":"895f4506-72c1-4c58-9c93-36d0a213f9a2","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T20:44:39.080786Z","iopub.execute_input":"2025-10-25T20:44:39.081443Z","iopub.status.idle":"2025-10-25T20:44:39.086232Z","shell.execute_reply.started":"2025-10-25T20:44:39.081400Z","shell.execute_reply":"2025-10-25T20:44:39.085254Z"}}},{"cell_type":"code","source":"from detectron2.engine import DefaultTrainer\nclass AugTrainer(DefaultTrainer):\n    def run_step(self):\n        import gc, torch\n        gc.collect()\n        torch.cuda.empty_cache()\n        self._trainer.iter = self.iter\n        self._trainer.run_step()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T21:28:51.376251Z","iopub.execute_input":"2025-10-25T21:28:51.376879Z","iopub.status.idle":"2025-10-25T21:28:51.381563Z","shell.execute_reply.started":"2025-10-25T21:28:51.376847Z","shell.execute_reply":"2025-10-25T21:28:51.380561Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if is_train:\n    # Instantiate trainer with memory-safe model\n    trainer = DefaultTrainer(cfg) if not is_augment else AugTrainer(cfg)\n\n    # Force gradient checkpointing if using AugTrainer\n    if is_augment:\n        trainer.model.backbone.bottom_up.use_checkpoint = True\n        print(\"Gradient checkpointing confirmed:\", trainer.model.backbone.bottom_up.use_checkpoint)\n\n    # Resume or start training\n    trainer.resume_or_load(resume=is_resume_training)\n\n    # Optional: clear cache before training loop\n    import gc, torch\n    gc.collect()\n    torch.cuda.empty_cache()\n\n    # Start training\n    trainer.train()\n\n    print(\"#### TRAINING COMPLETE ####\")\n\n    # Switch model to eval mode\n    _ = trainer.model.train(False)\n\n    # Link to final model\n    FileLink(str(OUTPUT_MODEL))","metadata":{"_uuid":"b2406ec3-c1e9-4081-af11-33e74c77588b","_cell_guid":"63025253-161c-4dd5-9c19-c729ec85aa4d","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-10-25T21:28:54.011040Z","iopub.execute_input":"2025-10-25T21:28:54.011372Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 7.2 Visualizing Training Metrics","metadata":{"_uuid":"9ff62687-f419-4b3c-98de-75624d429a1f","_cell_guid":"56de2062-2fc4-489e-9f7d-21460d4508a5","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"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    print(mdf.head(10).T)\n\n    # Plot loss\n    fig, ax = plt.subplots()\n\n    mdf1 = mdf[~mdf[\"total_loss\"].isna()]\n    ax.plot(mdf1[\"iteration\"], mdf1[\"total_loss\"], c=\"C0\", label=\"train\")\n\n    if \"validation_loss\" in mdf.columns:\n        mdf2 = mdf[~mdf[\"validation_loss\"].isna()]\n        ax.plot(mdf2[\"iteration\"], mdf2[\"validation_loss\"],\n                c=\"C1\", label=\"validation\")\n\n    ax.legend()\n    ax.set_title(\"Loss curve\")\n    plt.show()\n\n    # Plot Accuracy\n    fig, ax = plt.subplots()\n\n    mdf1 = mdf[~mdf[\"fast_rcnn/cls_accuracy\"].isna()]\n    ax.plot(mdf1[\"iteration\"], mdf1[\"fast_rcnn/cls_accuracy\"],\n            c=\"C0\", label=\"train\")\n\n    ax.legend()\n    ax.set_title(\"Accuracy curve\")\n    plt.show()\n\n    # Plot Bounding Box regressor loss\n    fig, ax = plt.subplots()\n\n    mdf1 = mdf[~mdf[\"loss_box_reg\"].isna()]\n    ax.plot(mdf1[\"iteration\"], mdf1[\"loss_box_reg\"], c=\"C0\", label=\"train\")\n\n    ax.legend()\n    ax.set_title(\"loss_box_reg\")\n    plt.show()","metadata":{"_uuid":"58417dc3-a312-4499-9941-55e95dd0bdf7","_cell_guid":"96755541-a960-48f6-bb8f-b0afbb779eac","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 8 Evaluation","metadata":{"_uuid":"6ccdcf28-9d92-4502-aead-cdd622061acf","_cell_guid":"a1156a65-66bb-41bc-9290-ac2085b4b8ad","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"markdown","source":"Can evaluate trained model on validation dataset to obtain different metric scores.","metadata":{"_uuid":"9909de03-6800-46c5-85c8-62a3f08540d5","_cell_guid":"14470f27-82b7-4bc8-a310-8658a96df1fd","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"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, cfg, False, output_dir=cfg.OUTPUT_DIR, 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, val_loader, evaluator=evaluator\n    )","metadata":{"_uuid":"241c6011-d39e-4243-8f98-fc1f5c57c8ed","_cell_guid":"7f7ef3b3-aa2f-497d-b7c8-eed39921fcee","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 9 Inference","metadata":{"_uuid":"d044c794-acc7-4bb0-9ab3-ab29a3f8cb3c","_cell_guid":"a608c4ab-7f69-4705-be50-7f2c0b39badc","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"markdown","source":"## 9.1 Setting Up Inference Model","metadata":{"_uuid":"aacef77e-cfb5-4b70-8101-ea9949c3abea","_cell_guid":"51aa8bd2-c98b-45ca-a208-b466231e3506","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"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":{"_uuid":"205c77d6-45c0-4328-ad5b-440bc9775103","_cell_guid":"85addb97-99d4-4fab-a5d7-a3005f194b95","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"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":{"_uuid":"17a9cece-d9c1-4d2b-90a0-8be7b6ae26b6","_cell_guid":"4fdb55d9-2c4a-4811-bc0a-ab3b2eefcca2","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"How confident should the model be for you to accept the predicted mask?","metadata":{"_uuid":"32da1549-a86c-4ec4-b0bf-f45c450003e5","_cell_guid":"17728649-0da7-4dc7-afe1-cb8d31b65b6f","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"ACCEPTANCE_THRESHOLD = 0.6  # for all categories","metadata":{"_uuid":"81caeaed-85ec-4fb2-97a7-00c1134c5660","_cell_guid":"6ccd5be3-2a4d-4428-99f7-b4b3d3733590","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(f\"#### MODEL: {inf_cfg.MODEL.WEIGHTS} FOR INFERENCE ####\")","metadata":{"_uuid":"c16d6142-fcac-427f-b1e2-203fe645164a","_cell_guid":"7254835f-4279-4416-9001-c74871ac6553","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 9.2 Sample Inference","metadata":{"_uuid":"bf040989-e739-4581-8694-2746e2e92f6c","_cell_guid":"83fa2ac9-d305-48e9-8119-c3ce5ed4d913","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"predictor = DefaultPredictor(inf_cfg)","metadata":{"_uuid":"35a04467-a698-4916-8035-54bb63763b32","_cell_guid":"eedcaeb1-6c07-4ddb-a847-58369136027b","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"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 = 4\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(\n        outputs[\"instances\"].to(\"cpu\")\n    )\n    display_img = visualizer.get_image()[:, :, ::-1]\n    indices[i].grid(False)\n    indices[i].imshow(display_img)","metadata":{"_uuid":"3ce13495-96e2-4ed5-912a-8b51943648e4","_cell_guid":"47749978-254f-4798-af6a-00e4381dac5b","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 9.3 Test Data Inference and Submission","metadata":{"_uuid":"cc18a46b-7fcc-48e1-8ad7-3dd19e8bf1f6","_cell_guid":"292c683d-fbc8-4b75-b1c5-d061d459a758","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"markdown","source":"### 9.3.1 Building Inference Model","metadata":{"_uuid":"5f4b1e16-e816-4e69-a268-0141561eaf60","_cell_guid":"b932d38c-0740-46a0-8ae7-7abf4fdfb574","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"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":{"_uuid":"c93ae3a0-e73a-4211-92e7-c16daad4cb6e","_cell_guid":"eb2cf7da-a45c-4dd9-8215-ef3e187b7a9f","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = rebuild_model()","metadata":{"_uuid":"e3736ca8-2289-4e29-a947-052dcc16a134","_cell_guid":"06e40828-c57e-408b-b790-757d767a0570","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 9.3.2 CUDA Problems","metadata":{"_uuid":"8f344fdb-d295-48e5-98ef-121bd0bfa3c8","_cell_guid":"52922d85-f34c-4ae0-a225-347eb330b51e","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"!export LRU_CACHE_CAPACITY=1\n!export 'PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512'","metadata":{"_uuid":"041f2252-7a39-4271-8ff3-808c4f446557","_cell_guid":"d9f73618-577b-45ba-bc45-d5351d1ccc21","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Occassionally CUDA memory overflow occurs. Trying to save as much VRAM as we can.","metadata":{"_uuid":"e6187a00-7e80-4219-9858-d7b7f5545c14","_cell_guid":"64d441bf-35de-420a-9008-6c47756956b8","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"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}\")\n        del globals()[v]\n    elif v in locals():\n        print(f\"Deleting {v}\")\n        del locals()[v]","metadata":{"_uuid":"33d71602-eac6-44c2-a2d3-c8e64b528bf4","_cell_guid":"e650e584-0258-4064-8241-f008d8e87aa5","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 9.3.3 Inference Utils","metadata":{"_uuid":"4391be71-47db-4697-801b-9b9157df9348","_cell_guid":"f7072866-95c1-4ab0-9aca-1c43d828fc1b","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"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":{"_uuid":"9968c66a-f97d-4e1b-8344-79b13eefa6fc","_cell_guid":"e6db3d1a-e89f-4ca6-93da-9882322f5749","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"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":{"_uuid":"cd687462-715c-44c5-b6d8-ee3fb31509df","_cell_guid":"94195aea-b2e0-4bb9-9f83-bebbf806820a","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"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    return results","metadata":{"_uuid":"edef6b1b-e39f-4251-9942-c26d8f418861","_cell_guid":"de8cdaf4-c2e5-4448-9eb7-f165ccbdbb73","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"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":{"_uuid":"f08b967f-2353-4930-82cd-4b1569c7f257","_cell_guid":"75e99457-755b-4388-8ad3-bcdab3bb5c73","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"print(\"#### RUNNING INFERENCE ON TEST DATA ####\")","metadata":{"_uuid":"5f4d687f-41ce-4b9b-ae47-cc44d0210807","_cell_guid":"3bbfe9ce-c7cb-489c-b860-8b291a39a2fb","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"torch.cuda.empty_cache()\ngc.collect()","metadata":{"_uuid":"ff849166-9359-4bcb-836a-6f423f89796c","_cell_guid":"c905faef-e0c3-4c75-a080-a74c15e50b19","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"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    submission_file.writelines(results)\n    submission_file.close()","metadata":{"_uuid":"93f3e61a-bd9a-4659-ae04-7c78c89ba250","_cell_guid":"a901a8e9-386f-4fb6-87d4-78a9d49706b8","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if Path(\"submission.csv\").exists:\n    display(FileLink(\"submission.csv\"))","metadata":{"_uuid":"28ba0d21-6ac5-498f-b33d-e84b46766116","_cell_guid":"ce816839-4ef9-4790-9ba3-dfa390f3bd1a","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!rm -r detectron2/","metadata":{"_uuid":"af6cfeb9-2095-44e5-93d4-c022624273ad","_cell_guid":"d21dac18-1ff1-4bb7-9c25-e4a003a0aa02","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null}]}