{"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":"## SwinTransformer RCNN training and evaluation code \nThis code is based on https://github.com/xiaohu2015/SwinT_detectron2<br>\n\nI splitted video 0 and 1 to training and video 2 to validation.<br>\nSo far, I could only train the model for 5 epochs and only achieved:<br>\n**bbox/AP = 15.2<br>**\n**bbox/AP50 = 35.7<br>**\n**bbox/AP75 = 8.0<br>**\non the validation set.\n\nHowever, I'd like to share my work and will continue working on it to further improve.","metadata":{}},{"cell_type":"markdown","source":"## Install requirements\ninstall Detectron2, timm and Detectron2 version implementation of SwinTransformer","metadata":{}},{"cell_type":"code","source":"\"\"\"cpu\"\"\"\n# !python -m pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cpu/torch1.9/index.html\n\"\"\"gpu\"\"\"    \n!pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 -f https://download.pytorch.org/whl/torch_stable.html\n!pip install detectron2==0.5 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu110/torch1.7/index.html\n\n!pip install timm\n!git clone https://github.com/emiz6413/SwinT_detectron2.git swin\n!curl -OL https://github.com/xiaohu2015/SwinT_detectron2/releases/download/v1.1/faster_rcnn_swint_T.pth  # pretrained\n!curl -OL https://github.com/emiz6413/SwinT_detectron2/releases/download/v1.3/model_0021209.pth  # trained","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2021-12-22T11:21:58.448802Z","iopub.execute_input":"2021-12-22T11:21:58.449135Z","iopub.status.idle":"2021-12-22T11:25:18.757792Z","shell.execute_reply.started":"2021-12-22T11:21:58.449057Z","shell.execute_reply":"2021-12-22T11:25:18.75699Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nfrom pathlib import Path\nfrom ast import literal_eval\n\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport torch\nimport matplotlib.pyplot as plt\nfrom PIL import Image\n\nfrom detectron2.data import (DatasetCatalog, \n                             MetadataCatalog, \n                             build_detection_test_loader\n                            )\nfrom detectron2.data.datasets.coco import convert_to_coco_json\nfrom detectron2.utils.visualizer import Visualizer, ColorMode\nfrom detectron2.structures import BoxMode\nfrom detectron2.config import get_cfg\nfrom detectron2.utils.logger import setup_logger\nfrom detectron2.engine import DefaultTrainer, default_setup, hooks\nfrom detectron2.modeling import GeneralizedRCNNWithTTA\nfrom detectron2.evaluation import COCOEvaluator, DatasetEvaluators, inference_on_dataset\n\nfrom swin.swint import add_swint_config\n\nlogger = setup_logger()","metadata":{"execution":{"iopub.status.busy":"2021-12-22T11:41:59.466568Z","iopub.execute_input":"2021-12-22T11:41:59.46682Z","iopub.status.idle":"2021-12-22T11:41:59.4738Z","shell.execute_reply.started":"2021-12-22T11:41:59.466792Z","shell.execute_reply":"2021-12-22T11:41:59.473058Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 1. Load data","metadata":{}},{"cell_type":"code","source":"data_root = Path('/kaggle/input/tensorflow-great-barrier-reef/')\ntrain_df = pd.read_csv(str(data_root/'train.csv'))\ntrain_df['annotations'] = train_df['annotations'].apply(literal_eval)","metadata":{"execution":{"iopub.status.busy":"2021-12-22T11:25:20.672207Z","iopub.execute_input":"2021-12-22T11:25:20.67362Z","iopub.status.idle":"2021-12-22T11:25:21.049356Z","shell.execute_reply.started":"2021-12-22T11:25:20.673577Z","shell.execute_reply":"2021-12-22T11:25:21.048525Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2.1 Define a dataset function","metadata":{}},{"cell_type":"code","source":"def gbl_dataset(df, img_root):\n    dataset = []\n    for i, row in df.iterrows():\n        file_name = str(img_root/\"video_{}/{}.jpg\".format(*row['image_id'].split('-')))\n        width, height = Image.open(file_name).size\n        image_id = i\n        annotations = [dict(bbox=[bbox['x'], bbox['y'], bbox['width'], bbox['height']],\n                            bbox_mode=BoxMode.XYWH_ABS,\n                            category_id=0)\n                       for bbox in row['annotations']]\n        \n        dataset.append(\n            dict(file_name=file_name,width=width,\n                 height=height,\n                 image_id=image_id,\n                 annotations=annotations\n                )\n        )\n    return dataset\n\ndef gbl_dataset_wrapper(df, img_root):\n    def wrapper():\n        return gbl_dataset(df, img_root)\n    return wrapper","metadata":{"execution":{"iopub.status.busy":"2021-12-22T11:25:21.05163Z","iopub.execute_input":"2021-12-22T11:25:21.051967Z","iopub.status.idle":"2021-12-22T11:25:21.061027Z","shell.execute_reply.started":"2021-12-22T11:25:21.05193Z","shell.execute_reply":"2021-12-22T11:25:21.060246Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"_train_df = train_df.query(\"video_id != 2\")\n_val_df = train_df.query(\"video_id == 2\")\ntrain_ds = gbl_dataset_wrapper(_train_df, Path('/kaggle/input/tensorflow-great-barrier-reef/train_images/'))\nval_ds = gbl_dataset_wrapper(_val_df, Path('/kaggle/input/tensorflow-great-barrier-reef/train_images/'))\nDatasetCatalog.pop(\"gbl_train_dataset\", None)\nDatasetCatalog.pop(\"gbl_val_dataset\", None)\nDatasetCatalog.register(\"gbl_train_dataset\", train_ds)\nDatasetCatalog.register(\"gbl_val_dataset\", val_ds)\nMetadataCatalog.get(\"gbl_train_dataset\").thing_classes = [\"starfish\"]\nMetadataCatalog.get(\"gbl_val_dataset\").thing_classes = [\"starfish\"]\n#Convert validation dataset to coco format and dump it for evaluation\nconvert_to_coco_json('gbl_val_dataset', output_file='./output/inference/gbl_val_dataset_coco_format.json', allow_cached=False)","metadata":{"execution":{"iopub.status.busy":"2021-12-22T11:25:21.062182Z","iopub.execute_input":"2021-12-22T11:25:21.062773Z","iopub.status.idle":"2021-12-22T11:27:00.367024Z","shell.execute_reply.started":"2021-12-22T11:25:21.062735Z","shell.execute_reply":"2021-12-22T11:27:00.366303Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2.2 Check the Dataset\nvisualize the dataset for verification","metadata":{}},{"cell_type":"code","source":"gbl_ds = DatasetCatalog.get(\"gbl_train_dataset\")\nmetadata = MetadataCatalog.get('gbl_train_dataset')","metadata":{"execution":{"iopub.status.busy":"2021-12-22T11:27:00.36844Z","iopub.execute_input":"2021-12-22T11:27:00.368722Z","iopub.status.idle":"2021-12-22T11:28:10.335397Z","shell.execute_reply.started":"2021-12-22T11:27:00.368686Z","shell.execute_reply":"2021-12-22T11:28:10.334085Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for data in gbl_ds:\n    if len(data['annotations']):\n        break\nim = cv2.cvtColor(cv2.imread(data['file_name']), cv2.COLOR_BGR2RGB)\nv = Visualizer(im, \n               metadata=MetadataCatalog.get('gbl_train_dataset'),\n               scale=0.5)\nout = v.draw_dataset_dict(data)\nim = Image.fromarray(out.get_image())\nim","metadata":{"execution":{"iopub.status.busy":"2021-12-22T11:28:10.336471Z","iopub.status.idle":"2021-12-22T11:28:10.337126Z","shell.execute_reply.started":"2021-12-22T11:28:10.336881Z","shell.execute_reply":"2021-12-22T11:28:10.336907Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 3.1 Define a custom Trainer to evaluate on custom dataset","metadata":{}},{"cell_type":"code","source":"class Trainer(DefaultTrainer):\n    @classmethod\n    def build_evaluator(cls, cfg, dataset_name, output_folder=None):\n        if output_folder is None:\n            output_folder = os.path.join(cfg.OUTPUT_DIR, \"inference\")\n        return COCOEvaluator(dataset_name=dataset_name,\n                             tasks=[\"bbox\"],\n                             distributed=True,\n                             output_dir=output_folder)\n    \n    @classmethod\n    def build_tta_model(cls, cfg, model):\n        return GeneralizedRCNNWithTTA(cfg, model)\n    \n    @classmethod\n    def test_with_TTA(cls, cfg, model):\n        # In the end of training, run an evaluation with TTA\n        # Only support some R-CNN models.\n        logger.info(\"Running inference with test-time augmentation ...\")\n        model = self.build_tta_model(cfg, model)\n        evaluators = [\n            cls.build_evaluator(\n                cfg, name, output_folder=os.path.join(cfg.OUTPUT_DIR, \"inference_TTA\")\n            )\n            for name in cfg.DATASETS.TEST\n        ]\n        res = cls.test(cfg, model, evaluators)\n        res = OrderedDict({k + \"_TTA\": v for k, v in res.items()})\n        return res","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 3.2 Set a config","metadata":{}},{"cell_type":"code","source":"TRAIN_STEPS = 4242 # only 4242 images with annotation\ncfg = get_cfg()\nadd_swint_config(cfg)\ncfg.merge_from_file('swin/configs/SwinT/faster_rcnn_swint_T_FPN_3x_.yaml')\ncfg.DATASETS.TRAIN = (\"gbl_train_dataset\",)\ncfg.DATASETS.TEST = (\"gbl_val_dataset\",)\ncfg.MODEL.ROI_HEADS.NUM_CLASSES = 1\ncfg.MODEL.WEIGHTS = None\ncfg.DATALOADER.NUM_WORKERS = 2\ncfg.SOLVER.IMS_PER_BATCH = 4\ncfg.SOLVER.MAX_ITER = TRAIN_STEPS * 10\ncfg.SOLVER.STEPS = []\ncfg.SOLVER.CHECKPOINT_PERIOD = TRAIN_STEPS\ncfg.TEST.EVAL_PERIOD = TRAIN_STEPS \nos.makedirs(cfg.OUTPUT_DIR, exist_ok=True)\ntrainer = Trainer(cfg)\ntrainer.resume_or_load(resume=False)","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2021-12-22T11:28:17.302351Z","iopub.execute_input":"2021-12-22T11:28:17.304317Z","iopub.status.idle":"2021-12-22T11:29:45.156664Z","shell.execute_reply.started":"2021-12-22T11:28:17.303234Z","shell.execute_reply":"2021-12-22T11:29:45.155906Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 3.3 Train\n\nI could only run for 5 epochs due to runtime quota","metadata":{}},{"cell_type":"code","source":"\"\"\"\n# load pretrained weights\nother_weights = torch.load('faster_rcnn_swint_T.pth')['model']\nself_weight = trainer.model.state_dict()\nfor name, param in self_weight.items():\n    if name in other_weights:\n        if other_weights[name].shape == param.shape:\n            self_weight[name] = other_weights[name]\n        else:\n            print(f\"size mismatch at {name}\")\n    else:\n        print(f\"layer {name} does not exist\")\ntrainer.model.load_state_dict(self_weight)\ntrainer.train()\n\"\"\"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 4.1 Evaluate\nevaluate on validation set","metadata":{}},{"cell_type":"code","source":"trainer.model.load_state_dict(torch.load('model_0021209.pth')['model'])\ntrainer.test(cfg, trainer.model)","metadata":{"execution":{"iopub.status.busy":"2021-12-22T11:29:54.939443Z","iopub.execute_input":"2021-12-22T11:29:54.941179Z","iopub.status.idle":"2021-12-22T11:29:56.986727Z","shell.execute_reply.started":"2021-12-22T11:29:54.941116Z","shell.execute_reply":"2021-12-22T11:29:56.985409Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 4.2 Visualize a few prediction examples","metadata":{}},{"cell_type":"code","source":"val_ds = DatasetCatalog.get(\"gbl_val_dataset\")\ntrainer.model.eval()\nmetadata = MetadataCatalog.get('gbl_train_dataset')","metadata":{"execution":{"iopub.status.busy":"2021-12-22T11:35:25.514546Z","iopub.execute_input":"2021-12-22T11:35:25.514817Z","iopub.status.idle":"2021-12-22T11:35:25.525685Z","shell.execute_reply.started":"2021-12-22T11:35:25.514787Z","shell.execute_reply":"2021-12-22T11:35:25.524939Z"},"scrolled":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for _ in range(5):\n    idx = np.random.randint(0, len(val_ds))\n    data = val_ds[idx]\n    print(data['file_name'])\n    im = cv2.imread(data['file_name'])\n    im_tensor = torch.from_numpy(im).permute(2,0,1)  # h, w, c -> c, h, w\n    h, w, _ = im.shape\n    with torch.no_grad():\n        pred = trainer.model([{\"image\": im_tensor.cuda(), \"width\": w, \"height\": h}])\n    v = Visualizer(im[:, :, ::-1],\n                   metadata=metadata, \n                   scale=0.5, \n                   instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels. This option is only available for segmentation models\n    )\n    out = v.draw_instance_predictions(pred[0][\"instances\"].to(\"cpu\"))\n    plt.figure()\n    plt.imshow(out.get_image())","metadata":{"execution":{"iopub.status.busy":"2021-12-22T11:46:40.767662Z","iopub.execute_input":"2021-12-22T11:46:40.768458Z","iopub.status.idle":"2021-12-22T11:46:43.058513Z","shell.execute_reply.started":"2021-12-22T11:46:40.768417Z","shell.execute_reply":"2021-12-22T11:46:43.057864Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}