{"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":"import numpy as np\nimport pandas as pd\nimport random\nimport os\nimport time\nimport collections\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nfrom sklearn.model_selection import StratifiedKFold, train_test_split\nfrom tqdm.auto import tqdm\n\nimport torch\nimport torchvision\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-20T14:45:44.906448Z","iopub.execute_input":"2021-11-20T14:45:44.90679Z","iopub.status.idle":"2021-11-20T14:45:44.914349Z","shell.execute_reply.started":"2021-11-20T14:45:44.906747Z","shell.execute_reply":"2021-11-20T14:45:44.913391Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install pycocotools\nimport pycocotools.mask as mask_util","metadata":{"execution":{"iopub.status.busy":"2021-11-20T14:45:44.916387Z","iopub.execute_input":"2021-11-20T14:45:44.916722Z","iopub.status.idle":"2021-11-20T14:45:52.071773Z","shell.execute_reply.started":"2021-11-20T14:45:44.916684Z","shell.execute_reply":"2021-11-20T14:45:52.070959Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"VAL_SPLIT = 0.1\nBATCH_SIZE = 4\nNUM_EPOCHS = 20\nBOX_DETECTIONS_PER_IMG = 540\nMIN_SCORE = 0.5\nTEST_FLAG = False\n\nMOMENTUM = 0.9\nLEARNING_RATE = 0.001\n#WEIGHT_DECAY = 0.0005\nWEIGHT_DECAY = 0.0\nUSE_SCHEDULER = True\n\nTRAIN_CSV = '../input/sartorius-cell-instance-segmentation/train.csv'\nTRAIN_PATH = '../input/sartorius-cell-instance-segmentation/train'\nTEST_PATH = '../input/sartorius-cell-instance-segmentation/test'\n\nORIG_WIDTH = 704\nORIG_HEIGHT = 520\nWIDTH = 520\nHEIGHT = 520\nMIN_BOX_SIDE = 5\n\nMASK_THRESHOLD = 0.5\nSEED = 2021\n\nSCORE_THRESHOLDS = {\n    'cort': 0.75,\n    'shsy5y': 0.50,\n    'astro': 0.55\n}\n\nMASK_THRESHOLDS = {\n    'cort': 0.75,\n    'shsy5y': 0.60,\n    'astro': 0.55\n}\n\nSCORE_THRESHOLDS = {\n    'cort': 0.50,\n    'shsy5y': 0.50,\n    'astro': 0.50\n}\n\nMASK_THRESHOLDS = {\n    'cort': 0.50,\n    'shsy5y': 0.50,\n    'astro': 0.50\n}\n\nMIN_PIXELS = 75","metadata":{"execution":{"iopub.status.busy":"2021-11-20T14:45:52.074183Z","iopub.execute_input":"2021-11-20T14:45:52.074451Z","iopub.status.idle":"2021-11-20T14:45:52.080731Z","shell.execute_reply.started":"2021-11-20T14:45:52.074415Z","shell.execute_reply":"2021-11-20T14:45:52.080044Z"},"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-11-20T14:45:52.083587Z","iopub.execute_input":"2021-11-20T14:45:52.083951Z","iopub.status.idle":"2021-11-20T14:45:54.032351Z","shell.execute_reply.started":"2021-11-20T14:45:52.083916Z","shell.execute_reply":"2021-11-20T14:45:54.031321Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def 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(SEED)","metadata":{"execution":{"iopub.status.busy":"2021-11-20T14:45:54.035104Z","iopub.execute_input":"2021-11-20T14:45:54.035398Z","iopub.status.idle":"2021-11-20T14:45:54.045119Z","shell.execute_reply.started":"2021-11-20T14:45:54.035357Z","shell.execute_reply":"2021-11-20T14:45:54.044382Z"},"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    \nclass ContrastFromMiddle:\n    def __call__(self, img, target):\n        m2 = (img > 0.5) * (img - 0.5)\n        m1 = (img <= 0.5) * (0.5 - img)\n        conv = (img > 0.5) * (np.sin(m2 * np.pi) / 2 + 0.5) \n        conv += (img <= 0.5) * (0.5 - np.sin(m1 * np.pi) / 2)\n        image = conv\n        return image, target\n\nclass Rotate:\n    def __init__(self, prob):\n        self.prob = prob\n\n    def __call__(self, image, target):\n        if random.random() < self.prob:\n            k = random.randint(1, 3)\n            height, width = image.shape[-2:]\n            image = image.rot90(k, [1, 2])\n            bbox = target[\"boxes\"]\n            x = bbox[:, [0, 2]]\n            y = bbox[:, [1, 3]]\n            if k == 1:\n                bbox[:, [0, 2]] = y\n                bbox[:, [3, 1]] = width - x\n            elif k == 2:\n                bbox[:, [2, 0]] = width - x\n                bbox[:, [3, 1]] = height - y\n            else:\n                bbox[:, [2, 0]] = height - y\n                bbox[:, [1, 3]] = x\n            target[\"boxes\"] = bbox\n            target[\"masks\"] = target[\"masks\"].rot90(k, [1, 2])\n        return image, target\n\nclass ColorChange:\n    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):\n        self.brightness = brightness\n        self.contrast = contrast\n        self.saturation = saturation\n        self.hue = hue\n        \n    def __call__(self, image, target):\n        if self.brightness > 0:\n            image = F.adjust_brightness(image, random.uniform(1 - self.brightness, 1 + self.brightness))\n        if self.contrast > 0:\n            image = F.adjust_contrast(image, random.uniform(1 - self.contrast, 1 + self.contrast))\n        if self.saturation > 0:\n            image = F.adjust_saturation(image, random.uniform(1 - self.saturation, 1 + self.saturation))\n        if self.hue > 0:\n            image = F.adjust_hue(image, random.uniform(-self.hue, self.hue))\n        return image, target\n    \ndef get_transform(train):\n    transforms = [ToTensor()]\n    # Data augmentation for train\n    if train: \n        transforms.append(HorizontalFlip(0.50))\n        transforms.append(VerticalFlip(0.50))\n        transforms.append(Rotate(0.75))\n        transforms.append(ColorChange(0.2, 0.2, 0.2, 0.1))\n\n    return Compose(transforms)","metadata":{"execution":{"iopub.status.busy":"2021-11-20T14:45:54.04647Z","iopub.execute_input":"2021-11-20T14:45:54.048185Z","iopub.status.idle":"2021-11-20T14:45:54.539504Z","shell.execute_reply.started":"2021-11-20T14:45:54.048148Z","shell.execute_reply":"2021-11-20T14:45:54.538627Z"},"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)\n\ndef rle_encode(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\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":{"execution":{"iopub.status.busy":"2021-11-20T14:45:54.542761Z","iopub.execute_input":"2021-11-20T14:45:54.544031Z","iopub.status.idle":"2021-11-20T14:45:54.555848Z","shell.execute_reply.started":"2021-11-20T14:45:54.543984Z","shell.execute_reply":"2021-11-20T14:45:54.555046Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def precision_at(threshold, iou):\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 image_metric(pred, targ):\n    enc_preds = [mask_util.encode(np.asarray(p, order='F')) for p in pred]\n    enc_targs = [mask_util.encode(np.asarray(p, order='F')) for p in targ]\n    ious = mask_util.iou(enc_preds, enc_targs, [0]*len(enc_targs))\n    prec = []\n    for t in np.arange(0.5, 1.0, 0.05):\n        tp, fp, fn = precision_at(t, ious)\n        p = tp / (tp + fp + fn)\n        prec.append(p)\n    return np.mean(prec)","metadata":{"execution":{"iopub.status.busy":"2021-11-20T14:45:54.557474Z","iopub.execute_input":"2021-11-20T14:45:54.557961Z","iopub.status.idle":"2021-11-20T14:45:54.570286Z","shell.execute_reply.started":"2021-11-20T14:45:54.557924Z","shell.execute_reply":"2021-11-20T14:45:54.569412Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CellDataset(Dataset):\n    def __init__(self, image_dir, df, trainFlag):\n        self.transforms = get_transform(trainFlag)\n        self.image_dir = image_dir\n        self.df = df\n        self.height = HEIGHT\n        self.width = WIDTH\n        self.trainFlag = trainFlag\n        \n        cell_type_list = {}\n        self.image_info = {}\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                    'cell_type': df[df['id'] == row['id']].iloc[0]['cell_type']\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        img_path = self.image_info[idx][\"image_path\"]\n        img = Image.open(img_path).convert(\"RGB\")\n        info = self.image_info[idx]\n        \n        if self.trainFlag:\n            shift = random.randint(0, ORIG_WIDTH - WIDTH)\n            n_img = np.array(img)\n            n_img = n_img[:, shift:shift + WIDTH, :]\n            img = Image.fromarray(n_img)\n\n        if self.trainFlag:\n            masks = np.zeros((0, HEIGHT, WIDTH), dtype=np.uint8)\n        else:\n            masks = np.zeros((0, ORIG_HEIGHT, ORIG_WIDTH), dtype=np.uint8)\n            \n        boxes = []\n        n_objects = 0\n        for i, annotation in enumerate(info['annotations']):\n            a_mask = rle_decode(annotation, (ORIG_HEIGHT, ORIG_WIDTH))\n            a_mask = a_mask > 0\n            if self.trainFlag:\n                a_mask = a_mask[:, shift:shift + WIDTH]\n            if a_mask.sum() > 0:\n                # have non-empty mask after crop\n                box = self.get_box(a_mask)\n                if (box[2] - box[0] >= MIN_BOX_SIDE) and (box[3] - box[1] >= MIN_BOX_SIDE):\n                    # ignore very small boxes\n                    boxes.append(box)\n                    a_mask = np.expand_dims(a_mask, axis=0)\n                    masks = np.concatenate((masks, a_mask), axis=0)\n                    n_objects += 1\n\n        # dummy labels\n        labels = [1] * 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-20T14:45:54.571792Z","iopub.execute_input":"2021-11-20T14:45:54.572275Z","iopub.status.idle":"2021-11-20T14:45:54.59413Z","shell.execute_reply.started":"2021-11-20T14:45:54.572235Z","shell.execute_reply":"2021-11-20T14:45:54.593378Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CellTestDataset(Dataset):\n    def __init__(self, image_dir, transforms=None):\n        self.transforms = transforms\n        self.image_dir = image_dir\n        self.image_ids = [f[:-4]for f in os.listdir(self.image_dir)]\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\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":{"execution":{"iopub.status.busy":"2021-11-20T14:45:54.598925Z","iopub.execute_input":"2021-11-20T14:45:54.599147Z","iopub.status.idle":"2021-11-20T14:45:54.60875Z","shell.execute_reply.started":"2021-11-20T14:45:54.599124Z","shell.execute_reply":"2021-11-20T14:45:54.608001Z"},"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    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True, 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","metadata":{"execution":{"iopub.status.busy":"2021-11-20T14:45:54.611758Z","iopub.execute_input":"2021-11-20T14:45:54.613752Z","iopub.status.idle":"2021-11-20T14:45:54.620325Z","shell.execute_reply.started":"2021-11-20T14:45:54.613716Z","shell.execute_reply":"2021-11-20T14:45:54.619668Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def validate_model(model, ds_valid):\n    # when in eval mode, model doesn't return loss_dict. It returns predictions instead (boxes, labels, scores, masks)\n    model.train()\n    \n    loss_list = []\n    loss_mask_list = []\n    loss_by_type = collections.defaultdict(list)\n    loss_by_type_mask = collections.defaultdict(list)\n    \n    for i in range(len(ds_valid)):\n        images, targets = ds_valid[i]\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        with torch.no_grad():\n            loss_dict = model(images, targets)\n            loss = sum(x for x in loss_dict.values())            \n            loss_mask = loss_dict['loss_mask'].item()\n            loss_list.append(loss.item())\n            loss_mask_list.append(loss_mask)\n            \n            cell_type = ds_valid.image_info[i]['cell_type']\n            loss_by_type[cell_type].append(loss.item())\n            loss_by_type_mask[cell_type].append(loss_mask)\n    \n    return np.mean(loss_list), np.mean(loss_mask_list), loss_by_type, loss_by_type_mask","metadata":{"execution":{"iopub.status.busy":"2021-11-20T14:45:54.621822Z","iopub.execute_input":"2021-11-20T14:45:54.622787Z","iopub.status.idle":"2021-11-20T14:45:54.634386Z","shell.execute_reply.started":"2021-11-20T14:45:54.622741Z","shell.execute_reply":"2021-11-20T14:45:54.633659Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def calc_metric(model, ds_valid):\n    model.eval()\n    metric_list = []\n    for i in range(len(ds_valid)):\n        image, target = ds_valid[i]\n        images = list(x.to(DEVICE) for x in [image])\n        gt = target['masks']\n        cell_type = ds_valid.image_info[i]['cell_type']\n        \n        with torch.no_grad():\n            result = model(images)[0]\n            found_masks = []\n            for j, mask in enumerate(result[\"masks\"]):\n                score = result[\"scores\"][j].cpu().item()\n                if score < SCORE_THRESHOLDS[cell_type]:\n                    continue\n        \n                mask = mask.cpu().numpy()[0]\n                # Keep only highly likely pixels\n                binary_mask = mask > MASK_THRESHOLDS[cell_type]\n                binary_mask = remove_overlapping_pixels(binary_mask, found_masks)\n                if binary_mask.sum() < MIN_PIXELS:\n                    continue\n                found_masks.append(binary_mask)\n\n        metric = image_metric(gt, found_masks)\n        metric_list.append(metric)\n    return np.mean(metric_list)","metadata":{"execution":{"iopub.status.busy":"2021-11-20T14:45:54.635857Z","iopub.execute_input":"2021-11-20T14:45:54.636653Z","iopub.status.idle":"2021-11-20T14:45:54.646305Z","shell.execute_reply.started":"2021-11-20T14:45:54.636617Z","shell.execute_reply":"2021-11-20T14:45:54.645503Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_sample(ds_valid, sample_index):\n    \n    cell_type = ds_valid.image_info[sample_index]['cell_type']\n    print('Cell type:', cell_type)\n    img, target = ds_valid[sample_index]\n    gt = target['masks']\n        \n    found_masks = []\n\n    plt.imshow(img.numpy().transpose((1,2,0)))\n    plt.title(\"Image\")\n    plt.show()\n    \n    masks = gt[0]\n    for mask in gt:\n        masks = np.logical_or(masks, mask)\n    plt.imshow(img.numpy().transpose((1,2,0)))\n    plt.imshow(masks, alpha=0.3)\n    plt.title(\"Ground truth\")\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2021-11-20T14:45:54.647635Z","iopub.execute_input":"2021-11-20T14:45:54.648372Z","iopub.status.idle":"2021-11-20T14:45:54.660111Z","shell.execute_reply.started":"2021-11-20T14:45:54.648335Z","shell.execute_reply":"2021-11-20T14:45:54.659239Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_all_train = pd.read_csv(TRAIN_CSV, nrows=5000 if TEST_FLAG else None)\nprint('Samples loaded:', df_all_train.shape[0])\n# remove images with bad masks: \n# https://www.kaggle.com/tolgadincer/sartorius-eda-general-overview-and-outliers\n#df_all_train = df_all_train[~df_all_train['id'].isin(['ce5d0de993bd', 'a9fc5e872671', 'db5260527117'])].copy()\ndf_all_train = df_all_train[~df_all_train['id'].isin(['ce5d0de993bd'])].copy()\nprint('Samples remaining:', df_all_train.shape[0])","metadata":{"execution":{"iopub.status.busy":"2021-11-20T14:45:54.663529Z","iopub.execute_input":"2021-11-20T14:45:54.663962Z","iopub.status.idle":"2021-11-20T14:45:54.998007Z","shell.execute_reply.started":"2021-11-20T14:45:54.663925Z","shell.execute_reply":"2021-11-20T14:45:54.997236Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_all_ids = df_all_train.groupby('id')['cell_type'].agg('max').reset_index()\ntrain_ids, valid_ids = train_test_split(df_all_ids['id'].values, test_size=VAL_SPLIT, random_state=SEED, stratify=df_all_ids['cell_type'].values)\ndf_train = df_all_train[df_all_train['id'].isin(train_ids)]\ndf_valid = df_all_train[df_all_train['id'].isin(valid_ids)]\n\nds_train = CellDataset(TRAIN_PATH, df_train, trainFlag=True)\ndl_train = DataLoader(ds_train, batch_size=BATCH_SIZE, shuffle=True, num_workers=2, collate_fn=lambda x: tuple(zip(*x)))\n\nds_valid = CellDataset(TRAIN_PATH, df_valid, trainFlag=False)\n\nprint('Train samples:', df_train.shape[0])\nprint('Validation samples:', df_valid.shape[0])","metadata":{"execution":{"iopub.status.busy":"2021-11-20T14:45:54.999695Z","iopub.execute_input":"2021-11-20T14:45:55.000231Z","iopub.status.idle":"2021-11-20T14:46:00.590326Z","shell.execute_reply.started":"2021-11-20T14:45:55.000191Z","shell.execute_reply":"2021-11-20T14:46:00.589577Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_sample(ds_train, 1)","metadata":{"execution":{"iopub.status.busy":"2021-11-20T14:46:02.381084Z","iopub.execute_input":"2021-11-20T14:46:02.381442Z","iopub.status.idle":"2021-11-20T14:46:02.960273Z","shell.execute_reply.started":"2021-11-20T14:46:02.381409Z","shell.execute_reply":"2021-11-20T14:46:02.959604Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DEVICE = torch.device('cuda') \nmodel = get_model()\nmodel.to(DEVICE);","metadata":{"execution":{"iopub.status.busy":"2021-11-18T14:54:07.308023Z","iopub.execute_input":"2021-11-18T14:54:07.308521Z","iopub.status.idle":"2021-11-18T14:54:11.986081Z","shell.execute_reply.started":"2021-11-18T14:54:07.308474Z","shell.execute_reply":"2021-11-18T14:54:11.984988Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for param in model.parameters():\n    param.requires_grad = True\n    \nparams = [p for p in model.parameters() if p.requires_grad]\n\noptimizer = torch.optim.SGD(params, lr=LEARNING_RATE, momentum=MOMENTUM, weight_decay=WEIGHT_DECAY)\nlr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.4)\n\nn_batches = len(dl_train)","metadata":{"execution":{"iopub.status.busy":"2021-11-18T14:54:11.987432Z","iopub.execute_input":"2021-11-18T14:54:11.988602Z","iopub.status.idle":"2021-11-18T14:54:11.998891Z","shell.execute_reply.started":"2021-11-18T14:54:11.988558Z","shell.execute_reply":"2021-11-18T14:54:11.997494Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"best_metric = 0\n\nfor epoch in range(NUM_EPOCHS):\n    print(f\"Starting epoch {epoch} of {NUM_EPOCHS}\")\n    model.train()    \n    \n    time_start = time.time()\n    loss_accum = 0.0\n    loss_mask_accum = 0.0\n    \n    for batch_idx, (images, targets) in enumerate(dl_train, 1):\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        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():5.3f}. Mask-only loss: {loss_mask:5.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#    val_loss, val_mask_loss, loss_by_type, loss_by_type_mask = validate_model(model, ds_valid)\n\n#    print('Mask-only_loss: {:.4f}, total loss: {:.4f}'.format(val_mask_loss, val_loss))\n#    for cell_type in loss_by_type.keys():\n#        print('Cell type: {:6}; mask-only_loss: {:.4f}, total loss: {:.4f}'.format(cell_type, np.mean(loss_by_type_mask[cell_type]), np.mean(loss_by_type[cell_type])))\n\n    metric = calc_metric(model, ds_valid)\n    print('Validation MAP IoU: {:.4f}'.format(metric))\n    if metric > best_metric:\n        best_metric = metric\n        print('Saving a better model at epoch:', epoch)\n        torch.save(model.state_dict(), 'pytorch_model.bin')\n\nprint('\\nBest validation MAP IoU: {:.4f}'.format(best_metric))   ","metadata":{"execution":{"iopub.status.busy":"2021-11-18T14:54:12.001725Z","iopub.execute_input":"2021-11-18T14:54:12.002308Z","iopub.status.idle":"2021-11-18T15:12:57.739025Z","shell.execute_reply.started":"2021-11-18T14:54:12.002264Z","shell.execute_reply":"2021-11-18T15:12:57.737902Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#torch.save(model.state_dict(), 'pytorch_model.bin')\nwith open('train_valid_ids.npz', 'wb') as outfile:\n    np.savez(outfile, train_ids=train_ids, valid_ids=valid_ids)","metadata":{"execution":{"iopub.status.busy":"2021-11-18T15:12:57.740995Z","iopub.execute_input":"2021-11-18T15:12:57.741616Z","iopub.status.idle":"2021-11-18T15:12:58.145006Z","shell.execute_reply.started":"2021-11-18T15:12:57.741569Z","shell.execute_reply":"2021-11-18T15:12:58.144091Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Plots: the image, The image + the ground truth mask, The image + the predicted mask\ndef analyze_train_sample(model, ds_train, sample_index):\n    \n    img, targets = ds_train[sample_index]\n    plt.imshow(img.numpy().transpose((1,2,0)))\n    plt.title(\"Image\")\n    plt.show()\n    \n    masks = np.zeros((HEIGHT, WIDTH))\n    for mask in targets['masks']:\n        masks = np.logical_or(masks, mask)\n    plt.imshow(img.numpy().transpose((1,2,0)))\n    plt.imshow(masks, alpha=0.3)\n    plt.title(\"Ground truth\")\n    plt.show()\n    \n    model.eval()\n    with torch.no_grad():\n        preds = model([img.to(DEVICE)])[0]\n\n    plt.imshow(img.cpu().numpy().transpose((1,2,0)))\n    all_preds_masks = np.zeros((HEIGHT, WIDTH))\n    for mask in preds['masks'].cpu().detach().numpy():\n        all_preds_masks = np.logical_or(all_preds_masks, mask[0] > MASK_THRESHOLD)\n    plt.imshow(all_preds_masks, alpha=0.4)\n    plt.title(\"Predictions\")\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2021-11-18T15:12:58.146385Z","iopub.execute_input":"2021-11-18T15:12:58.146699Z","iopub.status.idle":"2021-11-18T15:12:58.158634Z","shell.execute_reply.started":"2021-11-18T15:12:58.146656Z","shell.execute_reply":"2021-11-18T15:12:58.157348Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}