{"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":"#  Mask RCNN w/ ResnetV2_50 Pre-Trained on LiveCell","metadata":{}},{"cell_type":"markdown","source":"# Libraries","metadata":{}},{"cell_type":"code","source":"!pip install pycocotools\n!pip install git+https://github.com/ildoonet/pytorch-gradual-warmup-lr.git","metadata":{"execution":{"iopub.status.busy":"2021-12-04T18:14:04.290157Z","iopub.execute_input":"2021-12-04T18:14:04.29049Z","iopub.status.idle":"2021-12-04T18:14:28.9361Z","shell.execute_reply.started":"2021-12-04T18:14:04.290408Z","shell.execute_reply":"2021-12-04T18:14:28.935203Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"import os\nimport json\nimport torch\nimport torchvision\nimport numpy as np\nimport pandas as pd\nimport albumentations as A\n\nfrom PIL import Image\nfrom pathlib import Path\nfrom pycocotools.coco import COCO\nfrom albumentations.pytorch import ToTensorV2\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection.mask_rcnn import MaskRCNNPredictor","metadata":{"execution":{"iopub.status.busy":"2021-12-04T18:14:28.938049Z","iopub.execute_input":"2021-12-04T18:14:28.938301Z","iopub.status.idle":"2021-12-04T18:14:32.525695Z","shell.execute_reply.started":"2021-12-04T18:14:28.938268Z","shell.execute_reply":"2021-12-04T18:14:32.524977Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Constants & Config\nThe LIVECell train annotations is not a `list` of `dict` but instead a `dict` with keys being the image ids.  COCO expects it to be a `list` of `dict`.  The function `fix_annotations` gets the annotations file and fixes this.  If the `annotations` key is an instance of `list` it just returns the same `json` file otherwise it creates a new one with the correct type to be loaded by COCO.","metadata":{}},{"cell_type":"code","source":"def fix_annotations(anns_file):\n    with open(anns_file, 'r') as f:\n        data = json.load(f)\n    imgs_anns = data[\"annotations\"]\n    if isinstance(imgs_anns, list):\n        print(\"data['annotations'] is a list so using original file\")\n        return anns_file\n    fixed_ann = []\n    for key in imgs_anns:\n        img_anns = imgs_anns[key]\n        fixed_ann.append(img_anns)\n    data[\"annotations\"] = fixed_ann\n    with open(f\"./{anns_file.split('/')[-1]}\", 'w') as f:\n        json.dump(data, f)\n    return os.path.abspath(f\"./{anns_file.split('/')[-1]}\")\n\nDATASET_PATHS = {\n    \"train\": \"../input/sartorius-cell-instance-segmentation/LIVECell_dataset_2021/annotations/LIVECell/livecell_coco_train.json\",\n    \"val\": \"../input/sartorius-cell-instance-segmentation/LIVECell_dataset_2021/annotations/LIVECell/livecell_coco_val.json\",\n    \"test\": \"../input/sartorius-cell-instance-segmentation/LIVECell_dataset_2021/annotations/LIVECell/livecell_coco_test.json\"\n}\n\nfor path in DATASET_PATHS:\n    DATASET_PATHS[path] = fix_annotations(DATASET_PATHS[path])\n    \nLIVECELL_IMAGES_ROOT = \"../input/sartorius-cell-instance-segmentation/LIVECell_dataset_2021/images\"","metadata":{"execution":{"iopub.status.busy":"2021-12-04T18:17:48.0872Z","iopub.execute_input":"2021-12-04T18:17:48.087453Z","iopub.status.idle":"2021-12-04T18:21:26.509961Z","shell.execute_reply.started":"2021-12-04T18:17:48.087425Z","shell.execute_reply":"2021-12-04T18:21:26.509181Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"IMAGE_RESIZE = (224, 224)\nRESNET_MEAN = (0.485, 0.456, 0.406)\nRESNET_STD = (0.229, 0.224, 0.225)\n\nBATCH_SIZE = 2\nDEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\n# Optimizer\nMOMENTUM = 0.9\nLEARNING_RATE = 3e-3\nWEIGHT_DECAY = 0.0005\n\nNUM_EPOCHS = 50","metadata":{"execution":{"iopub.status.busy":"2021-12-04T18:21:32.082489Z","iopub.execute_input":"2021-12-04T18:21:32.082749Z","iopub.status.idle":"2021-12-04T18:21:32.137376Z","shell.execute_reply.started":"2021-12-04T18:21:32.08272Z","shell.execute_reply":"2021-12-04T18:21:32.136583Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def rle_decode(mask_rle):\n    \n    mask_rle = np.array(mask_rle.split(), dtype=np.int)\n    pixels = mask_rle.reshape(-1, 2)\n#     assert len(start) == len(length)\n    pixels[:, 0] -= 1\n    mask = np.zeros(IMAGE_SHAPE[0] * IMAGE_SHAPE[1])\n    for pixel in pixels:\n        mask[pixel[0]:pixel[0] + pixel[1]] = 1\n    return mask.reshape(IMAGE_SHAPE)\n\ndef prepare_image_mask(mask_annotations):\n    \n    mask = np.zeros(IMAGE_SHAPE)\n    for mask_annotation in mask_annotations:\n        mask += rle_decode(mask_annotation)\n        \n    mask = mask.clip(0, 1)\n    return mask\n\ndef compute_iou(labels, y_pred):\n    \"\"\"\n    Computes the IoU for instance labels and predictions.\n\n    Args:\n        labels (np array): Labels.\n        y_pred (np array): predictions\n\n    Returns:\n        np array: IoU matrix, of size true_objects x pred_objects.\n    \"\"\"\n\n    true_objects = len(np.unique(labels))\n    pred_objects = len(np.unique(y_pred))\n\n    # Compute intersection between all objects\n    intersection = np.histogram2d(\n        labels.flatten(), y_pred.flatten(), bins=(true_objects, pred_objects)\n    )[0]\n\n    # Compute areas (needed for finding the union between all objects)\n    area_true = np.histogram(labels, bins=true_objects)[0]\n    area_pred = np.histogram(y_pred, bins=pred_objects)[0]\n    area_true = np.expand_dims(area_true, -1)\n    area_pred = np.expand_dims(area_pred, 0)\n\n    # Compute union\n    union = area_true + area_pred - intersection\n    iou = intersection / union\n    \n    return iou[1:, 1:]  # exclude background\n\ndef precision_at(threshold, iou):\n    \"\"\"\n    Computes the precision at a given threshold.\n\n    Args:\n        threshold (float): Threshold.\n        iou (np array): IoU matrix.\n\n    Returns:\n        int: Number of true positives,\n        int: Number of false positives,\n        int: Number of false negatives.\n    \"\"\"\n    matches = iou > threshold\n    true_positives = np.sum(matches, axis=1) == 1  # Correct objects\n    false_positives = np.sum(matches, axis=0) == 0  # Missed objects\n    false_negatives = np.sum(matches, axis=1) == 0  # Extra objects\n    tp, fp, fn = (\n        np.sum(true_positives),\n        np.sum(false_positives),\n        np.sum(false_negatives),\n    )\n    return tp, fp, fn\n\ndef iou_map(truths, preds, verbose=0):\n    \"\"\"\n    Computes the metric for the competition.\n    Masks contain the segmented pixels where each object has one value associated,\n    and 0 is the background.\n\n    Args:\n        truths (list of masks): Ground truths.\n        preds (list of masks): Predictions.\n        verbose (int, optional): Whether to print infos. Defaults to 0.\n\n    Returns:\n        float: mAP.\n    \"\"\"\n    ious = [compute_iou(truth, pred) for truth, pred in zip(truths, preds)]\n\n    if verbose:\n        print(\"Thresh\\tTP\\tFP\\tFN\\tPrec.\")\n\n    prec = []\n    for t in np.arange(0.5, 1.0, 0.05):\n        tps, fps, fns = 0, 0, 0\n        for iou in ious:\n            tp, fp, fn = precision_at(t, iou)\n            tps += tp\n            fps += fp\n            fns += fn\n\n        p = tps / (tps + fps + fns)\n        prec.append(p)\n\n        if verbose:\n            print(\"{:1.3f}\\t{}\\t{}\\t{}\\t{:1.3f}\".format(t, tps, fps, fns, p))\n\n    if verbose:\n        print(\"AP\\t-\\t-\\t-\\t{:1.3f}\".format(np.mean(prec)))\n\n    return np.mean(prec)","metadata":{"execution":{"iopub.status.busy":"2021-12-04T18:21:38.50721Z","iopub.execute_input":"2021-12-04T18:21:38.507647Z","iopub.status.idle":"2021-12-04T18:21:38.525169Z","shell.execute_reply.started":"2021-12-04T18:21:38.507607Z","shell.execute_reply":"2021-12-04T18:21:38.524261Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data Augmentations","metadata":{}},{"cell_type":"code","source":"transforms_train = A.Compose([\n#   A.HorizontalFlip(p=0.5),\n#   A.VerticalFlip(p=0.5),\n  A.Resize(IMAGE_RESIZE[0], IMAGE_RESIZE[1]),\n  A.Normalize(mean=RESNET_MEAN, std=RESNET_STD),\n  ToTensorV2()\n], bbox_params=A.BboxParams(min_visibility=0.1, format=\"pascal_voc\", label_fields=[\"class_labels\"]))\n\ntransforms_valid = A.Compose([\n  A.Resize(IMAGE_RESIZE[0], IMAGE_RESIZE[1]),\n  A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n  ToTensorV2()\n])\n    \ndef get_transform(train):\n    if train: return transforms_train\n    return transforms_valid","metadata":{"execution":{"iopub.status.busy":"2021-12-04T18:21:40.312177Z","iopub.execute_input":"2021-12-04T18:21:40.31283Z","iopub.status.idle":"2021-12-04T18:21:40.31939Z","shell.execute_reply.started":"2021-12-04T18:21:40.312792Z","shell.execute_reply":"2021-12-04T18:21:40.318717Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torchvision import transforms\n\nclass SartoriusDataset(Dataset):\n\n    def __init__(self, root, anns_file, transforms=None):\n        self.root = root\n        self.coco = COCO(anns_file)\n        self.transforms = transforms\n        self.img_ids = list(sorted(self.coco.imgs.keys()))\n\n    def __len__(self):\n        return len(self.img_ids)\n\n    def __getitem__(self, idx):\n        img_id = self.img_ids[idx]\n        anns_ids = self.coco.getAnnIds(imgIds=img_id)\n        img_name = self.coco.loadImgs(ids=img_id)[0][\"file_name\"]\n        img_folder = img_name.split('_')[0]\n        img = Image.open(os.path.join(self.root, img_folder, img_name)).convert(\"RGB\")\n        img_anns = self.coco.loadAnns(ids=anns_ids)\n\n        num_objs = len(img_anns)\n        masks = []\n        bboxes = []\n        labels = []\n        areas = []\n        masks = []\n        iscrowd = []\n\n        for i in range(num_objs):\n            xmin = img_anns[i][\"bbox\"][0]\n            ymin = img_anns[i][\"bbox\"][1]\n            xmax = xmin + img_anns[i][\"bbox\"][2]\n            ymax = ymin + img_anns[i][\"bbox\"][3]\n\n            bboxes.append([xmin, ymin, xmax, ymax])\n            areas.append(img_anns[i][\"area\"])\n            masks.append(self.coco.annToMask(img_anns[i]))\n\n        is_crowd = torch.zeros((num_objs,), dtype=torch.int64)\n        image_id = torch.tensor([img_id])\n        labels = torch.ones((num_objs,), dtype=torch.int64)\n        areas = torch.as_tensor(areas, dtype=torch.float32)\n\n        if self.transforms:\n            augmented = self.transforms(image=img, masks=masks, bboxes=bboxes, class_labels=[\"cell\"] * len(bboxes))\n            image = augmented[\"image\"]\n            masks = augmented[\"masks\"]\n            bboxes = augmented[\"bboxes\"]\n        else:\n            image = transforms.Compose([\n                transforms.ToTensor(),\n                transforms.Resize(IMAGE_RESIZE)\n            ])(img)\n        masks = torch.as_tensor(masks, dtype=torch.uint8)\n        bboxes = torch.as_tensor(bboxes, dtype=torch.float32)\n        target = {}\n        target[\"boxes\"] = bboxes\n        target[\"masks\"] = masks\n        target[\"labels\"] = labels\n        target[\"image_id\"] = image_id\n        target[\"area\"] = areas\n        target[\"iscrowd\"] = is_crowd\n\n        return image, target","metadata":{"execution":{"iopub.status.busy":"2021-12-04T18:21:41.997754Z","iopub.execute_input":"2021-12-04T18:21:41.998332Z","iopub.status.idle":"2021-12-04T18:21:42.014454Z","shell.execute_reply.started":"2021-12-04T18:21:41.998295Z","shell.execute_reply":"2021-12-04T18:21:42.013794Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset = SartoriusDataset(os.path.join(LIVECELL_IMAGES_ROOT, \"livecell_train_val_images\"), DATASET_PATHS[\"train\"])\ndl_train = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, \n                      num_workers=0, collate_fn=lambda x: tuple(zip(*x)))","metadata":{"execution":{"iopub.status.busy":"2021-12-04T18:26:45.147652Z","iopub.execute_input":"2021-12-04T18:26:45.148565Z","iopub.status.idle":"2021-12-04T18:27:15.500656Z","shell.execute_reply.started":"2021-12-04T18:26:45.148512Z","shell.execute_reply":"2021-12-04T18:27:15.499935Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Override pythorch checkpoint with an \"offline\" version of the file\n!mkdir -p /root/.cache/torch/hub/checkpoints/\n!cp ../input/cocopre/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth /root/.cache/torch/hub/checkpoints/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth","metadata":{"execution":{"iopub.status.busy":"2021-12-04T18:27:15.504783Z","iopub.execute_input":"2021-12-04T18:27:15.505476Z","iopub.status.idle":"2021-12-04T18:27:20.15574Z","shell.execute_reply.started":"2021-12-04T18:27:15.505435Z","shell.execute_reply":"2021-12-04T18:27:20.154745Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_model():\n    # This is just a dummy value for the classification head\n    NUM_CLASSES = 2\n    \n    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True, image_mean=RESNET_MEAN, image_std=RESNET_STD)\n\n    # get the number of input features for the classifier\n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n    # replace the pre-trained head with a new one\n    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, NUM_CLASSES)\n\n    # now get the number of input features for the mask classifier\n    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels\n    hidden_layer = 256\n    # and replace the mask predictor with a new one\n    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, NUM_CLASSES)\n    return model\n\n\n# Get the Mask R-CNN model\n# The model does classification, bounding boxes and MASKs for individuals, all at the same time\n# We only care about MASKS\nmodel = get_model()\nmodel.to(DEVICE)\n\n# TODO: try removing this for\nfor param in model.parameters():\n    param.requires_grad = True\n    \nmodel.train()","metadata":{"execution":{"iopub.status.busy":"2021-12-04T18:27:20.157408Z","iopub.execute_input":"2021-12-04T18:27:20.157681Z","iopub.status.idle":"2021-12-04T18:27:21.183527Z","shell.execute_reply.started":"2021-12-04T18:27:20.157642Z","shell.execute_reply":"2021-12-04T18:27:21.182856Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import time\nimport gc\nfrom warmup_scheduler import GradualWarmupScheduler\nparams = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.Adam(params, lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)\n\nlr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=5)\nscheduler_warmup = GradualWarmupScheduler(optimizer, multiplier=1, total_epoch=5, after_scheduler=lr_scheduler)\nn_batches = len(dl_train)\n# Zero gradient removes a warning\noptimizer.zero_grad()\noptimizer.step()\nfor epoch in range(1, NUM_EPOCHS + 1):\n    print(f\"Starting epoch {epoch} of {NUM_EPOCHS}\")\n    \n    time_start = time.time()\n    loss_accum = 0.0\n    loss_mask_accum = 0.0\n    scheduler_warmup.step()\n    for batch_idx, (images, targets) in enumerate(dl_train, 1):\n    \n        # Predict\n        images = list(image.to(DEVICE) for image in images)\n        targets = [{k: v.to(DEVICE) for k, v in t.items()} for t in targets]\n\n        loss_dict = model(images, targets)\n        loss = sum(loss for loss in loss_dict.values())\n        \n        # Backprop518\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n        \n        # Logging\n        loss_mask = loss_dict['loss_mask'].item()\n        loss_accum += loss.item()\n        loss_mask_accum += loss_mask\n        \n        if batch_idx % 50 == 0:\n            print(f\"    [Batch {batch_idx:3d} / {n_batches:3d}] Batch train loss: {loss.item():7.3f}. Mask-only loss: {loss_mask:7.3f}\")\n        del images\n        del targets\n        gc.collect()\n    lr_scheduler.step()\n    \n    # Train losses\n    train_loss = loss_accum / n_batches\n    train_loss_mask = loss_mask_accum / n_batches\n    \n    \n    elapsed = time.time() - time_start\n    \n    \n    torch.save(model.state_dict(), f\"pytorch_model-e{epoch}.bin\")\n    prefix = f\"[Epoch {epoch:2d} / {NUM_EPOCHS:2d}]\"\n    print(f\"{prefix} Train mask-only loss: {train_loss_mask:7.3f}\")\n    print(f\"{prefix} Train loss: {train_loss:7.3f}. [{elapsed:.0f} secs]\")","metadata":{"execution":{"iopub.status.busy":"2021-12-04T18:27:21.185392Z","iopub.execute_input":"2021-12-04T18:27:21.18582Z","iopub.status.idle":"2021-12-04T18:30:48.687035Z","shell.execute_reply.started":"2021-12-04T18:27:21.185781Z","shell.execute_reply":"2021-12-04T18:30:48.685731Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}