{"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":"markdown","source":"# Install Detectron2","metadata":{}},{"cell_type":"code","source":"!rm -rf detectron2/\n!python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'","metadata":{"execution":{"iopub.status.busy":"2023-08-06T13:14:18.453393Z","iopub.execute_input":"2023-08-06T13:14:18.454435Z","iopub.status.idle":"2023-08-06T13:17:57.343838Z","shell.execute_reply.started":"2023-08-06T13:14:18.454396Z","shell.execute_reply":"2023-08-06T13:17:57.342717Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Download Pretrained Model","metadata":{}},{"cell_type":"code","source":"!wget 'https://www.dropbox.com/scl/fi/hrhdo9ev957hbllbnzctt/model_final.pth?rlkey=lfn9ju29g1xtpcxz3pl2pj8do&dl=0' -O model.pth","metadata":{"execution":{"iopub.status.busy":"2023-08-06T13:18:12.136945Z","iopub.execute_input":"2023-08-06T13:18:12.137354Z","iopub.status.idle":"2023-08-06T13:18:20.286103Z","shell.execute_reply.started":"2023-08-06T13:18:12.137308Z","shell.execute_reply":"2023-08-06T13:18:20.284913Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Set Some Constants","metadata":{}},{"cell_type":"code","source":"!export LRU_CACHE_CAPACITY=1\n# inference constants\nACCEPTANCE_THRESHOLD=0.45\nBATCH = 2","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Make all necessary imports","metadata":{}},{"cell_type":"code","source":"from datetime import datetime\nimport logging\nfrom pathlib import Path\nfrom detectron2.checkpoint import DetectionCheckpointer\nfrom detectron2.config import LazyConfig, instantiate\nfrom detectron2.engine import (\n    AMPTrainer,\n    SimpleTrainer,\n    default_argument_parser,\n    default_setup,\n    default_writers,\n    hooks,\n    launch\n)\nfrom detectron2.engine.defaults import _try_get_key\nfrom detectron2.engine.defaults import create_ddp_model\nfrom detectron2.evaluation import inference_on_dataset, print_csv_format\nfrom detectron2.utils import comm\n\nfrom functools import partial\nfrom detectron2.utils.file_io import PathManager\nfrom omegaconf import OmegaConf\nimport torch.nn as nn\nfrom fvcore.common.param_scheduler import MultiStepParamScheduler\n\nfrom detectron2 import model_zoo\nfrom detectron2.config import LazyCall as L\nfrom detectron2.solver import WarmupParamScheduler\nfrom detectron2.modeling import MViT\nfrom detectron2.layers import ShapeSpec\nfrom detectron2.modeling.box_regression import Box2BoxTransform\nfrom detectron2.modeling.matcher import Matcher\nfrom detectron2.modeling.roi_heads import (\n    FastRCNNOutputLayers,\n    FastRCNNConvFCHead,\n    CascadeROIHeads,\n)\nfrom detectron2.utils.env import seed_all_rng\nfrom detectron2.data.datasets import register_coco_instances\nimport detectron2.data.transforms as T\nfrom detectron2 import model_zoo\nfrom detectron2.config import LazyCall as L\nimport detectron2.data.transforms as T\nfrom detectron2.config import LazyCall as L\nfrom detectron2.data import (\n    DatasetMapper,\n    build_detection_test_loader,\n    build_detection_train_loader,\n    get_detection_dataset_dicts,\n)\nfrom detectron2.evaluation import COCOEvaluator\nfrom detectron2.utils.logger import setup_logger\nfrom detectron2.checkpoint import DetectionCheckpointer\nfrom detectron2.config import LazyConfig, instantiate\nfrom detectron2.data.detection_utils import read_image\nfrom detectron2.utils.logger import setup_logger\nimport matplotlib.pyplot as plt\nimport time\n\nimport atexit\nimport bisect\nfrom copy import copy\nimport multiprocessing as mp\nfrom collections import deque\nimport cv2\nimport torch\n\nimport detectron2.data.transforms as T\nfrom detectron2.data import MetadataCatalog\nfrom detectron2.structures import Instances\nfrom detectron2.utils.video_visualizer import VideoVisualizer\nfrom detectron2.utils.visualizer import ColorMode, Visualizer\nimport numpy as np\nfrom torchvision import transforms","metadata":{"execution":{"iopub.status.busy":"2023-08-06T13:18:20.28792Z","iopub.execute_input":"2023-08-06T13:18:20.288582Z","iopub.status.idle":"2023-08-06T13:18:22.732371Z","shell.execute_reply.started":"2023-08-06T13:18:20.288537Z","shell.execute_reply":"2023-08-06T13:18:22.731175Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Register Datasets","metadata":{}},{"cell_type":"code","source":"register_coco_instances(\"badlad_train\", {}, '/kaggle/input/dlsprint2/badlad/labels/coco_format/train/badlad-train-coco.json', '/kaggle/input/dlsprint2/badlad/images/train')\nregister_coco_instances(\"badlad_test\", {}, '/kaggle/input/dlsprint2/badlad/badlad-test-metadata.json', '/kaggle/input/dlsprint2/badlad/images/test')","metadata":{"execution":{"iopub.status.busy":"2023-08-06T13:18:22.735852Z","iopub.execute_input":"2023-08-06T13:18:22.736452Z","iopub.status.idle":"2023-08-06T13:18:22.74554Z","shell.execute_reply.started":"2023-08-06T13:18:22.736423Z","shell.execute_reply":"2023-08-06T13:18:22.741731Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Configure Model, Trainer and Dataloader","metadata":{}},{"cell_type":"code","source":"def setup():\n    model = model_zoo.get_config(\"common/models/mask_rcnn_fpn.py\").model\n    constants = model_zoo.get_config(\"common/data/constants.py\").constants\n    model.pixel_mean = constants.imagenet_rgb256_mean\n    model.pixel_std = constants.imagenet_rgb256_std\n    model.input_format = \"RGB\"\n    model.backbone.bottom_up = L(MViT)(\n        embed_dim=96,\n        depth=24,\n        num_heads=1,\n        last_block_indexes=(1, 4, 20, 23),\n        residual_pooling=True,\n        drop_path_rate=0.4,\n        norm_layer=partial(nn.LayerNorm, eps=1e-6),\n        out_features=(\"scale2\", \"scale3\", \"scale4\", \"scale5\"),\n    )\n    model.backbone.in_features = \"${.bottom_up.out_features}\"\n    model.backbone.square_pad = 1024\n\n    # New heads and LN\n    model.backbone.norm = \"LN\"  # Use LN in FPN\n    model.roi_heads.box_head.conv_norm = model.roi_heads.mask_head.conv_norm = \"LN\"\n\n    # 2conv in RPN:\n    model.proposal_generator.head.conv_dims = [-1, -1]\n\n    # arguments that don't exist for Cascade R-CNN\n    [model.roi_heads.pop(k) for k in [\"box_head\", \"box_predictor\", \"proposal_matcher\"]]\n    model.roi_heads.update(\n        _target_=CascadeROIHeads,\n        box_heads=[\n            L(FastRCNNConvFCHead)(\n                input_shape=ShapeSpec(channels=256, height=7, width=7),\n                conv_dims=[256, 256, 256, 256],\n                fc_dims=[1024],\n                conv_norm=\"LN\",\n            )\n            for _ in range(3)\n        ],\n        box_predictors=[\n            L(FastRCNNOutputLayers)(\n                input_shape=ShapeSpec(channels=1024),\n                test_score_thresh=0.05,\n                box2box_transform=L(Box2BoxTransform)(weights=(w1, w1, w2, w2)),\n                cls_agnostic_bbox_reg=True,\n                num_classes=\"${...num_classes}\",\n                test_topk_per_image=1000\n            )\n            for (w1, w2) in [(10, 5), (20, 10), (30, 15)]\n        ],\n        proposal_matchers=[\n            L(Matcher)(thresholds=[th], labels=[0, 1], allow_low_quality_matches=False)\n            for th in [0.5, 0.6, 0.7]\n        ],\n    )\n    model.roi_heads.num_classes = 4\n    model.roi_heads.batch_size_per_image = 512\n\n    dataloader = OmegaConf.create()\n\n    image_size = 1024\n    dataloader.train = L(build_detection_train_loader)(\n        dataset=L(get_detection_dataset_dicts)(names=\"badlad_train\"),\n        mapper=L(DatasetMapper)(\n            is_train=True,\n            augmentations=[\n                L(T.RandomBrightness)(intensity_min=0.8,intensity_max=1.2),\n                L(T.RandomContrast)(intensity_min=0.5,intensity_max=1.5),\n                L(T.RandomSaturation)(intensity_min=0.5,intensity_max=1.0),\n                L(T.RandomRotation)(angle=[-5, 5], sample_style=\"range\"),\n                L(T.ResizeScale)(\n                    min_scale=0.1, max_scale=2.0, target_width=image_size, target_height=image_size\n                ),\n                L(T.FixedSizeCrop)(crop_size=(image_size, image_size), pad=False),\n            ],\n            image_format=\"RGB\",\n            use_instance_mask=True,\n        ),\n        total_batch_size=16,\n        num_workers=4,\n    )\n\n    dataloader.test = L(build_detection_test_loader)(\n        dataset=L(get_detection_dataset_dicts)(names=\"badlad_test\", filter_empty=False),\n        mapper=L(DatasetMapper)(\n            is_train=False,\n            augmentations=[\n                L(T.ResizeShortestEdge)(short_edge_length=image_size, max_size=image_size),\n            ],\n            image_format=\"RGB\",\n        ),\n        batch_size=BATCH,\n        num_workers=2,\n    )\n\n    dataloader.evaluator = L(COCOEvaluator)(\n        dataset_name=\"${..test.dataset.names}\",\n    )\n\n    dataloader.train.num_workers = 2\n    dataloader.train.total_batch_size = 16\n    # recompute boxes due to cropping\n    dataloader.train.mapper.recompute_boxes = True\n\n    # Initialization and trainer settings\n    train = model_zoo.get_config(\"common/train.py\").train\n    train.amp.enabled = True\n    train.ddp.fp16_compression = True\n    train.init_checkpoint = \"detectron2://ImageNetPretrained/mvitv2/MViTv2_B_in21k.pyth\"\n    train.output_dir = \"./output/mvit2b\"\n    # Schedule\n    # 36 epoch = 20365/16 * 36 = 45821 iterations \n    train.max_iter = 45821\n    train.eval_period = 1000\n    train.log_period = 20\n    train.checkpointer.period = 2000\n    train.device = \"cuda\"\n\n    lr_multiplier = L(WarmupParamScheduler)(\n        scheduler=L(MultiStepParamScheduler)(\n            values=[1.0, 0.1, 0.01],\n            milestones=[40730, 44447],\n            num_updates=train.max_iter,\n        ),\n        warmup_length=50 / train.max_iter,\n        warmup_factor=0.001,\n    )\n\n    optimizer = model_zoo.get_config(\"common/optim.py\").AdamW\n    optimizer.params.overrides = {\"pos_embed\": {\"weight_decay\": 0.0}}\n    optimizer.lr = 0.00008\n\n    dataloader.evaluator.output_dir = train.output_dir\n\n    cfg = OmegaConf.create()\n    cfg.model = model\n    cfg.dataloader = dataloader\n    cfg.train = train\n    cfg.optimizer = optimizer\n    cfg.lr_multiplier = lr_multiplier\n\n    return cfg","metadata":{"execution":{"iopub.status.busy":"2023-08-06T13:18:22.749107Z","iopub.execute_input":"2023-08-06T13:18:22.749889Z","iopub.status.idle":"2023-08-06T13:18:22.985684Z","shell.execute_reply.started":"2023-08-06T13:18:22.749852Z","shell.execute_reply":"2023-08-06T13:18:22.98462Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Limit CUDA allocation size to avoid OOM.","metadata":{}},{"cell_type":"code","source":"import os\nos.environ['PYTORCH_CUDA_ALLOC_CONF']='max_split_size_mb:512'","metadata":{"execution":{"iopub.status.busy":"2023-08-06T13:18:23.962219Z","iopub.execute_input":"2023-08-06T13:18:23.965783Z","iopub.status.idle":"2023-08-06T13:18:23.970802Z","shell.execute_reply.started":"2023-08-06T13:18:23.965746Z","shell.execute_reply":"2023-08-06T13:18:23.969646Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Instantiate Model and Load Weights","metadata":{}},{"cell_type":"code","source":"logger = setup_logger()\ncfg = setup()\nmodel = instantiate(cfg.model)\nmodel.to(cfg.train.device)\ncheckpointer = DetectionCheckpointer(model)\ncheckpointer.load('model.pth')\nmodel.eval()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Use RLE encoding for submission format.","metadata":{}},{"cell_type":"code","source":"def rle_encode(mask):\n    pixels = mask.T.flatten()\n    use_padding = False\n    if pixels[0] or pixels[-1]:\n        use_padding = True\n        pixel_padded = np.zeros([len(pixels) + 2], dtype=pixels.dtype)\n        pixel_padded[1:-1] = pixels\n        pixels = pixel_padded\n    rle = np.where(pixels[1:] != pixels[:-1])[0] + 2\n    if use_padding:\n        rle = rle - 1\n    rle[1::2] = rle[1::2] - rle[:-1:2]\n    return ' '.join(str(x) for x in rle)","metadata":{"execution":{"iopub.status.busy":"2023-08-06T13:18:26.315577Z","iopub.status.idle":"2023-08-06T13:18:26.316268Z","shell.execute_reply.started":"2023-08-06T13:18:26.316021Z","shell.execute_reply":"2023-08-06T13:18:26.316044Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from detectron2.utils.memory import retry_if_cuda_oom\n\nthing_classes = ['paragraph', 'text_box', 'image', 'table']\n@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    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":{"iopub.status.busy":"2023-08-06T13:18:26.317547Z","iopub.status.idle":"2023-08-06T13:18:26.31823Z","shell.execute_reply.started":"2023-08-06T13:18:26.317985Z","shell.execute_reply":"2023-08-06T13:18:26.318007Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Run Inference on Test Data","metadata":{}},{"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":{"iopub.status.busy":"2023-08-06T13:18:26.319502Z","iopub.status.idle":"2023-08-06T13:18:26.320194Z","shell.execute_reply.started":"2023-08-06T13:18:26.319946Z","shell.execute_reply":"2023-08-06T13:18:26.319968Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import gc\nfrom tqdm import tqdm\nprint(\"#### RUNNING INFERENCE ON TEST DATA ####\")\ntorch.cuda.empty_cache()\ngc.collect()\n\nsubmission_file = open(\"submission.csv\", \"w\")\nsubmission_file.write(\"Id,Predicted\\n\")\n\nresults: list[str] = []\ntest_loader = instantiate(cfg.dataloader.test)\nfor 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\nsubmission_file.writelines(results)\nsubmission_file.close()","metadata":{"execution":{"iopub.status.busy":"2023-08-06T13:18:26.321443Z","iopub.status.idle":"2023-08-06T13:18:26.322144Z","shell.execute_reply.started":"2023-08-06T13:18:26.321901Z","shell.execute_reply":"2023-08-06T13:18:26.321923Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Acknowledgement\n[MViTv2: Improved Multiscale Vision Transformers for Classification and Detection](https://arxiv.org/abs/2112.01526)","metadata":{}}]}