{"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":"markdown","source":"# HuBMAP - Efficient Sampling Ensemble (deepflash2, pytorch, fastai) [train]\n\n> Kernel for model training with efficient region based sampling.\n\n***\n\n## 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/cptac-pda-to-zarr converted images (downscaled with factor 2)\n- https://www.kaggle.com/matjes/cptac-pda-pancreas-efficient-sampling-deepflash2 masks and weights for sampling\n","metadata":{"id":"OcsetTMwKXqC"}},{"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,"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":{"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":{"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        for key, grp in self.df_stats.groupby('idx'):\n            for (idx, i), row in grp.iterrows():\n                for _ in range(5):\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=False):\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":{"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":{"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":{"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":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Config","metadata":{}},{"cell_type":"code","source":"class CONFIG():\n    # paths\n    path = Path('../input/cptacpda')\n    data_path = Path('../input/cptac-pda-to-zarr/images_scale2')\n    annotations_path = Path('../input/cptac-pda-pancreas-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 = 50 # Sample 100 tiles from each image, per epoch\n    val_length = None # Randomly sample 500 validation tiles\n    # Using 2*std\n    stats = np.array([0.5727, 0.4154, 0.6088]), np.array([0.1694, 0.1797, 0.1316])*2 \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 = 3\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.RandomGamma(p=.25),\n    A.RandomBrightnessContrast(p=.25),\n    A.ColorJitter(brightness=0, contrast=0, saturation=0.2, hue=.1, p=0.25),\n    A.OneOf([\n        A.Blur(blur_limit=3, p=1),\n        A.MedianBlur(blur_limit=3, p=1),\n        A.GaussNoise(0.002, p=1),\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":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_info = pd.read_csv(cfg.path/\"dataset_information.csv\", index_col='image_file')\nfiles = L([cfg.data_path/x for x in df_info.index])\nfiles","metadata":{"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":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Calculating Stats\n\nOnly perfomed once to get the dataset stats.","metadata":{}},{"cell_type":"code","source":"CALC_STATS = False\nif CALC_STATS:\n    ds = HubmapRandomTileDataset(files, **ds_kwargs)\n    dls = DataLoaders.from_dsets(ds, bs=cfg.batch_size)\n\n    mean_sum, var_sum = 0., 0.\n    i = 0\n    for b in dls.train:\n        for img in b[0]:\n            mean_sum += torch.mean(img, (1,2)).cpu()\n            var_sum += torch.var(img, (1,2)).cpu()\n            i += 1\n    mean = mean_sum/i\n    std = torch.sqrt(var_sum/i)\n    print(mean, std)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}