{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"#!git clone https://github.com/facebookresearch/detectron2.git\n#!python -m pip install -e detectron2","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-07-07T18:21:06.232299Z","iopub.execute_input":"2023-07-07T18:21:06.232918Z","iopub.status.idle":"2023-07-07T18:21:06.264437Z","shell.execute_reply.started":"2023-07-07T18:21:06.232879Z","shell.execute_reply":"2023-07-07T18:21:06.263163Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%capture\n!pip install 'git+https://github.com/facebookresearch/detectron2.git@5aeb252b194b93dc2879b4ac34bc51a31b5aee13'","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:21:06.266071Z","iopub.execute_input":"2023-07-07T18:21:06.266558Z","iopub.status.idle":"2023-07-07T18:22:46.145788Z","shell.execute_reply.started":"2023-07-07T18:21:06.266532Z","shell.execute_reply":"2023-07-07T18:22:46.144215Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%capture\n# 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\nimport seaborn as sns\n\n!pip install -q pycocotools\nfrom pycocotools.coco import COCO\nfrom PIL import Image\nimport random\nfrom pathlib import Path\n%matplotlib inline\nsns.set_theme(style='darkgrid', palette='deep', font='sans-serif', font_scale=1)\n\nfrom IPython.display import FileLink\n# torch\n#import os\n#os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:24'\nimport torch\nimport gc\nimport warnings\n# Ignore \"future\" warnings and Data-Frame-Slicing warnings.\nwarnings.filterwarnings('ignore')\n\nsetup_logger()","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:22:46.147327Z","iopub.execute_input":"2023-07-07T18:22:46.147716Z","iopub.status.idle":"2023-07-07T18:22:58.633670Z","shell.execute_reply.started":"2023-07-07T18:22:46.147679Z","shell.execute_reply":"2023-07-07T18:22:58.632766Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"### Data load\ntorch.cuda.empty_cache()\nfrom 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":{"execution":{"iopub.status.busy":"2023-07-07T18:22:58.635965Z","iopub.execute_input":"2023-07-07T18:22:58.636771Z","iopub.status.idle":"2023-07-07T18:22:58.641780Z","shell.execute_reply.started":"2023-07-07T18:22:58.636745Z","shell.execute_reply":"2023-07-07T18:22:58.641127Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"### Coco Annotation\n\nwith 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    \ntrain_coco_labels=COCO(annotation_file=TRAIN_COCO_PATH)\n\nprint(\"#### LABELS AND METADATA LOADED ####\")","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:22:58.642893Z","iopub.execute_input":"2023-07-07T18:22:58.643149Z","iopub.status.idle":"2023-07-07T18:23:07.630515Z","shell.execute_reply.started":"2023-07-07T18:22:58.643129Z","shell.execute_reply":"2023-07-07T18:23:07.629598Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Decisions\n\nfrom 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())\n\n# Model path based on Decisions\nMODEL_PATH = OUTPUT_MODEL if is_train else PRETRAINED_PATH","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:07.631506Z","iopub.execute_input":"2023-07-07T18:23:07.631720Z","iopub.status.idle":"2023-07-07T18:23:07.637559Z","shell.execute_reply.started":"2023-07-07T18:23:07.631701Z","shell.execute_reply":"2023-07-07T18:23:07.636166Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"print(\"There are \" + str(len(train_dict['categories'])) + \" categories.\\n\")\nprint(\"There are \" + str(len(test_dict['images']) + len(train_dict['images'])) + \" images in the dataset.\")\nprint(\"There are \" + str(len(train_dict['images'])) + \" images in the train set.\")\nprint(\"There are \" + str(len(test_dict['images'])) + \" images in the test set.\\n\")\nprint(\"There are \" + str(len(train_dict['annotations'])) + \" annotations in the train set.\\n\")\n\nprint(\"We will focus on mainly categories, images and annotations.\")","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:07.639071Z","iopub.execute_input":"2023-07-07T18:23:07.639362Z","iopub.status.idle":"2023-07-07T18:23:07.655072Z","shell.execute_reply.started":"2023-07-07T18:23:07.639334Z","shell.execute_reply":"2023-07-07T18:23:07.654004Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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) coordinnates\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":{"execution":{"iopub.status.busy":"2023-07-07T18:23:07.656777Z","iopub.execute_input":"2023-07-07T18:23:07.657161Z","iopub.status.idle":"2023-07-07T18:23:07.667971Z","shell.execute_reply.started":"2023-07-07T18:23:07.657134Z","shell.execute_reply":"2023-07-07T18:23:07.667045Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"thing_classes, images_metadata, data_annotations = organize_coco_data(train_dict)\n\nthing_classes_test, images_metadata_test, _ = organize_coco_data(test_dict)\n\nprint(thing_classes)","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:07.669272Z","iopub.execute_input":"2023-07-07T18:23:07.669520Z","iopub.status.idle":"2023-07-07T18:23:08.971570Z","shell.execute_reply.started":"2023-07-07T18:23:07.669500Z","shell.execute_reply":"2023-07-07T18:23:08.970221Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":{"execution":{"iopub.status.busy":"2023-07-07T18:23:08.974936Z","iopub.execute_input":"2023-07-07T18:23:08.975257Z","iopub.status.idle":"2023-07-07T18:23:09.046814Z","shell.execute_reply.started":"2023-07-07T18:23:08.975235Z","shell.execute_reply":"2023-07-07T18:23:09.045993Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":{"execution":{"iopub.status.busy":"2023-07-07T18:23:09.048005Z","iopub.execute_input":"2023-07-07T18:23:09.048363Z","iopub.status.idle":"2023-07-07T18:23:09.889509Z","shell.execute_reply.started":"2023-07-07T18:23:09.048336Z","shell.execute_reply":"2023-07-07T18:23:09.888218Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":{"execution":{"iopub.status.busy":"2023-07-07T18:23:09.890753Z","iopub.execute_input":"2023-07-07T18:23:09.891175Z","iopub.status.idle":"2023-07-07T18:23:09.925298Z","shell.execute_reply.started":"2023-07-07T18:23:09.891150Z","shell.execute_reply":"2023-07-07T18:23:09.924227Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(15, 15))\nfor i in range(6):\n    image_id = random.randint(2, 20364)\n    image_file = train_coco_labels.loadImgs([image_id])[0]['file_name']\n    image = Image.open(TRAIN_IMG_DIR/image_file)\n    code = 231 + i;\n    plt.subplot(code)\n    plt.axis('off')\n    plt.imshow(np.asarray(image))","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:09.926592Z","iopub.execute_input":"2023-07-07T18:23:09.926877Z","iopub.status.idle":"2023-07-07T18:23:12.671101Z","shell.execute_reply.started":"2023-07-07T18:23:09.926855Z","shell.execute_reply":"2023-07-07T18:23:12.669821Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_image(_image_id):\n    image_file = train_coco_labels.loadImgs([_image_id])[0]['file_name']\n    image = Image.open(TRAIN_IMG_DIR/image_file)\n    plt.figure(figsize=(15, 15))\n    plt.axis('off')\n    plt.imshow(np.asarray(image))","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:12.672652Z","iopub.execute_input":"2023-07-07T18:23:12.672951Z","iopub.status.idle":"2023-07-07T18:23:12.679696Z","shell.execute_reply.started":"2023-07-07T18:23:12.672927Z","shell.execute_reply":"2023-07-07T18:23:12.678364Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_segmentations(_image_id):\n    show_image(_image_id)\n    annotation_ids = train_coco_labels.getAnnIds(imgIds=[_image_id])\n    annotations = train_coco_labels.loadAnns(annotation_ids)\n    train_coco_labels.showAnns(annotations)","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:12.681186Z","iopub.execute_input":"2023-07-07T18:23:12.681485Z","iopub.status.idle":"2023-07-07T18:23:12.691083Z","shell.execute_reply.started":"2023-07-07T18:23:12.681462Z","shell.execute_reply":"2023-07-07T18:23:12.690195Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_image(51)","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:12.692037Z","iopub.execute_input":"2023-07-07T18:23:12.692346Z","iopub.status.idle":"2023-07-07T18:23:13.197123Z","shell.execute_reply.started":"2023-07-07T18:23:12.692320Z","shell.execute_reply":"2023-07-07T18:23:13.196089Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_segmentations(51)","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:13.198840Z","iopub.execute_input":"2023-07-07T18:23:13.199210Z","iopub.status.idle":"2023-07-07T18:23:13.740623Z","shell.execute_reply.started":"2023-07-07T18:23:13.199186Z","shell.execute_reply":"2023-07-07T18:23:13.739655Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Exploring Annonations\ntrain_annotations = pd.DataFrame(train_dict['annotations'])\ntrain_annotations.head()","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:13.741625Z","iopub.execute_input":"2023-07-07T18:23:13.742287Z","iopub.status.idle":"2023-07-07T18:23:14.563321Z","shell.execute_reply.started":"2023-07-07T18:23:13.742264Z","shell.execute_reply":"2023-07-07T18:23:14.562424Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# adding bbox_area and bbox_aspect_ratio\nprint(\"\\nRenaming id to annotation_id.\\nAdding bbox_area, bbox_aspect_ratio.\\n\")\ntrain_annotations.rename(columns={\"id\":\"annotation_id\"}, inplace=True)\nbbox_area=[]\nbbox_aspect_ratio=[]\nfor idx in train_annotations.index:\n    bbox_area.append(train_annotations['bbox'][idx][3]*train_annotations['bbox'][idx][2])\n    bbox_aspect_ratio.append(train_annotations['bbox'][idx][3]/train_annotations['bbox'][idx][2])\ntrain_annotations['bbox_area']=bbox_area\ntrain_annotations['bbox_aspect_ratio']=bbox_aspect_ratio\nprint(\"train_annotations shape: \" + str(train_annotations.shape) + \"\\n\")\ntrain_annotations.head()","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:14.564451Z","iopub.execute_input":"2023-07-07T18:23:14.565353Z","iopub.status.idle":"2023-07-07T18:23:25.036443Z","shell.execute_reply.started":"2023-07-07T18:23:14.565324Z","shell.execute_reply":"2023-07-07T18:23:25.035100Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_annotations.describe()","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:25.038184Z","iopub.execute_input":"2023-07-07T18:23:25.038558Z","iopub.status.idle":"2023-07-07T18:23:25.125631Z","shell.execute_reply.started":"2023-07-07T18:23:25.038527Z","shell.execute_reply":"2023-07-07T18:23:25.124511Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Categorical Calssification\ntrain_categories = pd.DataFrame(train_dict['categories'])\n\nprint(\"Dropping supercategory.\\nRenaming id to category_id.\\n\")\ntrain_categories.drop('supercategory', axis=1, inplace=True)\ntrain_categories.rename(columns={\"id\":\"category_id\"}, inplace=True)\ntrain_categories","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:25.126981Z","iopub.execute_input":"2023-07-07T18:23:25.127371Z","iopub.status.idle":"2023-07-07T18:23:25.140172Z","shell.execute_reply.started":"2023-07-07T18:23:25.127348Z","shell.execute_reply":"2023-07-07T18:23:25.139081Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"img_cat_count = train_annotations[['image_id', 'category_id']].copy()\ncategory_names=[]\nfor idx in img_cat_count.index:\n    category_names.append(train_categories['name'][img_cat_count['category_id'][idx]])\nimg_cat_count['category_name']=category_names\nimg_cat_count[train_categories['name'][0]] = img_cat_count['category_id'] == train_categories['category_id'][0]\nimg_cat_count[train_categories['name'][1]] = img_cat_count['category_id'] == train_categories['category_id'][1]\nimg_cat_count[train_categories['name'][2]] = img_cat_count['category_id'] == train_categories['category_id'][2]\nimg_cat_count[train_categories['name'][3]] = img_cat_count['category_id'] == train_categories['category_id'][3]\nimg_cat_count.head()","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:25.141742Z","iopub.execute_input":"2023-07-07T18:23:25.142089Z","iopub.status.idle":"2023-07-07T18:23:30.229620Z","shell.execute_reply.started":"2023-07-07T18:23:25.142064Z","shell.execute_reply":"2023-07-07T18:23:30.228701Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\nCategorywise Document Image Count.\\n\")\ncategorywise_image_count=img_cat_count.groupby('category_id', as_index=False)['image_id'].nunique()\ncategorywise_image_count['category_name']=train_categories['name']\ncategorywise_image_count.rename(columns={'image_id':'image_count'}, inplace=True)\ncategorywise_image_count = categorywise_image_count[['category_id', 'category_name', 'image_count']]\ncategorywise_image_count","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:30.230740Z","iopub.execute_input":"2023-07-07T18:23:30.231049Z","iopub.status.idle":"2023-07-07T18:23:30.269594Z","shell.execute_reply.started":"2023-07-07T18:23:30.231000Z","shell.execute_reply":"2023-07-07T18:23:30.268476Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.title(\"Categories vs Number of document images they appear\")\nsns.barplot(x=categorywise_image_count['category_name'], y=categorywise_image_count['image_count'])","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:30.270794Z","iopub.execute_input":"2023-07-07T18:23:30.271098Z","iopub.status.idle":"2023-07-07T18:23:30.540264Z","shell.execute_reply.started":"2023-07-07T18:23:30.271076Z","shell.execute_reply":"2023-07-07T18:23:30.539157Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\nImagewise Category Count.\\n\")\nimagewise_category_count=img_cat_count.groupby('image_id', as_index=False)[['paragraph', 'text_box', 'image', 'table']].sum()\nimagewise_category_count.head()","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:30.541944Z","iopub.execute_input":"2023-07-07T18:23:30.542363Z","iopub.status.idle":"2023-07-07T18:23:30.584053Z","shell.execute_reply.started":"2023-07-07T18:23:30.542330Z","shell.execute_reply":"2023-07-07T18:23:30.582789Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(10, 20))\nplt.subplot(411)\nplt.title(\"Number of paragraphs per document image\")\nsns.lineplot(x=imagewise_category_count['image_id'], y=imagewise_category_count['paragraph'])\n\nplt.subplot(412)\nplt.title(\"Number of text boxes per document image\")\nsns.lineplot(x=imagewise_category_count['image_id'], y=imagewise_category_count['text_box'])\n\nplt.subplot(413)\nplt.title(\"Number of images per document image\")\nsns.lineplot(x=imagewise_category_count['image_id'], y=imagewise_category_count['image'])\n\nplt.subplot(414)\nplt.title(\"Number of tables per document image\")\nsns.lineplot(x=imagewise_category_count['image_id'], y=imagewise_category_count['table'])","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:30.585313Z","iopub.execute_input":"2023-07-07T18:23:30.585603Z","iopub.status.idle":"2023-07-07T18:23:32.212344Z","shell.execute_reply.started":"2023-07-07T18:23:30.585579Z","shell.execute_reply":"2023-07-07T18:23:32.211200Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Preparing Data for Training\n","metadata":{}},{"cell_type":"code","source":"TRAIN_SPLIT = 0.80","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:32.213574Z","iopub.execute_input":"2023-07-07T18:23:32.213889Z","iopub.status.idle":"2023-07-07T18:23:32.220616Z","shell.execute_reply.started":"2023-07-07T18:23:32.213862Z","shell.execute_reply":"2023-07-07T18:23:32.219090Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":{"execution":{"iopub.status.busy":"2023-07-07T18:23:32.225003Z","iopub.execute_input":"2023-07-07T18:23:32.225377Z","iopub.status.idle":"2023-07-07T18:23:32.235059Z","shell.execute_reply.started":"2023-07-07T18:23:32.225346Z","shell.execute_reply":"2023-07-07T18:23:32.233718Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Formatting Data for detectron2\ndef 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        # 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\n\n","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:32.236311Z","iopub.execute_input":"2023-07-07T18:23:32.236579Z","iopub.status.idle":"2023-07-07T18:23:32.248647Z","shell.execute_reply.started":"2023-07-07T18:23:32.236559Z","shell.execute_reply":"2023-07-07T18:23:32.247078Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Registering and Loading Data for detectron2\nDATA_REGISTER_TRAINING = \"badlad_train\"\nDATA_REGISTER_VALID    = \"badlad_valid\"\nDATA_REGISTER_TEST     = \"badlad_test\"","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:32.250408Z","iopub.execute_input":"2023-07-07T18:23:32.250956Z","iopub.status.idle":"2023-07-07T18:23:32.265032Z","shell.execute_reply.started":"2023-07-07T18:23:32.250896Z","shell.execute_reply":"2023-07-07T18:23:32.263547Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_inds","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:23:32.266778Z","iopub.execute_input":"2023-07-07T18:23:32.267229Z","iopub.status.idle":"2023-07-07T18:23:32.279435Z","shell.execute_reply.started":"2023-07-07T18:23:32.267193Z","shell.execute_reply":"2023-07-07T18:23:32.278116Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":{"execution":{"iopub.status.busy":"2023-07-07T18:23:32.281209Z","iopub.execute_input":"2023-07-07T18:23:32.281547Z","iopub.status.idle":"2023-07-07T18:25:04.704236Z","shell.execute_reply.started":"2023-07-07T18:23:32.281522Z","shell.execute_reply":"2023-07-07T18:25:04.702805Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":{"execution":{"iopub.status.busy":"2023-07-07T18:25:04.705920Z","iopub.execute_input":"2023-07-07T18:25:04.706662Z","iopub.status.idle":"2023-07-07T18:26:39.140293Z","shell.execute_reply.started":"2023-07-07T18:25:04.706629Z","shell.execute_reply":"2023-07-07T18:26:39.139007Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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)\nmetadata_dicts_test = MetadataCatalog.get(DATA_REGISTER_TEST)\nprint(\"#### DATA REGISTERED ####\")","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:26:39.141887Z","iopub.execute_input":"2023-07-07T18:26:39.142824Z","iopub.status.idle":"2023-07-07T18:26:39.898496Z","shell.execute_reply.started":"2023-07-07T18:26:39.142790Z","shell.execute_reply":"2023-07-07T18:26:39.897801Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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    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":{"execution":{"iopub.status.busy":"2023-07-07T18:26:39.899343Z","iopub.execute_input":"2023-07-07T18:26:39.899600Z","iopub.status.idle":"2023-07-07T18:26:39.907413Z","shell.execute_reply.started":"2023-07-07T18:26:39.899578Z","shell.execute_reply":"2023-07-07T18:26:39.906026Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":{"execution":{"iopub.status.busy":"2023-07-07T18:26:39.908875Z","iopub.execute_input":"2023-07-07T18:26:39.909206Z","iopub.status.idle":"2023-07-07T18:26:39.923246Z","shell.execute_reply.started":"2023-07-07T18:26:39.909181Z","shell.execute_reply":"2023-07-07T18:26:39.921918Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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\n    cfg.DATALOADER.NUM_WORKERS = 2\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\n    cfg.SOLVER.BASE_LR = 0.001\n\n    cfg.SOLVER.WARMUP_ITERS = 5\n\n    # Maximum number of iterations\n    cfg.SOLVER.MAX_ITER = 5000\n\n    # cfg.SOLVER.STEPS = (500, 1000) # must be less than MAX_ITER\n\n    cfg.SOLVER.GAMMA = 0.09\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\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":{"execution":{"iopub.status.busy":"2023-07-07T18:26:39.924741Z","iopub.execute_input":"2023-07-07T18:26:39.925447Z","iopub.status.idle":"2023-07-07T18:26:39.953087Z","shell.execute_reply.started":"2023-07-07T18:26:39.925418Z","shell.execute_reply":"2023-07-07T18:26:39.951974Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"!nvidia-smi","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:33:28.035577Z","iopub.execute_input":"2023-07-07T18:33:28.035969Z","iopub.status.idle":"2023-07-07T18:33:28.039847Z","shell.execute_reply.started":"2023-07-07T18:33:28.035943Z","shell.execute_reply":"2023-07-07T18:33:28.039271Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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    FileLink(str(OUTPUT_MODEL))","metadata":{"execution":{"iopub.status.busy":"2023-07-07T18:33:55.953870Z","iopub.execute_input":"2023-07-07T18:33:55.954227Z","iopub.status.idle":"2023-07-07T18:34:08.301123Z","shell.execute_reply.started":"2023-07-07T18:33:55.954203Z","shell.execute_reply":"2023-07-07T18:34:08.299987Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Visualizing Training Metrics\nif 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":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Evaluation on validation set","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, 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":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Inference","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":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BATCH = 8  # lower if CUDA overflow occurs\ntest_loader = build_detection_test_loader(inf_cfg, DATA_REGISTER_TEST, batch_size=BATCH)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ACCEPTANCE_THRESHOLD = 0.6 ","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(f\"#### MODEL: {inf_cfg.MODEL.WEIGHTS} FOR INFERENCE ####\")","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"predictor = DefaultPredictor(inf_cfg)\nfig, 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":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Test Data Inference\n#Building Inference Model\n\ndef rebuild_model():\n    model = build_model(inf_cfg)\n    _ = DetectionCheckpointer(model).load(inf_cfg.MODEL.WEIGHTS)\n    return model","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = rebuild_model()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!export LRU_CACHE_CAPACITY=1\n!export 'PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512'","metadata":{},"execution_count":null,"outputs":[]},{"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":{},"execution_count":null,"outputs":[]},{"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":{},"execution_count":null,"outputs":[]},{"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":{},"execution_count":null,"outputs":[]},{"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":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"#### RUNNING INFERENCE ON TEST DATA ####\")\ntorch.cuda.empty_cache()\ngc.collect()","metadata":{},"execution_count":null,"outputs":[]},{"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":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if Path(\"submission.csv\").exists:\n    display(FileLink(\"submission.csv\"))","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!rm -r detectron2/","metadata":{},"execution_count":null,"outputs":[]}]}