{"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":"# The notebooks is self-contained\n# It has very few imports\n# No external dependencies (only the model weights)\n# No train - inference notebooks\n# We only rely on Pytorch\nimport os\nimport time\nimport random\nimport collections\n\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nfrom sklearn.model_selection import train_test_split\n\nimport torch\nimport torchvision\nfrom torchvision.transforms import ToPILImage\nfrom torchvision.transforms import functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection.mask_rcnn import MaskRCNNPredictor","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-11-05T13:00:33.028767Z","iopub.execute_input":"2021-11-05T13:00:33.029128Z","iopub.status.idle":"2021-11-05T13:00:33.036455Z","shell.execute_reply.started":"2021-11-05T13:00:33.029087Z","shell.execute_reply":"2021-11-05T13:00:33.035590Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Fix randomness\n\ndef fix_all_seeds(seed):\n    np.random.seed(seed)\n    torch.backends.cudnn.benchmark = False\n    torch.backends.cudnn.deterministic = True\n    torch.manual_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\n    \nfix_all_seeds(2021)","metadata":{"execution":{"iopub.status.busy":"2021-11-05T13:00:33.038545Z","iopub.execute_input":"2021-11-05T13:00:33.038809Z","iopub.status.idle":"2021-11-05T13:00:33.050915Z","shell.execute_reply.started":"2021-11-05T13:00:33.038775Z","shell.execute_reply":"2021-11-05T13:00:33.050145Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"TRAIN_CSV = \"../input/sartorius-cell-instance-segmentation/train.csv\"\nTRAIN_PATH = \"../input/sartorius-cell-instance-segmentation/train\"\nTEST_PATH = \"../input/sartorius-cell-instance-segmentation/test\"\nWIDTH = 704\nHEIGHT = 520\n\n# Reduced the train dataset to 5000 rows\nTEST = False\n\nDEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\nRESNET_MEAN = (0.485, 0.456, 0.406)\nRESNET_STD = (0.229, 0.224, 0.225)\n\nBATCH_SIZE = 2\n\n# No changes tried with the optimizer yet.\nMOMENTUM = 0.9\nLEARNING_RATE = 0.0008\nWEIGHT_DECAY = 0.0005\n\n# Changes the confidence required for a pixel to be kept for a mask. \n# Only used 0.5 till now.\nMASK_THRESHOLD = 0.5\n\n# Normalize to resnet mean and std if True.\nNORMALIZE = False\n\n\n# Use a StepLR scheduler if True. Not tried yet.\nUSE_SCHEDULER = False\n\n# Amount of epochs\nNUM_EPOCHS = 13\n\n\nBOX_DETECTIONS_PER_IMG = 539\n\n\nMIN_SCORE = 0.59","metadata":{"execution":{"iopub.status.busy":"2021-11-05T13:00:33.052977Z","iopub.execute_input":"2021-11-05T13:00:33.053279Z","iopub.status.idle":"2021-11-05T13:00:33.103345Z","shell.execute_reply.started":"2021-11-05T13:00:33.053243Z","shell.execute_reply":"2021-11-05T13:00:33.102131Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# These are slight redefinitions of torch.transformation classes\n# The difference is that they handle the target and the mask\n# Copied from Abishek, added new ones\nclass Compose:\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 VerticalFlip:\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(-2)\n            bbox = target[\"boxes\"]\n            bbox[:, [1, 3]] = height - bbox[:, [3, 1]]\n            target[\"boxes\"] = bbox\n            target[\"masks\"] = target[\"masks\"].flip(-2)\n        return image, target\n\nclass HorizontalFlip:\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            target[\"masks\"] = target[\"masks\"].flip(-1)\n        return image, target\n\nclass Normalize:\n    def __call__(self, image, target):\n        image = F.normalize(image, RESNET_MEAN, RESNET_STD)\n        return image, target\n\nclass ToTensor:\n    def __call__(self, image, target):\n        image = F.to_tensor(image)\n        return image, target\n    \n\ndef get_transform(train):\n    transforms = [ToTensor()]\n    if NORMALIZE:\n        transforms.append(Normalize())\n    \n    # Data augmentation for train\n    if train: \n        transforms.append(HorizontalFlip(0.5))\n        transforms.append(VerticalFlip(0.5))\n\n    return Compose(transforms)","metadata":{"execution":{"iopub.status.busy":"2021-11-05T13:00:33.105827Z","iopub.execute_input":"2021-11-05T13:00:33.106344Z","iopub.status.idle":"2021-11-05T13:00:33.122691Z","shell.execute_reply.started":"2021-11-05T13:00:33.106306Z","shell.execute_reply":"2021-11-05T13:00:33.121749Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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    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)","metadata":{"execution":{"iopub.status.busy":"2021-11-05T13:00:33.124373Z","iopub.execute_input":"2021-11-05T13:00:33.124745Z","iopub.status.idle":"2021-11-05T13:00:33.136077Z","shell.execute_reply.started":"2021-11-05T13:00:33.124707Z","shell.execute_reply":"2021-11-05T13:00:33.135238Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CellDataset(Dataset):\n    def __init__(self, image_dir, df, transforms=None, resize=False):\n        self.transforms = transforms\n        self.image_dir = image_dir\n        self.df = df\n        \n        self.should_resize = resize is not False\n        if self.should_resize:\n            self.height = int(HEIGHT * resize)\n            self.width = int(WIDTH * resize)\n        else:\n            self.height = HEIGHT\n            self.width = WIDTH\n        \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 get_box(self, a_mask):\n        ''' Get the bounding box of a given mask '''\n        pos = np.where(a_mask)\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        return [xmin, ymin, xmax, ymax]\n\n    def __getitem__(self, idx):\n        ''' Get the image and the target'''\n        \n        img_path = self.image_info[idx][\"image_path\"]\n        img = Image.open(img_path).convert(\"RGB\")\n        \n        if self.should_resize:\n            img = img.resize((self.width, self.height), resample=Image.BILINEAR)\n\n        info = self.image_info[idx]\n\n        n_objects = len(info['annotations'])\n        masks = np.zeros((len(info['annotations']), self.height, self.width), dtype=np.uint8)\n        boxes = []\n        \n        for i, annotation in enumerate(info['annotations']):\n            a_mask = rle_decode(annotation, (HEIGHT, WIDTH))\n            a_mask = Image.fromarray(a_mask)\n            \n            if self.should_resize:\n                a_mask = a_mask.resize((self.width, self.height), resample=Image.BILINEAR)\n            \n            a_mask = np.array(a_mask) > 0\n            masks[i, :, :] = a_mask\n            \n            boxes.append(self.get_box(a_mask))\n\n        # dummy labels\n        labels = [1 for _ in range(n_objects)]\n        \n        boxes = torch.as_tensor(boxes, dtype=torch.float32)\n        labels = torch.as_tensor(labels, dtype=torch.int64)\n        masks = torch.as_tensor(masks, dtype=torch.uint8)\n\n        image_id = torch.tensor([idx])\n        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n        iscrowd = torch.zeros((n_objects,), dtype=torch.int64)\n\n        # This is the required target for the Mask R-CNN\n        target = {\n            'boxes': boxes,\n            'labels': labels,\n            'masks': masks,\n            'image_id': image_id,\n            'area': area,\n            'iscrowd': iscrowd\n        }\n\n        if self.transforms is not None:\n            img, target = self.transforms(img, target)\n\n        return img, target\n\n    def __len__(self):\n        return len(self.image_info)","metadata":{"execution":{"iopub.status.busy":"2021-11-05T13:00:33.185203Z","iopub.execute_input":"2021-11-05T13:00:33.185725Z","iopub.status.idle":"2021-11-05T13:00:33.207883Z","shell.execute_reply.started":"2021-11-05T13:00:33.185684Z","shell.execute_reply":"2021-11-05T13:00:33.207129Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# df_train = pd.read_csv(TRAIN_CSV, nrows=5000 if TEST else None)\n# ds_train = CellDataset(TRAIN_PATH, df_train, resize=False, transforms=get_transform(train=True))\n# dl_train = DataLoader(ds_train, batch_size=BATCH_SIZE, shuffle=True, \n#                       num_workers=2, collate_fn=lambda x: tuple(zip(*x)))","metadata":{"execution":{"iopub.status.busy":"2021-11-05T13:00:33.211848Z","iopub.execute_input":"2021-11-05T13:00:33.212084Z","iopub.status.idle":"2021-11-05T13:00:33.220658Z","shell.execute_reply.started":"2021-11-05T13:00:33.212060Z","shell.execute_reply":"2021-11-05T13:00:33.219702Z"},"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    if NORMALIZE:\n        model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True, \n                                                                   box_detections_per_img=BOX_DETECTIONS_PER_IMG,\n                                                                   image_mean=RESNET_MEAN, \n                                                                   image_std=RESNET_STD)\n    else:\n        model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True,\n                                                                  box_detections_per_img=BOX_DETECTIONS_PER_IMG)\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-11-05T13:00:33.224040Z","iopub.execute_input":"2021-11-05T13:00:33.224221Z","iopub.status.idle":"2021-11-05T13:00:43.628949Z","shell.execute_reply.started":"2021-11-05T13:00:33.224199Z","shell.execute_reply":"2021-11-05T13:00:43.628014Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from sklearn.model_selection import train_test_split\ndf_all = pd.read_csv(TRAIN_CSV, nrows=5000 if TEST else None)\ndf_train, df_valid = train_test_split(df_all['id'].unique(), test_size=0.2, random_state=42)\ndf_valid, df_test = train_test_split(df_valid, test_size=0.5, random_state=42)\n\nds_train = CellDataset(TRAIN_PATH, df_all[df_all['id'].isin(df_train)], resize=False, \n                       transforms=get_transform(train=True))\nds_valid = CellDataset(TRAIN_PATH, df_all[df_all['id'].isin(df_valid)], resize=False, \n                       transforms=get_transform(train=False))\nds_test = CellDataset(TRAIN_PATH, df_all[df_all['id'].isin(df_test)], resize=False,\n                     transforms=get_transform(train=False))\n                       \ndl_train = DataLoader(ds_train, batch_size=BATCH_SIZE, shuffle=True, \n                      num_workers=2, collate_fn=lambda x: tuple(zip(*x)))\ndl_valid = DataLoader(ds_valid, batch_size=BATCH_SIZE, shuffle=True, \n                      num_workers=2, collate_fn=lambda x: tuple(zip(*x))) \ndl_test = DataLoader(ds_test, batch_size=BATCH_SIZE, shuffle=True, \n                      num_workers=2, collate_fn=lambda x: tuple(zip(*x))) \n\nparams = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.SGD(params, lr=LEARNING_RATE, momentum=MOMENTUM, weight_decay=WEIGHT_DECAY)\n\nmin_lr = 0.0001\nscheduler_step = NUM_EPOCHS // 5\nlr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, scheduler_step, eta_min=min_lr)\n# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)\n\nn_batches = len(dl_train)\n\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    model.train()\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        # Backprop\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\n    if USE_SCHEDULER:\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    elapsed = time.time() - time_start\n\n    # Valid\n#         model.eval()\n    loss_accum_v = 0.0\n    loss_mask_accum_v = 0.0\n    for batch_idx,(images_v, targets_v) in enumerate(dl_valid, 1):\n        images_v = list(image.to(DEVICE) for image in images_v)\n        targets_v = [{k: v.to(DEVICE) for k, v in t.items()} for t in targets_v]\n        with torch.no_grad():\n            loss_dict_v = model(images_v, targets_v)\n            loss_v = sum(loss for loss in loss_dict_v.values())\n            # Logging\n            loss_mask_v = loss_dict_v['loss_mask'].item()\n            loss_accum_v += loss_v.item()\n            loss_mask_accum_v += loss_mask_v\n    valid_loss = loss_accum_v / len(dl_valid)\n    valid_loss_mask = loss_mask_accum_v / len(dl_valid)\n    \n        # test\n#         model.eval()\n    loss_accum_t = 0.0\n    loss_mask_accum_t = 0.0\n    for batch_idx,(images_t, targets_t) in enumerate(dl_test, 1):\n        images_t = list(image.to(DEVICE) for image in images_t)\n        targets_t = [{k: v.to(DEVICE) for k, v in t.items()} for t in targets_t]\n        with torch.no_grad():\n            loss_dict_t = model(images_t, targets_t)\n            loss_t = sum(loss for loss in loss_dict_t.values())\n            # Logging\n            loss_mask_t = loss_dict_t['loss_mask'].item()\n            loss_accum_t += loss_t.item()\n            loss_mask_accum_t += loss_mask_t\n    test_loss = loss_accum_t / len(dl_test)\n    test_loss_mask = loss_mask_accum_t / len(dl_test)\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]\")\n\n    print(f\"{prefix} Valid mask-only loss: {valid_loss_mask:7.3f}\")\n    print(f\"{prefix} Valid loss: {valid_loss:7.3f}\")     \n\n    print(f\"{prefix} Test mask-only loss: {test_loss_mask:7.3f}\")\n    print(f\"{prefix} Test loss: {test_loss:7.3f}\")  ","metadata":{"execution":{"iopub.status.busy":"2021-11-05T13:00:43.630341Z","iopub.execute_input":"2021-11-05T13:00:43.630615Z","iopub.status.idle":"2021-11-05T13:59:21.990802Z","shell.execute_reply.started":"2021-11-05T13:00:43.630582Z","shell.execute_reply":"2021-11-05T13:59:21.989904Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}