{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Multi-label segmentation using fastai "},{"metadata":{},"cell_type":"markdown","source":"It is gravel or sugar? Part of this project is to  find out if humans even agree on these cloud patterns. Different persons will label clouds differently: the same area can be seen has gravel by one person and sugar by another one. Consequently, there is overlap in the labels, i.e. each pixel can belong to several classes. \n\n[Fastai](https://docs.fast.ai/) integrates convenient features for data augmentation and training which should be useful later. So it would be nice to use it in this project. However, multi-label segmentation isn't supported yet.\n\nFor now, I focused on adding multi-label support to fastai segmentation. To achieve this, the `MultiLabelSegmentationLabelList` and `MultiLabelImageSegment` classes have been created. Each RLE-encoded mask is loaded into a separate image segment channel. This way overlapping label information is preserved.\n\nA simple Unet with ResNet18 model is trained to verify that everything works correctly.\n\nThe `test_images` are loaded in the databunch `test` dataset and utilized at the end of this notebook to make predictions and write submission file.\n\nThis is only a start, but I hope that this will be helpful to others."},{"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","trusted":true},"cell_type":"code","source":"# Python 3 environment defined by the kaggle/python docker image: https://github.com/kaggle/docker-python\n\nfrom datetime import datetime\nfrom time import time\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\nfrom tqdm import tqdm\n\nimport matplotlib\nimport matplotlib.pyplot as plt\nfrom matplotlib.lines import Line2D\n\nimport warnings\n\nimport fastai\nfrom fastai.vision import *\nfrom fastai.callbacks import *\nfrom fastai.callbacks.hooks import *\nfrom fastai.utils.mem import *","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"from fastai.utils.show_install import show_install; show_install()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"!nvidia-smi","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Code utils"},{"metadata":{"trusted":false},"cell_type":"code","source":"def fmt_now():\n    return datetime.today().strftime('%Y%m%d-%H%M%S')","execution_count":null,"outputs":[]},{"metadata":{"toc-hr-collapsed":false},"cell_type":"markdown","source":"## EDA"},{"metadata":{"trusted":false},"cell_type":"code","source":"path = Path(\"/kaggle/input/understanding_cloud_organization\")\npath.ls()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"path_img = path/\"train_images\"\n\nfnames_train = get_image_files(path_img)\nfnames_train[:3]\nprint(len(fnames_train))","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"path_test = path/\"test_images\"\n\nfnames_test = get_image_files(path_test)\nfnames_test[:3]\nprint(len(fnames_test))","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"img_f = fnames_train[1]\nimg = open_image(img_f)\nimg.show(figsize=(10, 10))","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def split_img_label(img_lbl):\n    \"\"\"Return image and label from file name like '0011165.jpg_Flower'\"\"\"\n    s = img_lbl.split(\"_\")\n    assert len(s) == 2\n    return s[0], s[1]","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"train = pd.read_csv(f'{path}/train.csv')\n\n# split Image_Label\ntrain[\"Image\"] = train[\"Image_Label\"].apply(lambda img_lbl: split_img_label(img_lbl)[0])\ntrain[\"Label\"] = train[\"Image_Label\"].apply(lambda img_lbl: split_img_label(img_lbl)[1])\ndel train[\"Image_Label\"]\n\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"train_with_mask = train.dropna(subset=[\"EncodedPixels\"])\nax = train_with_mask[\"Label\"].value_counts().plot(kind=\"pie\", autopct='%1.1f%%', title=\"Shares of each classes\", figsize=(10, 6))","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"class_counts = train.dropna(subset=[\"EncodedPixels\"]).groupby(\"Image\")[\"Label\"].nunique()\nax = class_counts.plot(kind=\"hist\", title=\"Number of classes per image\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"# pivot to have one row per image and masks as columns\ntrain = train.pivot(index='Image', columns='Label', values='EncodedPixels')\nassert len(train) == len(fnames_train) # sanity check\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Broken images\n\nImages which look incorrect."},{"metadata":{"trusted":false},"cell_type":"code","source":"def show_img_fn(fname, figsize=(10, 10)):\n    img = open_image(fname)\n    img.show(figsize=figsize)    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def show_img_info(fname):\n    show_img_fn(path_img/fname)\n    display(train.loc[[fname]])   ","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"unusual_imgs = [\"1588d4c.jpg\", \"c0306e5.jpg\", \"c26c635.jpg\", \"fa645da.jpg\", \"41f92e5.jpg\", \"e5f2f24.jpg\"]","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":false},"cell_type":"code","source":"for fname in unusual_imgs:\n    img = open_image(path_img/fname)\n    img.show(figsize=(5, 5), title=fname)     ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Convert masks from/to RLE"},{"metadata":{"trusted":false},"cell_type":"code","source":"train_img_dims = (1400, 2100)  # Train and test images are 1400x2100 pixels","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"From [evaluation](https://www.kaggle.com/c/understanding_cloud_organization/overview/evaluation):\n> start position and a run length. E.g. '1 3' implies starting at pixel 1 and running a total of 3 pixels (1,2,3).\n\n> The pixels are numbered from top to bottom, then left to right: 1 is pixel (1,1), 2 is pixel (2,1), etc."},{"metadata":{"trusted":false},"cell_type":"code","source":"def rle_to_mask(rle, shape):\n    mask_img = open_mask_rle(rle, shape)\n    mask = mask_img.px.permute(0, 2, 1)\n    return mask","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def mask_to_rle(mask):\n    \"\"\"Convert binary `mask` to RLE string\"\"\"\n    return rle_encode(mask.numpy().T)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def test_mask_rle():\n    \"\"\"test case for mask RLE encode/decode\"\"\"\n    mask_rle = train.iloc[0][\"Fish\"]    \n    mask = rle_to_mask(mask_rle, train_img_dims)\n    mask_rle_enc = mask_to_rle(mask)\n    assert mask_rle_enc == mask_rle\n    \n    print(mask.shape)\n    Image(mask).show()\n    \ntest_mask_rle()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load images"},{"metadata":{"trusted":false},"cell_type":"code","source":"# TODO remove use_partial_data()\nitem_list = (SegmentationItemList.\n             from_df(df=train.reset_index(), path=path_img, cols=\"Image\")\n             .use_partial_data(sample_pct=0.1)  # use only a subset of data to speedup tests\n             .split_by_rand_pct(0.2))","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"class MultiLabelImageSegment(ImageSegment):\n    \"\"\"Store overlapping masks in separate image channels\"\"\"\n\n    def show(self, ax:plt.Axes=None, figsize:tuple=(3,3), title:Optional[str]=None, hide_axis:bool=True,\n        cmap:str='tab20', alpha:float=0.5, class_names=None, **kwargs):\n        \"Show the masks on `ax`.\"\n             \n        # put all masks into a single channel\n        flat_masks = self.px[0:1, :, :].clone()\n        for idx in range(1, self.shape[0]): # shape CxHxW\n            mask = self.px[idx:idx+1, :, :] # slice tensor to a single mask channel\n            # use powers of two for class codes to keep them distinguishable after sum \n            flat_masks += mask * 2**idx\n        \n        # use same color normalization in image and legend\n        norm = matplotlib.colors.Normalize(vmin=0, vmax=2**self.shape[0]-1)\n        ax = show_image(Image(flat_masks), ax=ax, hide_axis=hide_axis, cmap=cmap, norm=norm,\n                        figsize=figsize, interpolation='nearest', alpha=alpha, **kwargs)\n        \n        # custom legend, see https://matplotlib.org/3.1.1/gallery/text_labels_and_annotations/custom_legends.html\n        cm = matplotlib.cm.get_cmap(cmap)\n        legend_elements = []\n        for idx in range(self.shape[0]):\n            c = 2**idx\n            label = class_names[idx] if class_names is not None else f\"class {idx}\"\n            line = Line2D([0], [0], color=cm(norm(c)), label=label, lw=4)\n            legend_elements.append(line)\n        ax.legend(handles=legend_elements)\n        \n        # debug info\n        # ax.text(10, 10, f\"px={self.px.size()}\", {\"color\": \"white\"})\n        \n        if title: ax.set_title(title)\n\n    def reconstruct(self, t:Tensor): \n        return MultiClassImageSegment(t)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"# source: https://forums.fast.ai/t/unet-how-to-get-4-channel-output/54674/4\ndef bce_logits_floatify(input, target, reduction='mean'):\n    return F.binary_cross_entropy_with_logits(input, target.float(), reduction=reduction)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"class MultiLabelSegmentationLabelList(SegmentationLabelList):\n    \"\"\"Return a single image segment with all classes\"\"\"\n    # adapted from https://forums.fast.ai/t/how-to-load-multiple-classes-of-rle-strings-from-csv-severstal-steel-competition/51445/2\n    \n    def __init__(self, items:Iterator, src_img_size=None, classes:Collection=None, **kwargs):\n        super().__init__(items=items, classes=classes, **kwargs)\n        self.loss_func = bce_logits_floatify\n        self.src_img_size = src_img_size\n        # add attributes to copy by new() \n        self.copy_new += [\"src_img_size\"]\n    \n    def open(self, rles):        \n        # load mask at full resolution\n        masks = torch.zeros((len(self.classes), *self.src_img_size)) # shape CxHxW\n        for i, rle in enumerate(rles):\n            if isinstance(rle, str):  # filter out NaNs\n                masks[i] = rle_to_mask(rle, self.src_img_size)\n        return MultiLabelImageSegment(masks)\n    \n    def analyze_pred(self, pred, thresh:float=0.0):\n        # binarize masks\n        return (pred > thresh).float()\n    \n    def reconstruct(self, t:Tensor): \n        return MultiLabelImageSegment(t)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"class_names = [\"Fish\", \"Flower\", \"Gravel\", \"Sugar\"]","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def get_masks_rle(img):\n    \"\"\"Get RLE-encoded masks for this image\"\"\"\n    img = img.split(\"/\")[-1]  # get filename only\n    return train.loc[img, class_names].to_list()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"# reduce image size\n# img_size = tuple(v // 16 for v in train_img_dims)\nimg_size = (84, 132)  # use multiple of 4\nimg_size","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"classes = [0, 1, 2, 3] # no need for a \"void\" class: if a pixel isn't in any mask, it is not labelled\nitem_list = item_list.label_from_func(func=get_masks_rle, label_cls=MultiLabelSegmentationLabelList, \n                                      classes=classes, src_img_size=train_img_dims)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"# add unlabelled test images\n# set empty RLE string as label to produce empty multi-label masks and allow reconstruct() and show()\nitem_list = item_list.add_test_folder(path_test, label=\"\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"batch_size = 8\n\n# TODO add data augmentation\ntfms = ([], [])\n# tfms = get_transforms()\n\nitem_list = item_list.transform(tfms, tfm_y=True, size=img_size)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"data = (item_list\n        .databunch(bs=batch_size)\n        .normalize(imagenet_stats) # use same stats as pretrained model\n       )  \nassert data.test_ds is not None","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"data.show_batch(2, figsize=(15, 10), class_names=class_names)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train"},{"metadata":{"trusted":false},"cell_type":"code","source":"# adapted from: https://www.kaggle.com/iafoss/unet34-dice-0-87\n# can use sigmoid on the input too, in this case the threshold would be 0.5\ndef dice_metric(pred, targs, threshold=0):\n    pred = (pred > threshold).float()\n    targs = targs.float()  # make sure target is float too\n    return 2.0 * (pred*targs).sum() / ((pred+targs).sum() + 1.0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"metrics = [dice_metric]\n\ncallback_fns = [\n    # update a graph of learner stats and metrics after each epoch\n    ShowGraph,\n\n    # save model at every metric improvement\n    partial(SaveModelCallback, every='improvement', monitor='dice_metric', name=f\"{fmt_now()}_unet_resnet18_stage1_best\"),\n    \n    # stop training if metric no longer improve\n    partial(EarlyStoppingCallback, monitor='dice_metric', min_delta=0.01, patience=2),\n]\n\nlearn = unet_learner(data, models.resnet18, metrics=metrics, wd=1e-2, callback_fns=callback_fns)\nlearn.model_dir = \"/kaggle/working/\"  # point to writable directory","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"learn.loss_func","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"learn.summary()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"learn.lr_find()\nlearn.recorder.plot()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"learn.fit_one_cycle(15, max_lr=1e-4)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"learn.recorder.plot_metrics()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"learn.save(f\"{fmt_now()}_unet_resnet18_stage1\", return_path=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"!ls -lth {learn.model_dir}","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"# learn = learn.load(Path(learn.model_dir)/\"20190924-095959_unet_resnet18_stage1_best\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Unfreeze and differential learning rate"},{"metadata":{"trusted":false},"cell_type":"code","source":"learn.unfreeze()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"learn.lr_find()\nlearn.recorder.plot()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"# slice(start,end) syntax: the first group's learning rate is start, the last is end, and the remaining are evenly geometrically spaced\nlearn.fit_one_cycle(15, max_lr=slice(1e-6, 1e-5))","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"learn.save(f\"{fmt_now()}_unet_resnet18_stage2\", return_path=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Predictions"},{"metadata":{"trusted":false},"cell_type":"code","source":"learn.show_results(imgsize=8, class_names=class_names)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Get predictions for test images"},{"metadata":{"trusted":false},"cell_type":"code","source":"preds, _ = learn.get_preds(ds_type=DatasetType.Test, with_loss=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"preds.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"learn.show_results(ds_type=DatasetType.Test, imgsize=8, class_names=class_names)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"for i in range(3):\n    pimg = MultiLabelImageSegment(preds[i] > 0)\n    pimg.show(figsize=(6, 6), class_names=class_names)   ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Resize predictions to submission size"},{"metadata":{},"cell_type":"markdown","source":"As per competition rules, predicted masks must have 1/4 of original size for submission"},{"metadata":{"trusted":false},"cell_type":"code","source":"def resize_pred_masks(preds, shape=(4, 350, 525)):\n    \"\"\"Resize predicted masks and return them as a generator\"\"\"\n    for p in range(preds.shape[0]):\n        mask = MultiLabelImageSegment(preds[p])\n        yield mask.resize(shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"pred_masks = resize_pred_masks(preds)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Write submission file"},{"metadata":{"trusted":false},"cell_type":"code","source":"test_fnames = [p.name for p in data.test_dl.items]\nlen(test_fnames)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def write_submission_file(filename, test_fnames, preds, threshold=0):\n    with open(filename, mode='w') as f:\n        f.write(\"Image_Label,EncodedPixels\\n\")\n\n        for img_name, masks in zip(tqdm(test_fnames), resize_pred_masks(preds)):\n            binary_masks = masks.px > threshold # TODO use activation instead\n            \n            for class_idx, class_name in enumerate(class_names):\n                rle = mask_to_rle(binary_masks[class_idx])\n                f.write(f\"{img_name}_{class_name},{rle}\\n\")\n\n    print(f\"Wrote '{f.name}'.\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"submission_file = f\"{fmt_now()}_submission.csv\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"write_submission_file(submission_file, test_fnames, preds)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"","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.7.3"}},"nbformat":4,"nbformat_minor":1}