{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install pandas pillow matplotlib sklearn torch torchvision albumentations -q\n!pip install git+https://github.com/qubvel/segmentation_models.pytorch -q\n!pip install catalyst -q\n\nimport os\n\nimport cv2\n\nimport albumentations as albu\nimport albumentations.pytorch as AT\n\nfrom catalyst.dl.runner import SupervisedRunner\nfrom catalyst.dl.callbacks import DiceCallback, EarlyStoppingCallback\n\nimport numpy as np\nimport pandas as pd\n\nimport segmentation_models_pytorch as smp\n\nimport torch\nfrom torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau, CosineAnnealingLR\nfrom torch.utils.data import TensorDataset, DataLoader, Dataset\n\nfrom sklearn.model_selection import train_test_split\n\nimport seaborn as sns\n\nimport torch.nn as nn\nfrom segmentation_models_pytorch.utils.base import Activation\n\nimport torch\n\n!pip install --upgrade pip -q\n\n%matplotlib inline","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Raw data"},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls ../input/understanding_cloud_organization","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"home = '../input/understanding_cloud_organization'\n\ndef load_image_data():\n    data = pd.read_csv('{}/train.csv'.format(home))\n    data.loc[:, 'id'] = data.loc[:, 'Image_Label'].apply(lambda x: x.split('_')[0])\n    data.loc[:, 'label'] = data.loc[:, 'Image_Label'].apply(lambda x: x.split('_')[1])\n    data.loc[data.loc[:, 'EncodedPixels'].isnull(), 'label'] = np.nan\n    data.set_index('id', drop=True, inplace=True)\n    return data\n\nraw = load_image_data()\n\ndata = raw.copy()\n\nnum_epochs = 2\nENCODER = 'resnet50'\nENCODER_WEIGHTS = 'imagenet'\nDEVICE = 'cuda'\nACTIVATION = None\n\nlogdir = \"./logs/segmentation\"\n\nbs = 32\nnum_workers = 1\n\ndata = data.iloc[:128, :]\n\nprint(data.shape)\ndata.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset analysis"},{"metadata":{"trusted":true},"cell_type":"code","source":"#  number of labels in each image\ndata.groupby('id').count().sort_values('label')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#  number of images that have a number of labels\ndata.groupby('id').count().sort_values('label').groupby('label').count()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"##  Masking images"},{"metadata":{"trusted":true},"cell_type":"code","source":"from PIL import Image\nimport matplotlib.pyplot as plt\nimport numpy as np\n\n%matplotlib inline\n\n\ndef make_single_mask_from_rle(mask_rle, shape=(1400, 2100)):\n    \" https://www.kaggle.com/artgor/segmentation-in-pytorch-using-convenient-tools \"\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.uint8)\n    for lo, hi in zip(starts, ends):\n        img[lo:hi] = 1\n    return img.reshape(shape, order='F')\n\n\ndef plot_images(images):\n    nrows = 4\n    sample = np.random.choice(images, nrows)\n    print(sample)\n\n    shape = (1400, 2100)\n\n    f, axes = plt.subplots(nrows, 4, figsize=(25, 16))\n\n    for idx, img_id in enumerate(sample):\n        sub = data.loc[data.index == img_id]\n        img = Image.open('{}/train_images/{}'.format(home, img_id))\n\n        for col, ax in enumerate(axes[:, idx]):\n            ax.imshow(img)\n            ax.set_title(sub.loc[:, 'label'].iloc[col])\n\n            try:\n                mask = make_single_mask_from_rle(sub.loc[:, 'EncodedPixels'].iloc[col], shape)\n\n            except AttributeError:\n                mask = np.zeros(shape)\n\n            ax.imshow(mask, alpha=0.5 , cmap='gray')\n            \nimages = list(set(data.index))\nplot_images(images)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_images(images)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Helpers"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_img(data, idx=None, folder='train_images'):\n    if idx is None:\n        idx = np.random.randint(0, data.shape[0])\n        \n    img_data = data.iloc[idx, :]\n    img_id = img_data.name\n    img = cv2.imread(os.path.join(home, folder, img_id))\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    mask = data.EncodedPixels\n    return img, img_id, mask\n\n\ndef get_masks(data, img_id):\n    return data.loc[data.index == img_id, :]\n\n\ndef plot_with_aug(aug, img):\n    f, axes = plt.subplots(ncols=2, figsize=(25, 16))\n    \n    axes[0].imshow(img)\n    axes[0].set_title('original')\n    \n    axes[1].imshow(aug(image=img)['image'])\n    axes[1].set_title(repr(aug)) \n    return f","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Augmentations\n\n### Horizontal flip\n\nYes"},{"metadata":{"trusted":true},"cell_type":"code","source":"f = plot_with_aug(albu.HorizontalFlip(p=0.5), get_img(data)[0])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Shift Scale Rotate\n\nNo"},{"metadata":{"trusted":true},"cell_type":"code","source":"f = plot_with_aug(albu.ShiftScaleRotate(\n    scale_limit=1, #  smaller or bigger\n    rotate_limit=10,\n    shift_limit=1, \n    p=1.0, \n    border_mode=0\n), get_img(data)[0])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Grid Distortion\n\nLooks ok"},{"metadata":{"trusted":true},"cell_type":"code","source":"f = plot_with_aug(albu.GridDistortion(\n    p=1.0,\n), get_img(data)[0])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Optical Distortion\n\nNo"},{"metadata":{"trusted":true},"cell_type":"code","source":"f = plot_with_aug(albu.OpticalDistortion(\n    p=1.0, distort_limit=2, shift_limit=0.5\n), get_img(data)[0])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Torch dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_all_masks(data, img_id, shape=(1400, 2100)):\n    label_order = ['Fish', 'Flower', 'Gravel', 'Sugar']\n    \n    mask_data = get_masks(data, img_id)\n    masks = np.zeros((*shape, len(label_order)))\n\n    for idx, expected_label in enumerate(label_order):\n        data = mask_data.iloc[idx]\n        lbl = data.loc['label']\n\n        if lbl is not np.nan:\n            assert lbl == expected_label\n            masks[:, :, idx] = make_single_mask_from_rle(data.loc['EncodedPixels'], shape)\n            \n    return masks\n\nmasks = make_all_masks(data, data.index[0])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CloudDataset(Dataset):\n    \n    def __init__(\n        self,\n        data,\n        dataset='train',\n        transform=albu.Compose([albu.Resize(320, 640)]),\n        preprocessing=None\n    ):\n        super().__init__()\n        self.data = data\n        self.folder = os.path.join(home, '{}_images'.format(dataset))\n        self.transform = transform\n        self.preprocessing = preprocessing\n        \n    def __len__(self):\n        return self.data.shape[0]\n    \n    def __getitem__(self, idx):\n        img_data = self.data.iloc[idx, :]\n        img_id = img_data.name\n                \n        img = cv2.imread(os.path.join(self.folder, img_id))\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        \n        mask = make_all_masks(self.data, img_id)\n\n        augmented = self.transform(image=img, mask=mask)\n        img = augmented['image']\n        mask = augmented['mask']\n        \n        if self.preprocessing:\n            preprocessed = self.preprocessing(image=img, mask=mask)\n            img = preprocessed['image']\n            mask = preprocessed['mask']\n            \n        return img, mask\n\n    \ndef to_tensor(x, **kwargs):\n    return x.transpose(2, 0, 1).astype('float32')\n\n\ndef get_preprocessing(preprocessing_fn):\n    _transform = [\n        albu.Lambda(image=preprocessing_fn),\n        albu.Lambda(image=to_tensor, mask=to_tensor),\n    ]\n    return albu.Compose(_transform)\n\n\ndef get_training_augmentation():\n    train_transform = [\n        albu.HorizontalFlip(p=0.5),\n        #albu.ShiftScaleRotate(scale_limit=0.5, rotate_limit=0, shift_limit=0.1, p=0.5, border_mode=0),\n        albu.GridDistortion(p=0.5),\n        #albu.OpticalDistortion(p=0.5, distort_limit=2, shift_limit=0.5),\n        albu.Resize(320, 640)\n    ]\n    return albu.Compose(train_transform)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nid_grp = data.groupby(data.index).count()\nsub = id_grp.iloc[:, :]\ntrain_id, valid_id = train_test_split(sub.index, test_size=0.2, stratify=sub.loc[:, 'label'])\nprint(train_id.shape, valid_id.shape)\n\nmodel = smp.Unet(\n    encoder_name=ENCODER, \n    encoder_weights=ENCODER_WEIGHTS, \n    classes=4,\n    activation=ACTIVATION,\n)\n\npreprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)\n\ntrain_dataset = CloudDataset(\n    data.loc[train_id, :], \n    dataset='train', \n    transform=get_training_augmentation(), \n    preprocessing=get_preprocessing(preprocessing_fn)\n)\n\nvalid_dataset = CloudDataset(\n    data.loc[valid_id, :], \n    dataset='train', \n    transform=get_training_augmentation(), \n    preprocessing=get_preprocessing(preprocessing_fn)\n)\n\nloaders = {\n    'train': DataLoader(train_dataset, batch_size=bs, shuffle=True, num_workers=num_workers),\n    'valid': DataLoader(valid_dataset, batch_size=bs, shuffle=True, num_workers=num_workers)\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(DataLoader(valid_dataset, batch_size=bs, shuffle=True, num_workers=num_workers))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# model, criterion, optimizer\noptimizer = torch.optim.Adam([\n    {'params': model.decoder.parameters(), 'lr': 1e-2}, \n    {'params': model.encoder.parameters(), 'lr': 1e-3},  \n])\n\n\ndef f_score(pr, gt, beta=1, eps=1e-7, threshold=None, activation='sigmoid'):\n    \"\"\"\n    Args:\n        pr (torch.Tensor): A list of predicted elements\n        gt (torch.Tensor):  A list of elements that are to be predicted\n        eps (float): epsilon to avoid zero division\n        threshold: threshold for outputs binarization\n    Returns:\n        float: IoU (Jaccard) score\n    \"\"\"\n\n    if activation is None or activation == \"none\":\n        activation_fn = lambda x: x\n    elif activation == \"sigmoid\":\n        activation_fn = torch.nn.Sigmoid()\n    elif activation == \"softmax2d\":\n        activation_fn = torch.nn.Softmax2d()\n    else:\n        raise NotImplementedError(\n            \"Activation implemented for sigmoid and softmax2d\"\n        )\n\n    pr = activation_fn(pr)\n\n    if threshold is not None:\n        pr = (pr > threshold).float()\n\n\n    tp = torch.sum(gt * pr)\n    fp = torch.sum(pr) - tp\n    fn = torch.sum(gt) - tp\n\n    score = ((1 + beta ** 2) * tp + eps) \\\n            / ((1 + beta ** 2) * tp + beta ** 2 * fn + fp + eps)\n\n    return score\n\n\nclass DiceLoss(nn.Module):\n    __name__ = 'dice_loss'\n\n    def __init__(self, eps=1e-7, activation='sigmoid'):\n        super().__init__()\n        self.activation = activation\n        self.eps = eps\n\n    def forward(self, y_pr, y_gt):\n        return 1 - f_score(y_pr, y_gt, beta=1., eps=self.eps, threshold=None, activation=self.activation)\n\nclass BCEDiceLoss(DiceLoss):\n    __name__ = 'bce_dice_loss'\n\n    def __init__(self, eps=1e-7, activation='sigmoid'):\n        super().__init__(eps, activation)\n        self.bce = nn.BCEWithLogitsLoss(reduction='mean')\n\n    def forward(self, y_pr, y_gt):\n        dice = super().forward(y_pr, y_gt)\n        bce = self.bce(y_pr, y_gt)\n        return dice + bce\n\nscheduler = ReduceLROnPlateau(optimizer, factor=0.15, patience=2)\ncriterion = BCEDiceLoss(eps=1.)  #  training on CE, reporting on DICE\nrunner = SupervisedRunner()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"runner.train(\n    model=model,\n    criterion=criterion,\n    optimizer=optimizer,\n    scheduler=scheduler,\n    loaders=loaders,\n    callbacks=[DiceCallback(), EarlyStoppingCallback(patience=5, min_delta=0.001)],\n    logdir=logdir,\n    num_epochs=num_epochs,\n    verbose=True\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from catalyst.dl import utils\n\nutils.plot_metrics(\n    logdir=logdir, \n    metrics=[\"loss\", \"dice\", 'lr', '_base/lr']\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from catalyst.dl.callbacks import DiceCallback, EarlyStoppingCallback, InferCallback, CheckpointCallback\nimport tqdm\n\n\nrunner.infer(\n    model=model,\n    loaders={\"valid\": loaders['valid']},\n    callbacks=[\n        CheckpointCallback(\n            resume=f\"{logdir}/checkpoints/best.pth\"),\n        InferCallback()\n    ],\n)\n\nvalid_masks = []\nprobabilities = np.zeros((2220, 350, 525))\nfor i, (batch, output) in enumerate(tqdm.tqdm(zip(\n        valid_dataset, runner.callbacks[0].predictions[\"logits\"]))):\n    image, mask = batch\n    for m in mask:\n        if m.shape != (350, 525):\n            m = cv2.resize(m, dsize=(525, 350), interpolation=cv2.INTER_LINEAR)\n        valid_masks.append(m)\n\n    for j, probability in enumerate(output):\n        if probability.shape != (350, 525):\n            probability = cv2.resize(probability, dsize=(525, 350), interpolation=cv2.INTER_LINEAR)\n        probabilities[i * 4 + j, :, :] = probability","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def post_process(probability, threshold, min_size):\n    \"\"\"\n    Post processing of each predicted mask, components with lesser number of pixels\n    than `min_size` are ignored\n    \"\"\"\n    # don't remember where I saw it\n    mask = cv2.threshold(probability, threshold, 1, cv2.THRESH_BINARY)[1]\n    num_component, component = cv2.connectedComponents(mask.astype(np.uint8))\n    predictions = np.zeros((350, 525), np.float32)\n    num = 0\n    for c in range(1, num_component):\n        p = (component == c)\n        if p.sum() > min_size:\n            predictions[p] = 1\n            num += 1\n    return predictions, num\n\n\nsigmoid = lambda x: 1 / (1 + np.exp(-x))\n\n\ndef dice(img1, img2):\n    img1 = np.asarray(img1).astype(np.bool)\n    img2 = np.asarray(img2).astype(np.bool)\n    intersection = np.logical_and(img1, img2)\n    return 2. * intersection.sum() / (img1.sum() + img2.sum())\n\nclass_params = {}\nfor class_id in range(4):\n    print(class_id)\n    attempts = []\n    for t in range(0, 100, 5):\n        t /= 100\n        for ms in [0, 100, 1200, 5000, 10000]:\n            masks = []\n            for i in range(class_id, len(probabilities), 4):\n                probability = probabilities[i]\n                predict, num_predict = post_process(sigmoid(probability), t, ms)\n                masks.append(predict)\n\n            d = []\n            for i, j in zip(masks, valid_masks[class_id::4]):\n                if (i.sum() == 0) & (j.sum() == 0):\n                    d.append(1)\n                else:\n                    d.append(dice(i, j))\n\n            attempts.append((t, ms, np.mean(d)))\n\n    attempts_df = pd.DataFrame(attempts, columns=['threshold', 'size', 'dice'])\n    attempts_df = attempts_df.sort_values('dice', ascending=False)\n    print(attempts_df.head())\n    best_threshold = attempts_df['threshold'].values[0]\n    best_size = attempts_df['size'].values[0]\n    \n    class_params[class_id] = (best_threshold, best_size)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(class_params)\nsns.lineplot(x='threshold', y='dice', hue='size', data=attempts_df);\nplt.title('Threshold and min size vs dice for one of the classes');","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.6.8"}},"nbformat":4,"nbformat_minor":4}