{"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":"# 🦠 Sartorius - Starter Torch Mask R-CNN\n\n_Forked from [julian3833](https://www.kaggle.com/julian3833/sartorius-starter-torch-mask-r-cnn-lb-0-173):_\n\n> Following [this discussion thread](https://www.kaggle.com/c/sartorius-cell-instance-segmentation/discussion/279790), in this notebook we build a base starter Mask R-CNN with pytorch.\n>\n> The code is an adapted version from [this notebook](https://www.kaggle.com/abhishek/mask-rcnn-using-torchvision-0-17/) by the first quadruple kaggle grandmaster [Abishek](https://www.kaggle.com/abhishek).\n>\n> The [previous U-net model](https://www.kaggle.com/julian3833/sartorius-starter-baseline-torch-u-net), which I was expecting to enter a steep improvement regime with quick-wins, hit a ceiling at `0.03`, no matter what changes I performed 🥲.\nData augmentation, changes in the architecture, and other changes didn't work. The suggestion that semantic segmentation doesn't work seems reasonable, since the individuals cannot be split by connected components, as they overlap heavily.\n>\n> This is a follow up notebook with a Mask R-CNN, which was proposed by one of the top competitors ([Inoichan](https://www.kaggle.com/inoueu1)) as a more suitable architecture for this task.\n>\n> I'm not very familiar with the architecture, but it seems that it is the state-of-the art for \"instance segmentation\".\n> It classifies individuals, gets bounding boxes around them and, most importantly, provide a separated mask for each of them.\n>\n> You can read more about it [here](https://viso.ai/deep-learning/mask-r-cnn/).\n>\n> This model predicts different masks for different individual, rather that an unique mask for the whole picture and thus is better for the address the problem at hand.\n>\n> At the end, any overlapping pixel is removed, to ensure the non-overlapping policy. That wasn't required with the U-net, since the output was only one unique mask and therefore no overlap could have happened.","metadata":{"id":"eERgvcFD-QEe"}},{"cell_type":"markdown","source":"## Notes from Inoichan\n\n[Tips in submission and baseline](https://www.kaggle.com/c/sartorius-cell-instance-segmentation/discussion/279790):\n\n1. This is an \"Instance Segmentation\" problem, not a \"semantic segmentation\". [Mask R-CNN](https://viso.ai/deep-learning/mask-r-cnn/) is a good writeup on the difference.\n2. The mask encoding is l->r then t->b, not t->b then l->r.\n3. The masks cannot overlap or the submission will fail.\n\n_There are some helper functions in this notebook for dealing with the RLE of the masks._","metadata":{"id":"k1e_IRJH-QEj"}},{"cell_type":"markdown","source":"## Evan's Changes\n\nThese are changes I'm making from the original notebook.\n\n1. Adding [Torch Lightning](https://pytorch-lightning.readthedocs.io/en/latest/), using [this tutorial](https://www.aicrowd.com/showcase/tutorial-with-pytorch-torchvision-and-pytorch-lightning) from [aicrowd.com](aicrowd.com). (2021-10-20)\n2. Update module to use [Logging](https://pytorch-lightning.readthedocs.io/en/stable/extensions/logging.html#logging-from-a-lightningmodule) API to log metrics.\n3. Work on training just the head first, then training the whole stack.\n  a. First just the head. Hopefully.\n4. Also actually use epochs limit.\n5. Update to work on colab as well as kaggle.\n6. Implement [mean of IOUs](https://www.kaggle.com/c/sartorius-cell-instance-segmentation/overview/evaluation), use it for `val_accuracy` instead of a `val_loss`.\n7. Monitor `val_loss` for checkpoints, optimizers\n8. Add [CLAHE](https://scikit-image.org/docs/dev/api/skimage.exposure.html?highlight=clahe#skimage.exposure.equalize_adapthist) preprocessing.\n\n## Notes\n\n- Batch size 8 -> OOM\n\n### ADAM\n\n- currently sort of just hovers around 1.3 loss, trying different things to get out of the hole.\n- lr 0.01 -> loss nan\n- switch to adam. I don't know why it wasn't adam\n- Increase batch size to 8, OOM\n- Batch size to 4, so far so good\n- lr -> 0.005: terrible convergence\n- Using `val_loss` as something to maximize isn't great...\n- 6 epochs instead of 3 or 9\n- Turn wandb into a context manager to reuse the session across training runs\n- AMSGRAD flag, lr scheduler, seems to do a bit better than SGD even.\n- Batch size 6 -> OOM\n\n### SGD\n\n- Trying [SGD](https://shaoanlu.wordpress.com/2017/05/29/sgd-all-which-one-is-the-best-optimizer-dogs-vs-cats-toy-experiment/) again, with lr 0.0001\n- LR -> 0.001 w SGD - With SGD the masks visually look better. ADAM had weird block artifacts, but SGD just has masks.\n- LR -> 0.01 - Let's see if it'll converge faster. (Basically has the same loss curve)\n- Try [CosineAnnealingWarmRestarts](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CosineAnnealingWarmRestarts.html) scheduler\n- Use [levbszabo's SGD params](https://github.com/levbszabo/mask-rcnn)\n- Qualitatively, SGD produces much sharper masks.\n- lr -> 0.005 nan loss\n- lr -> 0.001\n- clahe the images before train/eval in the loader\n- Add option for including difference image","metadata":{"id":"jFV-h7cB-QEj"}},{"cell_type":"markdown","source":"### Configuration\n\nHere's the overall configuration for this run. It's sent to wandb as part of the report so this run can be compared to others.","metadata":{"id":"oRbH1hqrghzq"}},{"cell_type":"code","source":"config = dict (\n    project = \"sartiorius-cell-instance-segmentation\",\n    architecture = \"maskrcnn_resnet50_fpn\",\n    dataset_id = \"sartorius-cell-instance-segmentation\",\n    infra = \"kaggle\",\n    lr=0.01,\n    min_lr=0.0000001,\n    epochs=15,\n    batch_size=4,\n    nesterov=True,\n    momentum=0.9,\n    weight_decay=0.0005,\n    clip_limit=0.25,\n    difference=False,\n    notes=\"Clahe.\"\n)","metadata":{"id":"7wzHTpvlggMj","execution":{"iopub.status.busy":"2021-11-10T15:33:15.279025Z","iopub.execute_input":"2021-11-10T15:33:15.280036Z","iopub.status.idle":"2021-11-10T15:33:15.332574Z","shell.execute_reply.started":"2021-11-10T15:33:15.279989Z","shell.execute_reply":"2021-11-10T15:33:15.331826Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%load_ext autoreload\n%autoreload 2\n\nimport os\nfrom pathlib import Path\n\nINTERNET = config[\"infra\"] == \"colab\"\nif INTERNET:\n    !pip install wandb clahe --upgrade\n    \n    import wandb","metadata":{"id":"rXdS6rT9-QEk","outputId":"921a09f7-c2ad-490b-e77a-18046b2bea3a","execution":{"iopub.status.busy":"2021-11-10T15:33:15.334423Z","iopub.execute_input":"2021-11-10T15:33:15.334682Z","iopub.status.idle":"2021-11-10T15:33:15.385292Z","shell.execute_reply.started":"2021-11-10T15:33:15.334648Z","shell.execute_reply":"2021-11-10T15:33:15.384466Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Input Configuration\n\nSince Colab and Kaggle have different input paths, we configure them appropriately here. Always use [`pathlib`](https://docs.python.org/3/library/pathlib.html).","metadata":{"id":"pVUTSqNQirYU"}},{"cell_type":"code","source":"if config[\"infra\"] == \"colab\":\n    from google.colab import drive\n    drive.mount(\"/content/gdrive\")\n    INPUT_ROOT = Path(\"/content/gdrive/MyDrive/kaggle/input\")\n    !pip install pytorch-lightning\n    with (INPUT_ROOT / Path(\"wandb.txt\")).open(\"r\") as wf:\n      wandb_key = wf.read().strip()\n\n    os.environ[\"WANDB_API_KEY\"] = wandb_key\nelse:\n  INPUT_ROOT = Path(\"../input\")\n\n!nvidia-smi ","metadata":{"id":"okuGpTG7BafE","outputId":"dac4abb8-9e56-4cb5-9617-bb095d5e22b9","execution":{"iopub.status.busy":"2021-11-10T15:33:15.386916Z","iopub.execute_input":"2021-11-10T15:33:15.387197Z","iopub.status.idle":"2021-11-10T15:33:16.136061Z","shell.execute_reply.started":"2021-11-10T15:33:15.387150Z","shell.execute_reply":"2021-11-10T15:33:16.135226Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from contextlib import contextmanager\n\n@contextmanager\ndef wandb_context(configuration=config):\n  if INTERNET:\n    run = wandb.init(reinit=True, config=config, project=config[\"project\"])\n    try:\n      yield run\n    finally:\n      wandb.finish()\n  else:\n    yield None\n","metadata":{"id":"Zw0bVjfptV80","execution":{"iopub.status.busy":"2021-11-10T15:33:16.137792Z","iopub.execute_input":"2021-11-10T15:33:16.138054Z","iopub.status.idle":"2021-11-10T15:33:16.190029Z","shell.execute_reply.started":"2021-11-10T15:33:16.138019Z","shell.execute_reply":"2021-11-10T15:33:16.189210Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"","metadata":{"id":"K8C-5sB0-QEl"}},{"cell_type":"markdown","source":"# Imports and constants","metadata":{"id":"s0df6o4I-QEm"}},{"cell_type":"markdown","source":"First define some transforms for the images. This used to be in a separate file but then I edited it and didn't want to recreate a dataset.","metadata":{"id":"h0Uj7WT6GKM_"}},{"cell_type":"code","source":"import random\nimport torch\nfrom torchvision.transforms import functional as F\nfrom skimage.exposure import equalize_adapthist\nimport numpy as np\n\ndef _flip_coco_person_keypoints(kps, width):\n    flip_inds = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]\n    flipped_data = kps[:, flip_inds]\n    flipped_data[..., 0] = width - flipped_data[..., 0]\n    # Maintain COCO convention that if visibility == 0, then x, y = 0\n    inds = flipped_data[..., 2] == 0\n    flipped_data[inds] = 0\n    return flipped_data\n\nclass Compose(object):\n    def __init__(self, transforms):\n        self.transforms = transforms\n\n    def __call__(self, image, target):\n        for t in self.transforms:\n            image, target = t(image, target)\n        return image, target\n\nclass ImageEqualize:    \n    def __init__(self, clip_limit, difference):\n        self._clip_limit = clip_limit\n        self._difference = difference\n\n    def equalize(self, img):\n        img = np.array(img)\n        eimg = equalize_adapthist(img, clip_limit=self._clip_limit)  \n\n        if self._difference:\n            img = img - eimg\n            img = img - img.min()\n            img = img / img.max()\n        else:\n            return eimg\n\n    def __call__(self, image, target):\n        return self.equalize(image), target\n\nclass RandomHorizontalFlip(object):\n    def __init__(self, prob):\n        self.prob = prob\n\n    def __call__(self, image, target):\n        if random.random() < self.prob:\n            height, width = image.shape[-2:]\n            image = image.flip(-1)\n            bbox = target[\"boxes\"]\n            bbox[:, [0, 2]] = width - bbox[:, [2, 0]]\n            target[\"boxes\"] = bbox\n            if \"masks\" in target:\n                target[\"masks\"] = target[\"masks\"].flip(-1)\n            if \"keypoints\" in target:\n                keypoints = target[\"keypoints\"]\n                keypoints = _flip_coco_person_keypoints(keypoints, width)\n                target[\"keypoints\"] = keypoints\n        return image, target\n\nclass ToTensor(object):\n    def __call__(self, image, target):\n        image = F.to_tensor(np.array(image, dtype=np.float))\n        return image, target\n","metadata":{"id":"bqBZg0BeGInb","execution":{"iopub.status.busy":"2021-11-10T15:33:16.192683Z","iopub.execute_input":"2021-11-10T15:33:16.192994Z","iopub.status.idle":"2021-11-10T15:33:16.249957Z","shell.execute_reply.started":"2021-11-10T15:33:16.192958Z","shell.execute_reply":"2021-11-10T15:33:16.249215Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nimport os\nimport random\nimport collections\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\nimport torch\nimport torchvision\nfrom torchvision.transforms import ToPILImage\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection.mask_rcnn import MaskRCNNPredictor\nfrom torchvision.models.detection._utils import Matcher\nfrom torchvision.ops.boxes import box_iou\n\nfrom PIL import Image, ImageFile\nfrom pathlib import Path\n\nMASKRCNN_UTILS_PATH = INPUT_ROOT / Path(\"maskrcnn-utils/\")\n\n# We only use 3 transformations from this package\nsys.path.append(str(MASKRCNN_UTILS_PATH))\n\ndef fix_all_seeds(seed):\n    np.random.seed(seed)\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    \nfix_all_seeds(2021)","metadata":{"id":"gJgH4K8R-QEm","execution":{"iopub.status.busy":"2021-11-10T15:33:16.251161Z","iopub.execute_input":"2021-11-10T15:33:16.251528Z","iopub.status.idle":"2021-11-10T15:33:16.301503Z","shell.execute_reply.started":"2021-11-10T15:33:16.251492Z","shell.execute_reply":"2021-11-10T15:33:16.300802Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\n\nSAMPLE_SUBMISSION  = str(INPUT_ROOT / Path('sartorius-cell-instance-segmentation/sample_submission.csv'))\nTRAIN_CSV = str(INPUT_ROOT / Path(\"sartorius-cell-instance-segmentation/train.csv\"))\nTRAIN_PATH = str(INPUT_ROOT / Path(\"sartorius-cell-instance-segmentation/train\"))\nTEST_PATH = str(INPUT_ROOT / Path(\"sartorius-cell-instance-segmentation/test\"))\n\n\nNUM_EPOCHS = config[\"epochs\"]","metadata":{"id":"MqQRKw1v-QEo","execution":{"iopub.status.busy":"2021-11-10T15:33:16.302764Z","iopub.execute_input":"2021-11-10T15:33:16.303247Z","iopub.status.idle":"2021-11-10T15:33:16.348162Z","shell.execute_reply.started":"2021-11-10T15:33:16.303210Z","shell.execute_reply":"2021-11-10T15:33:16.347191Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Traning Dataset","metadata":{"id":"nV1IZjGu-QEp"}},{"cell_type":"markdown","source":"## Utilities","metadata":{"id":"46elCGvL-QEp"}},{"cell_type":"code","source":"def rle_decode(mask_rle, shape, color=1):\n    '''\n    mask_rle: run-length as string formated (start length)\n    shape: (height,width) of array to return \n    Returns numpy array, 1 - mask, 0 - background\n\n    '''\n    s = mask_rle.split()\n    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n    starts -= 1\n    ends = starts + lengths\n    img = np.zeros(shape[0] * shape[1], dtype=np.float32)\n    for lo, hi in zip(starts, ends):\n        img[lo : hi] = color\n    return img.reshape(shape)\n\n\ndef get_transform(train):\n    transforms = [ImageEqualize(clip_limit=config[\"clip_limit\"], difference=config[\"difference\"]),\n                  ToTensor()]\n\n    if train:\n        # during training, randomly flip the training images\n        # and ground-truth for data augmentation\n        transforms.append(RandomHorizontalFlip(0.5))\n    return Compose(transforms)","metadata":{"id":"Tkn3gRO9-QEq","execution":{"iopub.status.busy":"2021-11-10T15:33:16.351595Z","iopub.execute_input":"2021-11-10T15:33:16.351831Z","iopub.status.idle":"2021-11-10T15:33:16.398638Z","shell.execute_reply.started":"2021-11-10T15:33:16.351806Z","shell.execute_reply":"2021-11-10T15:33:16.397949Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training Dataset and DataLoader","metadata":{"id":"2XHVn45D-QEq"}},{"cell_type":"markdown","source":"First a helper function to run clahe on the images to enhance the contrast.","metadata":{"id":"DfzkdaeiB7PM"}},{"cell_type":"code","source":"class CellDataset(Dataset):\n    def __init__(self, image_dir, df_path, height, width, transforms=None):\n        self.transforms = transforms\n        self.image_dir = image_dir\n        self.df = pd.read_csv(df_path)\n        self.height = height\n        self.width = width\n        self.image_info = collections.defaultdict(dict)\n        temp_df = self.df.groupby('id')['annotation'].agg(lambda x: list(x)).reset_index()\n        for index, row in temp_df.iterrows():\n            self.image_info[index] = {\n                    'image_id': row['id'],\n                    'image_path': os.path.join(self.image_dir, row['id'] + '.png'),\n                    'annotations': row[\"annotation\"]\n                    }\n            \n    def __getitem__(self, idx):\n        # load images ad masks\n        img_path = self.image_info[idx][\"image_path\"]\n        img = Image.open(img_path).convert(\"RGB\")\n        img = np.array(img, dtype=np.float)\n        img = img - img.min()\n        img = img / img.max()\n        #img = img.resize((self.width, self.height), resample=Image.BILINEAR)\n\n        info = self.image_info[idx]\n\n        mask = np.zeros((len(info['annotations']), self.width, self.height), dtype=np.uint8)\n        labels = []\n        \n        for m, annotation in enumerate(info['annotations']):\n            sub_mask = rle_decode(annotation, (520, 704))\n            sub_mask = Image.fromarray(sub_mask)\n            #sub_mask = sub_mask.resize((self.width, self.height), resample=Image.BILINEAR)\n            sub_mask = np.array(sub_mask) > 0\n            mask[m, :, :] = sub_mask\n            labels.append(1)\n\n        num_objs = len(labels)\n        boxes = []\n        new_labels = []\n        new_masks = []\n\n        for i in range(num_objs):\n            try:\n                pos = np.where(mask[i, :, :])\n                xmin = np.min(pos[1])\n                xmax = np.max(pos[1])\n                ymin = np.min(pos[0])\n                ymax = np.max(pos[0])\n                boxes.append([xmin, ymin, xmax, ymax])\n                new_labels.append(labels[i])\n                new_masks.append(mask[i, :, :])\n            except ValueError:\n                print(\"Error in xmax xmin\")\n                pass\n\n        if len(new_labels) == 0:\n            boxes.append([0, 0, 20, 20])\n            new_labels.append(0)\n            new_masks.append(mask[0, :, :])\n\n        nmx = np.zeros((len(new_masks), self.width, self.height), dtype=np.uint8)\n        for i, n in enumerate(new_masks):\n            nmx[i, :, :] = n\n\n        boxes = torch.as_tensor(boxes, dtype=torch.float32)\n        labels = torch.as_tensor(new_labels, dtype=torch.int64)\n        masks = torch.as_tensor(nmx, dtype=torch.uint8)\n\n        image_id = torch.tensor([idx])\n        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)\n\n        target = {}\n        target[\"boxes\"] = boxes\n        target[\"labels\"] = labels\n        target[\"masks\"] = masks\n        target[\"image_id\"] = image_id\n        target[\"area\"] = area\n        target[\"iscrowd\"] = iscrowd\n\n        if self.transforms is not None:\n            img, target = self.transforms(img, target)\n        img = img.float()\n        return img, target\n\n    def __len__(self):\n        return len(self.image_info)","metadata":{"id":"40TtJ7Eu-QEr","execution":{"iopub.status.busy":"2021-11-10T15:33:16.400131Z","iopub.execute_input":"2021-11-10T15:33:16.400737Z","iopub.status.idle":"2021-11-10T15:33:16.461814Z","shell.execute_reply.started":"2021-11-10T15:33:16.400693Z","shell.execute_reply":"2021-11-10T15:33:16.461098Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We create the val dataloader with shuffle=False for consistent val runs. Also set num_workers to 0 to [prevent bus errors](https://stackoverflow.com/questions/51536114/pytorch-dataloader-killed).","metadata":{"id":"WlnsaOzz-QEs"}},{"cell_type":"code","source":"dataset = CellDataset(TRAIN_PATH, TRAIN_CSV, 704, 520, transforms=get_transform(train=True))\ntrain_size = int(len(dataset)*0.9)\nval_size = len(dataset)-train_size\ntrain_set, val_set = torch.utils.data.random_split(dataset, [train_size, val_size]) # We sample 10% of the images as a validation dataset\n\ndl_train = DataLoader(train_set, batch_size=config[\"batch_size\"], shuffle=True, \n                      num_workers=2, collate_fn=lambda x: tuple(zip(*x)))\ndl_val = DataLoader(val_set, batch_size=config[\"batch_size\"], shuffle=False, \n                    num_workers=2, collate_fn=lambda x: tuple(zip(*x)))","metadata":{"id":"VZkL9goy-QEs","execution":{"iopub.status.busy":"2021-11-10T15:33:16.464052Z","iopub.execute_input":"2021-11-10T15:33:16.465306Z","iopub.status.idle":"2021-11-10T15:33:16.887776Z","shell.execute_reply.started":"2021-11-10T15:33:16.465268Z","shell.execute_reply":"2021-11-10T15:33:16.887083Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"A function to combine all of the mask labels into one array with the sequence numbers in place of just the \"present\" flag.","metadata":{"id":"ft4VGCuZvw95"}},{"cell_type":"code","source":"from skimage.color import label2rgb\n\ndef combine_masks(masks):\n    \"\"\"Combine the masks labeled with their sequence number.\"\"\"\n\n    masks = masks.numpy()\n    all_masks = np.zeros((520, 704))\n    for i, mask in enumerate(masks, 1):\n        all_masks[mask == True] = i\n    return all_masks\n\ndef visualize_masks(img, masks, title, model=None):\n    \"\"\"Show the original image, then superimpose it with the provided masks.\n    \n    1   2\n    3   4\n    5   6\n    \n    1: Original Image\n    2: Ground Truth masks\n    3: Ground Truth masks\n    4: Predicted masks (if applicable)\n    5: Predicted masks\n    \"\"\"\n\n    fig, axs = plot.subplots(3, 2)\n    \n    fig.title(\"Mask Visualization\")\n\n    axs[0][0]\n","metadata":{"id":"uVr2tFB_oZHx","execution":{"iopub.status.busy":"2021-11-10T15:33:16.889268Z","iopub.execute_input":"2021-11-10T15:33:16.889549Z","iopub.status.idle":"2021-11-10T15:33:16.937251Z","shell.execute_reply.started":"2021-11-10T15:33:16.889516Z","shell.execute_reply":"2021-11-10T15:33:16.936518Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"enhance = ImageEqualize(config[\"clip_limit\"], config[\"difference\"])\n\ndef analyze(img, masks, model=None):\n    masks = combine_masks(masks)\n    fig, axs = plt.subplots(3, 2, figsize=(20, 20))\n    axs[0][0].imshow(img)\n\n    axs[0][1].imshow(label2rgb(masks, img, bg_label=0))\n\n    eimg = enhance.equalize(img)\n    print(eimg.shape)\n    axs[1][0].imshow(eimg)\n    axs[1][1].imshow(label2rgb(masks, eimg, bg_label=0))\n \n    if config[\"difference\"]:\n        dimg = img - eimg\n        dimg = dimg - dimg.min()\n        dimg = dimg / dimg.max()\n    else:\n        dimg = img\n    axs[2][0].imshow(dimg)\n    axs[2][1].imshow(label2rgb(masks, dimg, bg_label=0))\n\n    plt.show()\n    return \n    plt.title(\"Image\")\n    plt.imshow(img)\n    plt.show()\n\n    all_masks = combine_masks(masks)    \n    plt.imshow(all_masks, alpha=0.3)\n    plt.title(\"GT Masks\")\n    plt.show()\n\n    return\n    plt.title(\"Mask and Enhance Visualization\")\n    fig, axs = plt.subplots(2, 2, figsize=(30, 30))\n    axs[0][0].imshow(img)\n\n    masks = b['masks'].numpy().astype(np.int)\n    labels = b['labels']\n\n    print(labels[labels > 1].any())\n    all_masks = np.zeros_like(masks[0], np.int)\n    for i, m in zip(labels, masks):\n        all_masks[m != 0] = i\n\n    axs[1][1].imshow(all_masks)\n\n    if model is not None:\n        model.eval()\n        with torch.no_grad():\n            preds = model([img])[0]\n\n        all_preds_masks = np.zeros((520, 704))\n        for mask in preds['masks'].cpu().detach().numpy():\n            all_preds_masks = np.logical_or(all_masks, mask[0])\n        plt.imshow(all_preds_masks, alpha=0.8)\n        plt.title(\"Predictions\")\n        plt.show()\n\ndef analyze_sample(ds, sample_index, model=None):\n    img, targets = ds[sample_index]\n\n    img = img.numpy().transpose((1, 2, 0))\n    analyze(img, targets[\"masks\"], model)","metadata":{"id":"0k8NT5H_pndG","execution":{"iopub.status.busy":"2021-11-10T15:33:16.938536Z","iopub.execute_input":"2021-11-10T15:33:16.938948Z","iopub.status.idle":"2021-11-10T15:33:16.997180Z","shell.execute_reply.started":"2021-11-10T15:33:16.938911Z","shell.execute_reply":"2021-11-10T15:33:16.996293Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"analyze_sample(train_set, 8)\n#img, d = train_set[2]\n#img = img.numpy().transpose(1,2,0)\n#plt.imshow(img)\n#masks = combine_masks(d[\"masks\"])\n#plt.imshow(masks, alpha=0.3)\n#plt.show()","metadata":{"id":"TJ5JQJRpq4uf","outputId":"6c647051-508f-4f40-8f7e-16de99c30963","execution":{"iopub.status.busy":"2021-11-10T15:33:16.998711Z","iopub.execute_input":"2021-11-10T15:33:16.999095Z","iopub.status.idle":"2021-11-10T15:33:20.935951Z","shell.execute_reply.started":"2021-11-10T15:33:16.999057Z","shell.execute_reply":"2021-11-10T15:33:20.935092Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train loop","metadata":{"id":"XUixWQFP-QEs"}},{"cell_type":"markdown","source":"## Modeds                  \nFirst we set up model checkpointing\n\n","metadata":{"id":"wk0tITRw-QEs"}},{"cell_type":"code","source":"# Override pythorch checkpoint with an \"offline\" version of the file\nos.environ[\"HOME\"] = os.environ.get(\"HOME\", os.getcwd)\nHOME = Path(os.environ[\"HOME\"])\nchkpt_path = HOME / Path(\".cache/torch/hub/checkpoints\" % os.environ)\nchkpt_path.mkdir(parents=True, exist_ok=True)\n\nmodel_src_path = Path(INPUT_ROOT) / Path(\"cocopre/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth\")\nmodel_tgt_path = chkpt_path / Path(\"maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth\")\nprint(model_src_path.exists())\nwith model_src_path.open(\"rb\") as s:\n    with model_tgt_path.open(\"wb+\") as t:\n        t.write(s.read())","metadata":{"id":"rVLW5auc-QEt","outputId":"1f45b370-9c28-4825-feb6-e100619b485b","execution":{"iopub.status.busy":"2021-11-10T15:33:20.939691Z","iopub.execute_input":"2021-11-10T15:33:20.940187Z","iopub.status.idle":"2021-11-10T15:33:21.465007Z","shell.execute_reply.started":"2021-11-10T15:33:20.940137Z","shell.execute_reply":"2021-11-10T15:33:21.464144Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"In many places we have to clear any prior GPU memory allocations by first telling torch to dump it's cuda cache, then we trigger a python GC run to sweep the dangling references and trigger the objects to release the GPU memory in their destruction.\n\nWe define a decorator to handle this and wrap the relevant methods.","metadata":{"id":"F2IFz7jrod-a"}},{"cell_type":"code","source":"from functools import wraps\nimport gc\n\ndef flush_and_gc(f):\n  @wraps(f)\n  def g(*args, **kwargs):\n    torch.cuda.empty_cache()\n    gc.collect()\n    return f(*args, **kwargs)\n  return g","metadata":{"id":"2rSRtUFxchje","execution":{"iopub.status.busy":"2021-11-10T15:33:21.466475Z","iopub.execute_input":"2021-11-10T15:33:21.466875Z","iopub.status.idle":"2021-11-10T15:33:21.511522Z","shell.execute_reply.started":"2021-11-10T15:33:21.466839Z","shell.execute_reply":"2021-11-10T15:33:21.510819Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### IOU Accuracy Metric\n\nThis is mainly based on the [competition accuracy metric](https://www.kaggle.com/c/sartorius-cell-instance-segmentation/overview/evaluation).\n\n$ \\frac{1}{|\\bf{t}|} \\sum_\\bf{t} \\frac{\\bf{TP}(t)}{\\bf{TP}(t)\\bf{FP}(t)\\bf{FN}(t)}$\n\nwhere:\n\n$ \\bf{t} = \\{.5, .55, .6, .65, .7, .75, .8, .85, .9, .95\\}$","metadata":{"id":"mqrQhzsuADWD"}},{"cell_type":"code","source":"iou_thresholds = [.5, .55, .6, .65, .7, .75, .8, .85, .9, .95]\niou_thresholds_mean = sum(iou_thresholds) / len(iou_thresholds)\n\ndef sartorius_iou(src_boxes, pred_boxes):\n      \"\"\"\n      The accuracy method is not the one used in the evaluator but very similar\n      \"\"\"\n\n      total_gt = len(src_boxes)\n      total_pred = len(pred_boxes)\n\n      thrshs = torch.tensor(iou_thresholds)\n      thrshs_mean = torch.mean(thrshs)\n\n      def iou(threshold):\n          # Define the matcher and distance matrix based on iou\n          matcher = Matcher(threshold,threshold,allow_low_quality_matches=False) \n          match_quality_matrix = box_iou(src_boxes,pred_boxes)\n\n          results = matcher(match_quality_matrix)\n\n          true_positive = torch.count_nonzero(results.unique() != -1)\n          matched_elements = results[results > -1]\n\n          #in Matcher, a pred element can be matched only twice \n          false_positive = torch.count_nonzero(results == -1) + ( len(matched_elements) - len(torch.unique(matched_elements)))\n          false_negative = total_gt - true_positive\n\n          acc = true_positive / ( true_positive + false_positive + false_negative )\n\n          return acc\n\n      if total_gt > 0 and total_pred > 0:\n        return torch.tensor(sum([iou(t) for t in iou_thresholds]) / iou_thresholds_mean)\n  \n      elif total_gt == 0:\n          if total_pred > 0:\n              return torch.tensor(0.)\n          else:\n              return torch.tensor(1.)\n      elif total_gt > 0 and total_pred == 0:\n            return torch.tensor(0.)","metadata":{"id":"Nmtv7IfIAA-n","execution":{"iopub.status.busy":"2021-11-10T15:33:21.514906Z","iopub.execute_input":"2021-11-10T15:33:21.515155Z","iopub.status.idle":"2021-11-10T15:33:22.153970Z","shell.execute_reply.started":"2021-11-10T15:33:21.515127Z","shell.execute_reply":"2021-11-10T15:33:22.153101Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Next we create the model. It's base is a pretrained [Mask RCNN](https://pytorch.org/vision/stable/models.html#id63). We replace the features with ","metadata":{"id":"K7Lz-EwC-QEt"}},{"cell_type":"code","source":"import pytorch_lightning as pl\n\n\nclass Model(pl.LightningModule):\n    def __init__(self, num_classes=2, hidden_layer=256):\n        super().__init__()\n\n        # We don't want any of the pretrained model layers trainable at first.\n        self.detector = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True, trainable_backbone_layers=0)\n\n        # get the number of input features for the classifier\n        in_features = self.detector.roi_heads.box_predictor.cls_score.in_features\n        # replace the pre-trained head with a new one\n        self.detector.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 = self.detector.roi_heads.mask_predictor.conv5_mask.in_channels\n\n        # and replace the mask predictor with a new one\n        self.detector.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)\n\n    def full_train(self):\n        self.detector.requires_grad = True\n    \n    def configure_optimizers(self):\n        optimizer = torch.optim.SGD(self.parameters(), lr=config[\"lr\"], momentum=config[\"momentum\"], weight_decay=config[\"weight_decay\"], nesterov=config[\"nesterov\"])\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, 3, 1, config[\"min_lr\"], verbose=True)\n        return {\"optimizer\": optimizer,\n                \"lr_scheduler\": {\n                    \"scheduler\": scheduler,\n                    \"monitor\": \"val_accuracy\"\n                }}\n\n    def forward(self, x):\n        self.detector.eval()\n        return self.detector(x)\n\n    @flush_and_gc\n    def training_step(self, batch, batch_idx):\n        images, targets = batch\n        \n        loss_dict = self.detector(images, targets)\n    \n        preds = self.forward(images)\n        self.detector.train()\n        \n        loss_dict[\"loss_sartorius\"] = 1 - torch.mean(torch.stack([sartorius_iou(b[\"boxes\"],pb[\"boxes\"]) for b,pb in zip(targets,preds)]))\n        #loss_dict = {k:v for k, v in loss_dict.items()  if k in [\"loss_sartorius\", \"loss_classifier\", \"loss_mask\"]}\n        loss = sum(loss_dict.values())\n        \n        loss_dict = {k:(v.detach() if hasattr(v, \"detach\") else v) for k, v in loss_dict.items()}        \n        self.log(\"loss\", loss)\n        self.log_dict(loss_dict)\n\n        return {\"loss\": loss, \"log\": loss_dict}\n    \n    @flush_and_gc\n    def validation_step(self, batch, batch_idx):\n        img, boxes = batch\n        pred_boxes = self.forward(img)\n\n        self.val_accuracy = torch.mean(torch.stack([sartorius_iou(b[\"boxes\"],pb[\"boxes\"]) for b,pb in zip(boxes,pred_boxes)]))\n    \n        self.log(\"val_accuracy\", self.val_accuracy)\n        return self.val_accuracy\n\n    @flush_and_gc\n    def test_step(self, batch, batch_idx):\n        img, boxes, metadata = batch\n        pred_boxes = self.forward(img) # in validation, faster rcnn return the boxes\n        self.test_accuracy = torch.mean(torch.stack([sartorius_iou(b,pb[\"boxes\"]) for b,pb in zip(boxes,pred_boxes)]))\n        r = {\"accuracy_test\": self.test_accuracy}\n        self.log_dict(r)\n        return r","metadata":{"id":"ifR3hc8L-QEt","execution":{"iopub.status.busy":"2021-11-10T15:33:22.155952Z","iopub.execute_input":"2021-11-10T15:33:22.156654Z","iopub.status.idle":"2021-11-10T15:33:22.216070Z","shell.execute_reply.started":"2021-11-10T15:33:22.156611Z","shell.execute_reply":"2021-11-10T15:33:22.215317Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model Training\n\nSince we train the model twice, we have a function to help us with that. It will configure [Weights and Biases](wandb.ai). It also creates the trainer, then runs it with checkpointing, and at the end will load the best checkpoint and return it.","metadata":{"id":"zjAoSgaKcDKk"}},{"cell_type":"code","source":"from pytorch_lightning import Trainer\nfrom pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint\n\ndef train_model(model, run_name, run=None):\n  wandb_logger = None\n  if run:\n    run.config[\"train_run_name\"] = run_name\n\n    from pytorch_lightning.loggers import WandbLogger\n    wandb_logger = WandbLogger()\n    \n  chkpt = ModelCheckpoint(f\"/kaggle/working/chkpt-{run_name}\", monitor=\"val_accuracy\", mode=\"max\")\n\n  trainer = pl.Trainer(gpus=1, logger=wandb_logger, max_epochs=config[\"epochs\"], callbacks=[chkpt])\n  trainer.fit(model, dl_train, dl_val)\n  \n  return Model.load_from_checkpoint(chkpt.best_model_path)\n  ","metadata":{"id":"eideo5bttbuR","execution":{"iopub.status.busy":"2021-11-10T15:33:22.219342Z","iopub.execute_input":"2021-11-10T15:33:22.219990Z","iopub.status.idle":"2021-11-10T15:33:22.267657Z","shell.execute_reply.started":"2021-11-10T15:33:22.219948Z","shell.execute_reply":"2021-11-10T15:33:22.266903Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training loop!\n\nFirst we train just the head, and then we train the full model. This holds most of the parameters constant while getting the head to converge, and then fine tunes the entire model for the new head.","metadata":{"id":"AM3_m8Om-QEu"}},{"cell_type":"code","source":"with wandb_context(config) as run:\n  model = Model()\n  if run:\n    run.watch(model)\n  model = train_model(model, \"head\", run)\n\n  # Set the model for fine-tuning training\n  model.full_train()\n  model = train_model(model, \"full\", run)","metadata":{"id":"8oRasTvg-QEu","outputId":"9e5baea2-729a-43aa-bd88-fda3e321ba00","execution":{"iopub.status.busy":"2021-11-10T15:33:22.270720Z","iopub.execute_input":"2021-11-10T15:33:22.272210Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Test Dataset and DataLoader","metadata":{"id":"ssOSvoxT-QEv"}},{"cell_type":"code","source":"class CellTestDataset(Dataset):\n    def __init__(self, image_dir, height, width, transforms=None):\n        self.transforms = transforms\n        \n        self.image_dir = image_dir\n        \n        self.image_ids = [f[:-4]for f in os.listdir(self.image_dir)]\n        self.num_samples = len(self.image_ids)\n        \n        self.height = height\n        self.width = width\n        self.image_info = collections.defaultdict(dict)\n            \n    def __getitem__(self, idx):\n        image_id = self.image_ids[idx]\n        image_path = os.path.join(self.image_dir, image_id + '.png')\n        image = Image.open(image_path).convert(\"RGB\")\n        #image = image.resize((self.width, self.height), resample=Image.BILINEAR)\n\n        if self.transforms is not None:\n            image, _ = self.transforms(image=image, target=None)\n        return {'image': image, 'image_id': image_id}\n\n    def __len__(self):\n        return len(self.image_ids)","metadata":{"id":"ZkTCZCsn-QEv","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ds_test = CellTestDataset(TEST_PATH, 704, 520, transforms=get_transform(train=False))\ndl_test = DataLoader(ds_test, batch_size=config[\"batch_size\"], shuffle=True, \n                      num_workers=0, collate_fn=lambda x: tuple(zip(*x)))","metadata":{"id":"Dat-POum-QEv","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Analyze prediction results for train set","metadata":{"id":"hEBaKkJO-QEw"}},{"cell_type":"code","source":"# NOTE: It puts the model in eval mode!! Revert for re-training\nanalyze_sample(train_set, 20, model)","metadata":{"id":"qwTx7kvk-QEw","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"analyze_sample(train_set, 100, model)","metadata":{"id":"u5H5DXkL-QEw","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"analyze_sample(train_set, 2, model)","metadata":{"id":"N33uzjsL-QEw","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Prediction","metadata":{"id":"pOd8V71r-QEw"}},{"cell_type":"markdown","source":"## Utilities","metadata":{"id":"ByLFcKAI-QEw"}},{"cell_type":"code","source":"# Stolen from: https://www.kaggle.com/arunamenon/cell-instance-segmentation-unet-eda\n# Run-length encoding stolen from https://www.kaggle.com/rakhlin/fast-run-length-encoding-python\n# Modified by me\ndef rle_encoding(x):\n    dots = np.where(x.flatten() == 1)[0]\n    run_lengths = []\n    prev = -2\n    for b in dots:\n        if (b>prev+1): run_lengths.extend((b + 1, 0))\n        run_lengths[-1] += 1\n        prev = b\n    return ' '.join(map(str, run_lengths))\n\n\ndef does_overlap(mask, other_masks):\n    for other_mask in other_masks:\n        if np.sum(np.logical_and(mask, other_mask)) > 0:\n            #import pdb; pdb.set_trace()\n            #print(\"Found overlapping masks!\")\n            return True\n    return False\n\n\ndef remove_overlapping_pixels(mask, other_masks):\n    for other_mask in other_masks:\n        if np.sum(np.logical_and(mask, other_mask)) > 0:\n            mask[np.logical_and(mask, other_mask)] = 0\n    return mask","metadata":{"id":"q2Q77NYF-QEw","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Submission","metadata":{"id":"x52dir37-QEx"}},{"cell_type":"code","source":"sublist = []\ncounter = 0\n\nwidth = 704\nheight = 520\nTHRESHOLD = 0.5\n\nmodel.eval()\n\nfor sample in ds_test:\n    img = sample['image'].float()\n    image_id = sample['image_id']\n    with torch.no_grad():\n        result = model([img])[0]\n    if len(result[\"masks\"]) > 0:\n        previous_masks = []\n        for j, m in enumerate(result[\"masks\"]):\n            original_mask = result[\"masks\"][j][0].cpu().numpy()\n            original_mask = remove_overlapping_pixels(original_mask, previous_masks)\n            previous_masks.append(original_mask)\n            rle = rle_encoding(original_mask > THRESHOLD)\n            sublist.append([image_id, rle])\n    else:\n        sublist.append([image_id, \"\"])\n\ndf_sub = pd.DataFrame(sublist, columns=['id', 'predicted'])\ndf_sub.to_csv(\"submission.csv\", index=False)\ndf_sub.head()","metadata":{"id":"g1_co7SL-QEx","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"id":"ZeliInho-QEx"},"execution_count":null,"outputs":[]}]}