{"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":"pip install torch==1.10.0+cu113 torchvision==0.11.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html","metadata":{"execution":{"iopub.status.busy":"2021-11-24T03:18:00.152757Z","iopub.execute_input":"2021-11-24T03:18:00.153072Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip -q install pycocotools","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### copy some files to current dir.\n#### these files are modified from: https://github.com/pytorch/vision/tree/main/references/detection","metadata":{}},{"cell_type":"code","source":"!cp ../input/train-utils/* ./","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport torch\nimport gc\nimport cv2\nimport time\nimport random\nfrom torch import nn\nimport pandas as pd\nimport numpy as np\nimport collections\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nfrom torchvision.transforms import functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport albumentations as albu\nfrom albumentations.pytorch import ToTensorV2\nimport torchvision\nfrom engine import train_one_epoch, evaluate\nimport torchvision.transforms as T\nfrom tqdm import tqdm_notebook as tqdm\nfrom skimage.color import label2rgb\nfrom sklearn.model_selection import GroupKFold\nfrom torchvision.models.detection import MaskRCNN\nfrom torchvision.models.detection.anchor_utils import AnchorGenerator\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection.mask_rcnn import MaskRCNNPredictor\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau\n%matplotlib inline","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.read_csv('../input/sartorius-cell-instance-segmentation/train.csv');df.shape","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"folds = df.copy()\nkf = GroupKFold(n_splits=5)\nfor f, (t_idx, v_idx) in enumerate(kf.split(df, groups=df.id.values)):\n    folds.loc[v_idx, 'fold'] = int(f)\nfolds['fold'] = folds['fold'].astype(int)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CFG:\n    debug = False \n    num_workers = 0\n    precision = 16\n    device = torch.device('cuda')\n    img_dir = '../input/sartorius-cell-instance-segmentation/train/'\n    epochs = 10 \n    patience = 4\n    height = 520\n    width = 704\n    T_max = 5\n    momentum = 0.9\n    eta_min = 1e-7\n    weight_decay = 5e-4\n    mask_threshold = 0.5\n    mean = np.array([0.485, 0.456, 0.406])\n    std = np.array([0.229, 0.224, 0.225])\n    box_detections_per_img = 539\n    min_score = 0.59\n    lr = 1e-3\n    batch_size = 2\n    seed = 42\n    n_folds = 5","metadata":{"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(CFG.seed)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"label_dict = {'shsy5y':0, 'astro':1, 'cort':2}","metadata":{"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":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dataset","metadata":{}},{"cell_type":"code","source":"# https://www.kaggle.com/julian3833/sartorius-starter-torch-mask-r-cnn-lb-0-273\nclass SartoriusDataset(Dataset):\n    def __init__(self, df, transforms=None):\n        self.df = df\n        self.transforms = transforms\n        self.image_info = collections.defaultdict(dict)\n        temp_df = self.df.groupby('id')[['annotation', 'cell_type']].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                    'label': label_dict[row['cell_type'][0]],\n                    'image_path': os.path.join(CFG.img_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        label = self.image_info[idx][\"label\"]\n        info = self.image_info[idx]\n\n        n_objects = len(info['annotations'])\n        masks = np.zeros((len(info['annotations']), CFG.height, CFG.width), dtype=np.uint8)\n        boxes = []\n        \n        for i, annotation in enumerate(info['annotations']):\n            a_mask = rle_decode(annotation, (CFG.height, CFG.width))\n            a_mask = Image.fromarray(a_mask)\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        labels = [label] * 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) # uint8\n\n        image_id = torch.tensor([idx])\n        #image_id = self.image_info[idx][\"image_id\"]\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#, image_id\n\n    def __len__(self):\n        return len(self.image_info)\n    \ndef collate_fn(batch):\n    return tuple(zip(*batch))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Augmentations","metadata":{}},{"cell_type":"code","source":"# https://www.kaggle.com/julian3833/sartorius-starter-torch-mask-r-cnn-lb-0-273\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","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_transform(train):\n    transforms = []\n    transforms.append(ToTensor())\n    if train:\n        transforms.append(HorizontalFlip(0.5))\n        transforms.append(VerticalFlip(0.5))\n    return Compose(transforms)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# https://www.kaggle.com/blondinka/how-to-do-augmentations-for-instance-segmentation\ndef visualize_bbox(img, bbox, color=(255, 0, 255), thickness=2):  \n    \"\"\"Helper to add bboxes to images \n    Args:\n        img : image as open-cv numpy array\n        bbox : boxes as a list or numpy array in pascal_voc fromat [x_min, y_min, x_max, y_max]  \n        color=(255, 255, 0): boxes color \n        thickness=2 : boxes line thickness\n    \"\"\"\n    x_min, y_min, x_max, y_max = bbox\n    x_min, y_min, x_max, y_max = int(x_min), int(y_min), int(x_max), int(y_max)\n    cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color=color, thickness=thickness)\n    return img\n\ndef plot_image_aug(image, masks, boxes):\n    # glue masks together\n    one_mask = np.zeros_like(masks[0])\n    for i, mask in enumerate(masks):\n        #one_mask += (mask > 0).astype(np.uint8) * (11-i)\n        one_mask = np.add(one_mask, (mask > 0).astype(np.uint8) * (11-i), out=one_mask, casting=\"unsafe\")\n    \n    for box in boxes:\n        #print(box)\n        image = visualize_bbox(np.ascontiguousarray(image), box)  \n        \n    # for binary masks we get one channel and need to convert to RGB for visualization\n    mask_rgb = label2rgb(one_mask, bg_label=0)            \n    \n    plt.figure(figsize=(12,12))\n    plt.figure(1)\n    ax1 = plt.subplot(121)\n    plt.imshow(image)\n    plt.title('image')   \n    ax2 = plt.subplot(122)\n    plt.imshow(mask_rgb)\n    plt.title('mask')   \n    plt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sample_dataset = SartoriusDataset(folds, transforms=get_transform(True))\nsample_dataloader = DataLoader(sample_dataset, batch_size=4, collate_fn=collate_fn, shuffle=False)\nsamples = iter(sample_dataloader).next()\nimages, targets = samples[0], samples[1]\nfor it, (image, target) in enumerate(zip(images, targets)):\n    plot_image_aug(((image.permute(1,2,0))*255.0).numpy().astype(np.uint8), target['masks'].numpy(), target['boxes'])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model","metadata":{}},{"cell_type":"code","source":"# def get_model():\n#     backbone = torchvision.models.efficientnet_b4(pretrained=True).features   #1792\n#     backbone.out_channels = 1792\n#     anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),), aspect_ratios=((0.5, 1.0, 2.0),))\n#     roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], output_size=7, sampling_ratio=2)\n#     mask_roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], output_size=14, sampling_ratio=2)\n#     # put the pieces together inside a MaskRCNN model\n#     model = MaskRCNN(backbone, num_classes=3, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler, mask_roi_pool=mask_roi_pooler)\n#     return model","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_model(num_classes=3):\n    # load an instance segmentation model pre-trained on COCO\n    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True,\n                                                              box_detections_per_img=600)\n\n    # get 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,\n                                                       hidden_layer,\n                                                       num_classes)\n\n    return model","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_loop(fold):\n    # ------------\n    # data\n    # ------------\n    trn_idx = folds[folds['fold'] != fold].index\n    val_idx = folds[folds['fold'] == fold].index\n\n    train_folds = folds.loc[trn_idx].reset_index(drop=True)\n    valid_folds = folds.loc[val_idx].reset_index(drop=True)\n    \n    train_dataset = SartoriusDataset(train_folds, transforms=get_transform(True))\n    valid_dataset = SartoriusDataset(valid_folds, transforms=get_transform(False))\n    \n    train_loader = DataLoader(train_dataset,\n                              batch_size=CFG.batch_size, \n                              shuffle=True,\n                              worker_init_fn=lambda id: np.random.seed(torch.initial_seed() // 2 ** 32 + id),\n                              collate_fn=collate_fn,\n                              num_workers=CFG.num_workers, pin_memory=True, drop_last=False)\n    valid_loader = DataLoader(valid_dataset, \n                              batch_size=CFG.batch_size, \n                              shuffle=False, \n                              collate_fn=collate_fn,\n                              worker_init_fn=lambda id: np.random.seed(torch.initial_seed() // 2 ** 32 + id),\n                              num_workers=CFG.num_workers, pin_memory=True, drop_last=False)\n    \n    # ------------\n    # model\n    # ------------\n    model = get_model()\n    for param in model.parameters():\n        param.requires_grad = True\n    params = [p for p in model.parameters() if p.requires_grad]\n    model.to(CFG.device)\n    #optimizer = torch.optim.Adam(model.parameters(), lr=CFG.lr, weight_decay=CFG.weight_decay)\n    optimizer = torch.optim.SGD(params, lr=CFG.lr, momentum=CFG.momentum, weight_decay=CFG.weight_decay)\n    #scheduler = CosineAnnealingLR(optimizer, T_max=CFG.T_max, eta_min=CFG.eta_min)\n    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)\n    \n    best_metric = -1.\n    for epoch in range(CFG.epochs):\n        # train for one epoch, printing every 10 iterations\n        train_one_epoch(model, optimizer, train_loader, CFG.device, epoch, print_freq=200)\n        # update the learning rate\n        scheduler.step()\n        # evaluate on the test dataset\n        coco_evaluator = evaluate(model, valid_loader, device=CFG.device)\n        metric = coco_evaluator.coco_eval['segm'].stats[0]\n        if best_metric < metric:\n            best_metric = metric\n            torch.save(model.state_dict(), os.path.join(f'./best_{best_metric}_fold{fold}.pth'))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for i in range(CFG.n_folds):\n    train_loop(i)\n    break  ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"==========================================================================================================","metadata":{}},{"cell_type":"code","source":"## plz upvote if you like it.\n## And do help me to improve this, I can't get a good score with these settings.","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}