{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.9","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":22990,"datasetId":1136396,"databundleVersionId":2048213},{"sourceType":"datasetVersion","sourceId":2060255,"datasetId":1185464,"databundleVersionId":2100431},{"sourceType":"kernelVersion","sourceId":61317232},{"sourceType":"kernelVersion","sourceId":62298880},{"sourceType":"kernelVersion","sourceId":62424641}],"dockerImageVersionId":30068,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# HuBMAP - Efficient Sampling Ensemble (deepflash2, pytorch, fastai) [train]\n\n> Kernel for model training with efficient region based sampling.\n\nRequires deepflash2 (git version), zarr, and segmentation-models-pytorch\n","metadata":{"id":"OcsetTMwKXqC"}},{"cell_type":"markdown","source":"## Overview\n\n1. Installation and package loading\n2. Helper functions and patches\n3. Configuration\n4. Training\n\n### Inputs\n- https://www.kaggle.com/matjes/hubmap-zarr converted images (downscaled with factor 2)\n- https://www.kaggle.com/matjes/hubmap-labels-pdf-0-5-0-25-0-01 masks and weights for sampling\n\n### Versions\n- V7: Fixed augmentations in deepflash2 `RandomTileDataset` config (random zoom) - LB 0.913\n- V8: Adding *albumentations* transforms, switching to Cross-entropy loss","metadata":{}},{"cell_type":"markdown","source":"## Motivation\n\n### Background\n\nA glomerulus is a network of small blood vessels located at the beginning of a nephron in the kidney ([Wikipedia](https://en.wikipedia.org/wiki/Glomerulus_(kidney))\n)). Glomeruli are mainly found in the renal **cortex**, while the renal **medulla** contains mainly the renal tubule. Since we are dealing with biological structures, the separation is not not absolute and the transitions are not always perfectly sharp.\n\n![Diagram of a nephron](http://s3-us-west-2.amazonaws.com/courses-images/wp-content/uploads/sites/1842/2017/05/26234530/m9skcbftjqzrokkkopam.png)\n[Diagram of a nephron from libretexts.org, Introductory and General Biology](https://bio.libretexts.org/Bookshelves/Introductory_and_General_Biology/Book%3A_General_Biology_(Boundless)/41%3A_Osmotic_Regulation_and_the_Excretory_System/41.4%3A_Human_Osmoregulatory_and_Excretory_Systems/41.4B%3A_Nephron%3A_The_Functional_Unit_of_the_Kidney)\n\n### Key Idea\n\nA common approach to deal with the very large (>500MB - 5GB) TIFF files in the dataset is to decompose the images in smaller patches/tiles, for instance by using a sliding window apporach.\n> **Knowing that the glomeruli are mainly found in the cortex, we should focus on this region during training**. \n\nInstead of preprocessing the images and saving them into fixed tiles, we sample tiles from the entire images with a higher probability on tiles that contain glumeroli and cortex. Have a look at [this kernel](https://www.kaggle.com/matjes/hubmap-labels-pdf-0-5-0-25-0-01) for more details.\n\n\n## Advantages of this approach\n\nIn combination with [deepflash2](https://github.com/matjesg/deepflash2/tree/master/) and the deepflash2 [pytorch datasets](https://matjesg.github.io/deepflash2/data.html#Datasets) in particular, this approach has several advantages:\n- no preprocessing of the data (only saving them to .zarr files for memory efficient loading)\n    - flexible tile shapes (input shapes, e.g. 1024, 512, 256) at runtime\n    - flexible scaling (e.g., by facors of 2,3,4)\n- faster convergence during traing (~30 min for training a competitive model)\n    - focusing on the relevant regions (e.g., tiles that contain glumeroli and cortex)\n    - \"additional\" data augmentation from random sampling (compared to fixed windows)","metadata":{}},{"cell_type":"markdown","source":"### Installation and package loading","metadata":{}},{"cell_type":"code","source":"!pip install -q ../input/deepflash2-lfs\n!git clone https://github.com/qubvel/segmentation_models.pytorch.git\n!pip install -q ./segmentation_models.pytorch","metadata":{"_kg_hide-output":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-01-17T03:55:19.684565Z","iopub.execute_input":"2024-01-17T03:55:19.684999Z","iopub.status.idle":"2024-01-17T03:55:57.279516Z","shell.execute_reply.started":"2024-01-17T03:55:19.684901Z","shell.execute_reply":"2024-01-17T03:55:57.278549Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# imports\nimport zarr, cv2, random, torch\nimport numpy as np, pandas as pd\nimport albumentations as A\nimport segmentation_models_pytorch as smp\nfrom fastai.vision.all import *\nfrom deepflash2.all import *\nfrom scipy import interpolate\nfrom torch.utils.data import Dataset\nfrom sklearn.model_selection import KFold\nfrom hubmap_loss_metrics import *","metadata":{"execution":{"iopub.status.busy":"2024-01-17T03:55:57.281687Z","iopub.execute_input":"2024-01-17T03:55:57.281970Z","iopub.status.idle":"2024-01-17T03:56:04.731007Z","shell.execute_reply.started":"2024-01-17T03:55:57.281941Z","shell.execute_reply":"2024-01-17T03:56:04.730118Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Helper functions and patches","metadata":{}},{"cell_type":"code","source":"# Patch for deepflash2 'DeformationField' class, see https://fastcore.fast.ai/basics.html#patch\n@patch\ndef apply(self:DeformationField, data, offset=(0, 0), pad=(0, 0), order=1):\n    \"Apply deformation field to image using interpolation\"\n    outshape = tuple(int(s - p) for (s, p) in zip(self.shape, pad))\n    coords = [np.squeeze(d).astype('float32').reshape(*outshape) for d in self.get(offset, pad)]\n    # Get slices to avoid loading all data (.zarr files)\n    sl = []\n    for i in range(len(coords)):\n        cmin, cmax = int(coords[i].min()), int(coords[i].max())\n        dmax = data.shape[i]\n        if cmin<0: \n            cmax = max(-cmin, cmax)\n            cmin = 0 \n        elif cmax>dmax:\n            cmin = min(cmin, 2*dmax-cmax)\n            cmax = dmax\n            coords[i] -= cmin\n        else: coords[i] -= cmin\n        sl.append(slice(cmin, cmax))    \n    if len(data.shape) == len(self.shape) + 1:\n        tile = np.empty((*outshape, data.shape[-1]))\n        for c in range(data.shape[-1]):\n            # Adding divide\n            tile[..., c] = cv2.remap(data[sl[0],sl[1], c]/255, coords[1],coords[0], interpolation=order, borderMode=cv2.BORDER_REFLECT)\n    else:\n        tile = cv2.remap(data[sl[0], sl[1]], coords[1], coords[0], interpolation=order, borderMode=cv2.BORDER_REFLECT)\n    return tile","metadata":{"execution":{"iopub.status.busy":"2024-01-17T03:56:04.732562Z","iopub.execute_input":"2024-01-17T03:56:04.732853Z","iopub.status.idle":"2024-01-17T03:56:04.749373Z","shell.execute_reply.started":"2024-01-17T03:56:04.732823Z","shell.execute_reply":"2024-01-17T03:56:04.747690Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class HubmapRandomTileDataset(Dataset):\n    \"\"\"\n    Pytorch Dataset that creates random tiles with augmentations from the input images.\n    \"\"\"\n    n_inp = 1\n    def __init__(self, \n                 files,\n                 label_path,\n                 cdf_path, \n                 df_stats, \n                 sample_multiplier=50,\n                 tile_shape = (512,512),\n                 scale = 1,\n                 flip = True,                                \n                 rotation_range_deg = (0, 360),     \n                 deformation_grid = (150,150), \n                 deformation_magnitude = (10,10),\n                 value_minimum_range = (0, 0), \n                 value_maximum_range = (1, 1), \n                 value_slope_range = (1, 1),\n                 albumentations_tfms=None,\n                 **kwargs\n                ):\n        store_attr('files, df_stats, sample_multiplier, tile_shape, scale, albumentations_tfms')\n        store_attr('flip, rotation_range_deg, deformation_grid, deformation_magnitude, value_minimum_range, value_maximum_range, value_slope_range')\n        \n        self.data = zarr.open_group(self.files[0].parent.as_posix(), mode='r')\n        self.labels = zarr.open_group(label_path)\n        self.cdfs = zarr.open_group(cdf_path)\n        \n        self.indices = []\n        self.center_indices = []\n        self.df_stats = self.df_stats[self.df_stats.index.isin([f.stem for f in self.files],  level=0)]\n        print('Preparing sampling')\n        for key, grp in self.df_stats.groupby('idx'):\n            for (idx, i), row in grp.iterrows():\n                self.indices.append(idx)\n                self.center_indices.append(i)\n            for _ in range(self.sample_multiplier):\n                self.indices.append(idx)\n                self.center_indices.append(None)         \n        self.on_epoch_end()\n    \n    def __len__(self):\n        return len(self.indices)\n\n    def __getitem__(self, idx):\n        \n        if torch.is_tensor(idx): idx = idx.tolist()       \n        file_name = self.indices[idx]\n        center_idx = self.center_indices[idx]\n\n        img = self.data[file_name]\n        n_channels = img.shape[-1]\n\n        lbl = self.labels[file_name]\n        cdf = self.cdfs[file_name]\n\n        center = self.random_center(cdf[:], lbl.shape, scale=512, file=file_name, center_idx=center_idx)\n        X = self.gammaFcn(self.deformationField.apply(img, center).flatten()).reshape((*self.tile_shape, n_channels))\n        Y = self.deformationField.apply(lbl, center, (0,0), 0)\n\n        if self.albumentations_tfms:\n            augmented = self.albumentations_tfms(image=(X*255).astype('uint8'),mask=Y.astype('uint8'))\n            X = (augmented['image']/255)\n            Y = augmented['mask']\n\n        X = X.transpose(2, 0, 1).astype('float32')\n        Y = Y.astype('int64')\n        \n        return  TensorImage(X), TensorMask(Y)\n        \n    def random_center(self, cdf, orig_shape, file, center_idx, scale=512):\n        'Sample random center'\n        if center_idx:\n            stats = self.df_stats.loc[file, center_idx]\n            cx = random.randrange(stats.top, stats.top+stats.height)\n            cy = random.randrange(stats.left, stats.left+stats.width)\n        else:\n            scale_y = int((orig_shape[1]/orig_shape[0])*scale)\n            cx, cy = np.unravel_index(np.argmax(cdf > np.random.random()), (scale,scale_y))\n            cx = int(cx*orig_shape[0]/scale)\n            cy = int(cy*orig_shape[1]/scale_y)\n        return cx, cy\n        \n    def on_epoch_end(self, verbose=True):\n\n        if verbose: print(\"Generating deformation field\")\n        self.deformationField = DeformationField(self.tile_shape, self.scale)\n\n        if self.rotation_range_deg[1] > self.rotation_range_deg[0]:\n            self.deformationField.rotate(\n                theta=np.pi * (np.random.random()\n                            * (self.rotation_range_deg[1] - self.rotation_range_deg[0])\n                            + self.rotation_range_deg[0])\n                            / 180.0)\n\n        if self.flip:\n            self.deformationField.mirror(np.random.choice((True,False),2))\n\n        if self.deformation_grid is not None:\n            self.deformationField.addRandomDeformation(\n                self.deformation_grid, self.deformation_magnitude)\n\n        if verbose: print(\"Generating value augmentation function\")\n        minValue = (self.value_minimum_range[0]\n            + (self.value_minimum_range[1] - self.value_minimum_range[0])\n            * np.random.random())\n\n        maxValue = (self.value_maximum_range[0]\n            + (self.value_maximum_range[1] - self.value_maximum_range[0])\n            * np.random.random())\n\n        intermediateValue = 0.5 * (\n            self.value_slope_range[0]\n            + (self.value_slope_range[1] - self.value_slope_range[0])\n            * np.random.random())\n\n        self.gammaFcn = interpolate.interp1d([0, 0.5, 1.0], [minValue, intermediateValue, maxValue], kind=\"quadratic\")  \n        ","metadata":{"execution":{"iopub.status.busy":"2024-01-17T03:56:04.751638Z","iopub.execute_input":"2024-01-17T03:56:04.752095Z","iopub.status.idle":"2024-01-17T03:56:04.808103Z","shell.execute_reply.started":"2024-01-17T03:56:04.752052Z","shell.execute_reply":"2024-01-17T03:56:04.807306Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class HubmapValidationDataset(Dataset):\n    \"Pytorch Dataset that creates random tiles for validation and prediction on new data.\"\n    n_inp = 1\n    def __init__(self, \n                 files, \n                 label_path, \n                 tile_shape = (512,512),\n                 scale=1,\n                 val_length=None, \n                 val_seed=42, \n                 **kwargs\n                ):\n        store_attr('files, label_path, tile_shape, scale, val_seed')\n        self.data = zarr.open_group(self.files[0].parent.as_posix())\n        self.labels = zarr.open_group(label_path)\n        self.output_shape = self.tile_shape\n        self.tiler = DeformationField(self.tile_shape, scale=self.scale)\n        self.image_indices = []\n        self.image_shapes = []\n        self.centers = []\n        self.valid_indices = None\n\n        j = 0\n        for i, file in enumerate(progress_bar(self.files, leave=False)):\n            img = self.data[file.name]\n            \n            # Tiling\n            data_shape = tuple(int(x//self.scale) for x in img.shape[:-1])\n            start_points = [o//2 for o in self.output_shape]\n            end_points = [(s - st) for s, st in zip(data_shape, start_points)]\n            n_points = [int((s)//(o))+1 for s, o in zip(data_shape, self.output_shape)]\n            center_points = [np.linspace(st, e, num=n, endpoint=True, dtype=np.int64) for st, e, n in zip(start_points, end_points, n_points)]\n            for cx in center_points[1]:\n                for cy in center_points[0]:\n                    self.centers.append((int(cy*self.scale), int(cx*self.scale)))\n                    self.image_indices.append(i)\n                    self.image_shapes.append(data_shape)\n                    j += 1\n        \n        if val_length:\n            if val_length>len(self.image_shapes):\n                print(f'Reducing validation from lenght {val_length} to {len(self.image_shapes)}')\n                val_length = len(self.image_shapes)\n            np.random.seed(self.val_seed)\n            choice = np.random.choice(len(self.image_indices), val_length, replace=False)\n            self.valid_indices = {i:idx for i, idx in  enumerate(choice)}\n\n    def __len__(self):\n        if self.valid_indices: return len(self.valid_indices)\n        else: return len(self.image_shapes)\n\n    def __getitem__(self, idx):\n        if torch.is_tensor(idx):\n            idx = idx.tolist()\n        if self.valid_indices: idx = self.valid_indices[idx]\n        img_path = self.files[self.image_indices[idx]]\n        img = self.data[img_path.name]\n        centerPos = self.centers[idx]\n        X = self.tiler.apply(img, centerPos)       \n        X = X.transpose(2, 0, 1).astype('float32')\n        \n        lbl = self.labels[img_path.name]\n        Y = self.tiler.apply(lbl, centerPos, (0,0), order=0).astype('int64')\n        \n        return  TensorImage(X), TensorMask(Y)   ","metadata":{"execution":{"iopub.status.busy":"2024-01-17T03:56:04.811846Z","iopub.execute_input":"2024-01-17T03:56:04.812375Z","iopub.status.idle":"2024-01-17T03:56:04.835867Z","shell.execute_reply.started":"2024-01-17T03:56:04.812344Z","shell.execute_reply":"2024-01-17T03:56:04.835001Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_batch(batch):\n    fig, axs = plt.subplots(4,4, figsize=(20,20))   \n    images = batch[0].cpu().numpy()\n    labels = batch[1].cpu().numpy()\n\n    for i in range(16):     \n        axs[i%4, i//4].imshow(images[i, 1])\n        axs[i%4, i//4].imshow(labels[i], alpha=0.5)\n    plt.show()\n    \n    plt.hist(batch[0][:,0].cpu().numpy().flatten(), bins=100, alpha=0.5)\n    plt.hist(batch[0][:,1].cpu().numpy().flatten(), bins=100, alpha=0.5)\n    plt.hist(batch[0][:,2].cpu().numpy().flatten(), bins=100, alpha=0.5)\n    plt.show()\n    \n","metadata":{"execution":{"iopub.status.busy":"2024-01-17T03:56:04.837436Z","iopub.execute_input":"2024-01-17T03:56:04.837695Z","iopub.status.idle":"2024-01-17T03:56:04.850773Z","shell.execute_reply.started":"2024-01-17T03:56:04.837670Z","shell.execute_reply":"2024-01-17T03:56:04.850013Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Loss","metadata":{}},{"cell_type":"code","source":"dc = TorchLoss(smp.losses.DiceLoss(mode='multiclass', classes=[1]))\nce = CrossEntropyLossFlat(axis=1) #TorchLoss(smp.losses.SoftCrossEntropyLoss(smooth_factor=0.))","metadata":{"execution":{"iopub.status.busy":"2024-01-17T03:56:04.852141Z","iopub.execute_input":"2024-01-17T03:56:04.852414Z","iopub.status.idle":"2024-01-17T03:56:04.864015Z","shell.execute_reply.started":"2024-01-17T03:56:04.852387Z","shell.execute_reply":"2024-01-17T03:56:04.863268Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Config","metadata":{}},{"cell_type":"code","source":"class CONFIG():\n    # paths\n    path = Path('../input/hubmap-kidney-segmentation')\n    data_path = Path('../input/hubmap-zarr/images_scale2')\n    annotations_path = Path('../input/hubmap-more-efficient-sampling-deepflash2/masks_scale2')\n    \n    # deepflash2 dataset\n    scale = 1.5 # data is already downscaled to 2, so absulute downscale is 3\n    tile_shape = (512, 512)\n    sample_multiplier = 100 # Sample 100 tiles from each image, per epoch\n    val_length = 500 # Randomly sample 500 validation tiles\n    stats = np.array([0.61561477, 0.4379893 , 0.64067212]), np.array([0.2915353 , 0.31549066, 0.28647661])\n        \n    # pytorch model (segmentation_models_pytorch)\n    encoder_name = \"efficientnet-b2\"\n    encoder_weights = 'imagenet'\n    in_channels = 3\n    classes = 2\n    \n    # Training\n    n_splits = 5\n    mixed_precision_training = True\n    batch_size = 16\n    weight_decay = 0.00\n    loss_func = JointLoss(dc, ce, 1, 1)\n    metrics = [Dice(), Iou(), Recall(), Precision()]\n    max_learning_rate = 1e-3\n    epochs = 10\n    \ncfg = CONFIG()\n\n# Albumentations augmentations\ntfms = A.Compose([\n    A.OneOf([\n        A.RandomContrast(),\n        A.RandomGamma(),\n        A.RandomBrightness(),\n        ], p=0.3),\n    A.OneOf([\n        A.Blur(blur_limit=3, p=1),\n        A.MedianBlur(blur_limit=3, p=1)\n    ], p=.1),\n    A.OneOf([\n        A.GaussNoise(0.002, p=.5),\n        A.IAAAffine(p=.5),\n    ], p=.1),\n    # Additional position augmentations\n    A.RandomRotate90(p=.5),\n    A.HorizontalFlip(p=.5),\n    A.VerticalFlip(p=.5),\n    A.Cutout(num_holes=10,fill_value=255, \n             max_h_size=int(.1 * cfg.tile_shape[0]), \n             max_w_size=int(.1 * cfg.tile_shape[0]), \n             p=.1),\n])\n\n# Position Augmentations\nposition_augmentation_kwargs = {\n    'flip':True,                                \n    'rotation_range_deg':(0, 360),     \n    'deformation_grid': (150,150), \n    'deformation_magnitude':(10,10),\n    'value_minimum_range':(0, 0), \n    'value_maximum_range':(1, 1), \n    'value_slope_range':(1, 1)}\n\n# Datasets\nds_kwargs = {\n    'label_path': (cfg.annotations_path/'labels').as_posix(),\n    'cdf_path': (cfg.annotations_path/'cdfs').as_posix(),\n    'df_stats': pd.read_csv(cfg.annotations_path/'roi_stats.csv', index_col=[0,1]),\n    'tile_shape':cfg.tile_shape,\n    'scale': cfg.scale,\n    'val_length':cfg.val_length, \n    'sample_multiplier':cfg.sample_multiplier,\n    'albumentations_tfms': tfms\n}\n","metadata":{"execution":{"iopub.status.busy":"2024-01-17T03:56:04.865200Z","iopub.execute_input":"2024-01-17T03:56:04.865488Z","iopub.status.idle":"2024-01-17T03:56:04.914175Z","shell.execute_reply.started":"2024-01-17T03:56:04.865461Z","shell.execute_reply":"2024-01-17T03:56:04.913448Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train = pd.read_csv(cfg.path/'train.csv')\ndf_info = pd.read_csv(cfg.path/'HuBMAP-20-dataset_information.csv')\nfiles = L([cfg.data_path/x for x in df_train.id])\nfiles","metadata":{"execution":{"iopub.status.busy":"2024-01-17T03:56:04.915205Z","iopub.execute_input":"2024-01-17T03:56:04.915454Z","iopub.status.idle":"2024-01-17T03:56:05.343382Z","shell.execute_reply.started":"2024-01-17T03:56:04.915430Z","shell.execute_reply":"2024-01-17T03:56:05.342594Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training","metadata":{}},{"cell_type":"code","source":"kf = KFold(cfg.n_splits, shuffle=True, random_state=42)\n\nfor i, (train_idx, val_idx) in enumerate(kf.split(files)):\n    files_train, files_val = files[train_idx], files[val_idx]\n    print('Validating on', [x.name for x in files_val])\n    \n    # Datasets\n    train_ds = HubmapRandomTileDataset(files_train, **ds_kwargs, **position_augmentation_kwargs)\n    valid_ds = HubmapValidationDataset(files_val, **ds_kwargs)\n    \n    # Model\n    model = smp.Unet(encoder_name=cfg.encoder_name, \n                     encoder_weights=cfg.encoder_weights, \n                     in_channels=cfg.in_channels, \n                     classes=cfg.classes)\n    \n    # Dataloader and learner\n    dls = DataLoaders.from_dsets(train_ds, valid_ds, bs=cfg.batch_size, after_batch=Normalize.from_stats(*cfg.stats))\n    if torch.cuda.is_available(): dls.cuda(), model.cuda()\n    \n    if i==0: show_batch(dls.one_batch())\n\n    cbs = [SaveModelCallback(monitor='dice'), ElasticDeformCallback]\n    learn = Learner(dls, model, metrics=cfg.metrics, wd=cfg.weight_decay, loss_func=cfg.loss_func, opt_func=ranger, cbs=cbs)\n    if cfg.mixed_precision_training: learn.to_fp16()\n        \n    # Fit\n    learn.fit_one_cycle(cfg.epochs, lr_max=cfg.max_learning_rate)\n    learn.recorder.plot_metrics()\n    \n    # Save Model\n    state = {'model': learn.model.state_dict(), 'stats':cfg.stats}\n    torch.save(state, f'unet_{cfg.encoder_name}_{i}.pth', pickle_protocol=2, _use_new_zipfile_serialization=False)","metadata":{"execution":{"iopub.status.busy":"2024-01-17T03:56:05.344472Z","iopub.execute_input":"2024-01-17T03:56:05.344754Z"},"trusted":true},"execution_count":null,"outputs":[]}]}