{"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":"<center>\n    <h1>[Training] - FastAI Baseline</h1>\n<center>","metadata":{}},{"cell_type":"markdown","source":"<center>\n<img src=\"https://hubmapconsortium.org/wp-content/uploads/2019/01/HuBMAP-Retina-Logo-Color.png\">\n</center>","metadata":{"execution":{"iopub.status.busy":"2022-06-22T21:17:17.163998Z","iopub.execute_input":"2022-06-22T21:17:17.164466Z","iopub.status.idle":"2022-06-22T21:17:17.295969Z","shell.execute_reply.started":"2022-06-22T21:17:17.164418Z","shell.execute_reply":"2022-06-22T21:17:17.292799Z"}}},{"cell_type":"markdown","source":"# Description \n\nWelcome to Human BioMolecular Atlas Program (HuBMAP) + Human Protein Atlas (HPA) competition. \nThe objective of this challenge is segmentation of functional tissue units (FTU. e.g., glomeruli in kidney or alveoli in the lung) in biopsy slides from several different organs. \nThe underlying data includes imagery from different sources prepared with different protocols at a variety of resolutions, reflecting typical challenges for working with medical data.\n\nThis notebook provides a fast.ai starter Pytorch code based on a U-shape network (UneXt50) that was used on multiple competitions in the past and includes several tricks from the previous segmentation competitions.\nIt is [dividing the images into tiles](https://www.kaggle.com/code/thedevastator/converting-to-256x256), selection of tiles with tissue, evaluation of the predictions of multiple models with TTA, combining the tile masks back into image level masks, and conversion into RLE. The [inference](https://www.kaggle.com/code/thedevastator/inference-fastai-baseline) is performed based on models trained in the [fast.ai training notebook](https://www.kaggle.com/code/thedevastator/training-fastai-baseline).\n\n**Inference & Dataset Creation**\n\n- #### Inference Notebook [here](https://www.kaggle.com/code/thedevastator/inference-fastai-baseline). \n- #### Dataset Creation [here](https://www.kaggle.com/code/thedevastator/converting-to-256x256). \n\n**Precomputed Datasets**\n\n- ##### [Dataset (512 x 512)](https://www.kaggle.com/datasets/thedevastator/hubmap-2022-512x512/)\n\n- ##### [Dataset (256 x 256)](https://www.kaggle.com/datasets/thedevastator/hubmap-2022-256x256/)\n\n- ##### [Dataset (128 x 128)](https://www.kaggle.com/datasets/thedevastator/hubmap-2022-128x128/settings)\n\n____\n\n#### Everything is based on the excellent [notebooks](https://www.kaggle.com/code/iafoss/hubmap-pytorch-fast-ai-starter) by [iafoss](https://www.kaggle.com/iafoss) \nAll credit to belongs to the original author!\n____","metadata":{"papermill":{"duration":0.011984,"end_time":"2021-03-11T18:12:50.627613","exception":false,"start_time":"2021-03-11T18:12:50.615629","status":"completed"},"tags":[]}},{"cell_type":"code","source":"\"\"\"\nLovasz-Softmax and Jaccard hinge loss in PyTorch\nMaxim Berman 2018 ESAT-PSI KU Leuven (MIT License)\n\"\"\"\n\nimport torch\nfrom torch.autograd import Variable\nimport torch.nn.functional as F\nimport numpy as np\ntry:\n    from itertools import  ifilterfalse\nexcept ImportError: # py3k\n    from itertools import  filterfalse\n\n\ndef lovasz_grad(gt_sorted):\n    \"\"\"\n    Computes gradient of the Lovasz extension w.r.t sorted errors\n    See Alg. 1 in paper\n    \"\"\"\n    p = len(gt_sorted)\n    gts = gt_sorted.sum()\n    intersection = gts - gt_sorted.float().cumsum(0)\n    union = gts + (1 - gt_sorted).float().cumsum(0)\n    jaccard = 1. - intersection / union\n    if p > 1: # cover 1-pixel case\n        jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]\n    return jaccard\n\n\ndef iou_binary(preds, labels, EMPTY=1., ignore=None, per_image=True):\n    \"\"\"\n    IoU for foreground class\n    binary: 1 foreground, 0 background\n    \"\"\"\n    if not per_image:\n        preds, labels = (preds,), (labels,)\n    ious = []\n    for pred, label in zip(preds, labels):\n        intersection = ((label == 1) & (pred == 1)).sum()\n        union = ((label == 1) | ((pred == 1) & (label != ignore))).sum()\n        if not union:\n            iou = EMPTY\n        else:\n            iou = float(intersection) / union\n        ious.append(iou)\n    iou = f_mean(ious)    # mean accross images if per_image\n    return 100 * iou\n\n\ndef iou(preds, labels, C, EMPTY=1., ignore=None, per_image=False):\n    \"\"\"\n    Array of IoU for each (non ignored) class\n    \"\"\"\n    if not per_image:\n        preds, labels = (preds,), (labels,)\n    ious = []\n    for pred, label in zip(preds, labels):\n        iou = []    \n        for i in range(C):\n            if i != ignore: # The ignored label is sometimes among predicted classes (ENet - CityScapes)\n                intersection = ((label == i) & (pred == i)).sum()\n                union = ((label == i) | ((pred == i) & (label != ignore))).sum()\n                if not union:\n                    iou.append(EMPTY)\n                else:\n                    iou.append(float(intersection) / union)\n        ious.append(iou)\n    ious = map(f_mean, zip(*ious)) # mean accross images if per_image\n    return 100 * np.array(ious)\n\n\n# --------------------------- BINARY LOSSES ---------------------------\n\n\ndef lovasz_hinge(logits, labels, per_image=True, ignore=None):\n    \"\"\"\n    Binary Lovasz hinge loss\n      logits: [B, H, W] Variable, logits at each pixel (between -\\infty and +\\infty)\n      labels: [B, H, W] Tensor, binary ground truth masks (0 or 1)\n      per_image: compute the loss per image instead of per batch\n      ignore: void class id\n    \"\"\"\n    if per_image:\n        loss = f_mean(lovasz_hinge_flat(*flatten_binary_scores(log.unsqueeze(0), lab.unsqueeze(0), ignore))\n                          for log, lab in zip(logits, labels))\n    else:\n        loss = lovasz_hinge_flat(*flatten_binary_scores(logits, labels, ignore))\n    return loss\n\n\ndef lovasz_hinge_flat(logits, labels):\n    \"\"\"\n    Binary Lovasz hinge loss\n      logits: [P] Variable, logits at each prediction (between -\\infty and +\\infty)\n      labels: [P] Tensor, binary ground truth labels (0 or 1)\n      ignore: label to ignore\n    \"\"\"\n    if len(labels) == 0:\n        # only void pixels, the gradients should be 0\n        return logits.sum() * 0.\n    signs = 2. * labels.float() - 1.\n    errors = (1. - logits * Variable(signs))\n    errors_sorted, perm = torch.sort(errors, dim=0, descending=True)\n    perm = perm.data\n    gt_sorted = labels[perm]\n    grad = lovasz_grad(gt_sorted)\n    #loss = torch.dot(F.relu(errors_sorted), Variable(grad))\n    loss = torch.dot(F.elu(errors_sorted)+1, Variable(grad))\n    return loss\n\n\ndef flatten_binary_scores(scores, labels, ignore=None):\n    \"\"\"\n    Flattens predictions in the batch (binary case)\n    Remove labels equal to 'ignore'\n    \"\"\"\n    scores = scores.view(-1)\n    labels = labels.view(-1)\n    if ignore is None:\n        return scores, labels\n    valid = (labels != ignore)\n    vscores = scores[valid]\n    vlabels = labels[valid]\n    return vscores, vlabels\n\n\nclass StableBCELoss(torch.nn.modules.Module):\n    def __init__(self):\n         super(StableBCELoss, self).__init__()\n    def forward(self, input, target):\n         neg_abs = - input.abs()\n         loss = input.clamp(min=0) - input * target + (1 + neg_abs.exp()).log()\n         return loss.mean()\n\n\ndef binary_xloss(logits, labels, ignore=None):\n    \"\"\"\n    Binary Cross entropy loss\n      logits: [B, H, W] Variable, logits at each pixel (between -\\infty and +\\infty)\n      labels: [B, H, W] Tensor, binary ground truth masks (0 or 1)\n      ignore: void class id\n    \"\"\"\n    logits, labels = flatten_binary_scores(logits, labels, ignore)\n    loss = StableBCELoss()(logits, Variable(labels.float()))\n    return loss\n\n\n# --------------------------- MULTICLASS LOSSES ---------------------------\n\n\ndef lovasz_softmax(probas, labels, only_present=False, per_image=False, ignore=None):\n    \"\"\"\n    Multi-class Lovasz-Softmax loss\n      probas: [B, C, H, W] Variable, class probabilities at each prediction (between 0 and 1)\n      labels: [B, H, W] Tensor, ground truth labels (between 0 and C - 1)\n      only_present: average only on classes present in ground truth\n      per_image: compute the loss per image instead of per batch\n      ignore: void class labels\n    \"\"\"\n    if per_image:\n        loss = f_mean(lovasz_softmax_flat(*flatten_probas(prob.unsqueeze(0), lab.unsqueeze(0), ignore), only_present=only_present)\n                          for prob, lab in zip(probas, labels))\n    else:\n        loss = lovasz_softmax_flat(*flatten_probas(probas, labels, ignore), only_present=only_present)\n    return loss\n\n\ndef lovasz_softmax_flat(probas, labels, only_present=False):\n    \"\"\"\n    Multi-class Lovasz-Softmax loss\n      probas: [P, C] Variable, class probabilities at each prediction (between 0 and 1)\n      labels: [P] Tensor, ground truth labels (between 0 and C - 1)\n      only_present: average only on classes present in ground truth\n    \"\"\"\n    C = probas.size(1)\n    losses = []\n    for c in range(C):\n        fg = (labels == c).float() # foreground for class c\n        if only_present and fg.sum() == 0:\n            continue\n        errors = (Variable(fg) - probas[:, c]).abs()\n        errors_sorted, perm = torch.sort(errors, 0, descending=True)\n        perm = perm.data\n        fg_sorted = fg[perm]\n        losses.append(torch.dot(errors_sorted, Variable(lovasz_grad(fg_sorted))))\n    return f_mean(losses)\n\n\ndef flatten_probas(probas, labels, ignore=None):\n    \"\"\"\n    Flattens predictions in the batch\n    \"\"\"\n    B, C, H, W = probas.size()\n    probas = probas.permute(0, 2, 3, 1).contiguous().view(-1, C)  # B * H * W, C = P, C\n    labels = labels.view(-1)\n    if ignore is None:\n        return probas, labels\n    valid = (labels != ignore)\n    vprobas = probas[valid.nonzero().squeeze()]\n    vlabels = labels[valid]\n    return vprobas, vlabels\n\ndef xloss(logits, labels, ignore=None):\n    \"\"\"\n    Cross entropy loss\n    \"\"\"\n    return F.cross_entropy(logits, Variable(labels), ignore_index=255)\n\n\n# --------------------------- HELPER FUNCTIONS ---------------------------\n\ndef f_mean(l, ignore_nan=False, empty=0):\n    \"\"\"\n    nanmean compatible with generators.\n    \"\"\"\n    l = iter(l)\n    if ignore_nan:\n        l = ifilterfalse(np.isnan, l)\n    try:\n        n = 1\n        acc = next(l)\n    except StopIteration:\n        if empty == 'raise':\n            raise ValueError('Empty mean')\n        return empty\n    for n, v in enumerate(l, 2):\n        acc += v\n    if n == 1:\n        return acc\n    return acc / n","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-07-25T03:30:56.560957Z","iopub.execute_input":"2022-07-25T03:30:56.561474Z","iopub.status.idle":"2022-07-25T03:30:58.892052Z","shell.execute_reply.started":"2022-07-25T03:30:56.561361Z","shell.execute_reply":"2022-07-25T03:30:58.891097Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%reload_ext autoreload\n%autoreload 2\n%matplotlib inline\n\nfrom fastai.vision.all import *\nfrom torch.utils.data import Dataset, DataLoader\nimport pandas as pd\nimport numpy as np\nimport os\nimport cv2\nimport gc\nimport random\nfrom albumentations import *\nfrom sklearn.model_selection import KFold\nimport matplotlib.pyplot as plt\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-input":true,"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","papermill":{"duration":3.556569,"end_time":"2021-03-11T18:12:54.195121","exception":false,"start_time":"2021-03-11T18:12:50.638552","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-07-25T03:30:58.894309Z","iopub.execute_input":"2022-07-25T03:30:58.895223Z","iopub.status.idle":"2022-07-25T03:31:01.428025Z","shell.execute_reply.started":"2022-07-25T03:30:58.895167Z","shell.execute_reply":"2022-07-25T03:31:01.427067Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"bs = 64\nnfolds = 4\nfold = 0\nSEED = 2021\nTRAIN = '../input/hubmap2022256x256-v1-xjtlu/train/'\nMASKS = '../input/hubmap2022256x256-v1-xjtlu/mask/'\nLABELS = '../input/hubmap-organ-segmentation/train.csv'\nNUM_WORKERS = 4","metadata":{"papermill":{"duration":0.045568,"end_time":"2021-03-11T18:12:54.252135","exception":false,"start_time":"2021-03-11T18:12:54.206567","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-07-25T03:31:01.429382Z","iopub.execute_input":"2022-07-25T03:31:01.430204Z","iopub.status.idle":"2022-07-25T03:31:01.476429Z","shell.execute_reply.started":"2022-07-25T03:31:01.430145Z","shell.execute_reply":"2022-07-25T03:31:01.475192Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    #the following line gives ~10% speedup\n    #but may lead to some stochasticity in the results \n    torch.backends.cudnn.benchmark = True\n    \nseed_everything(SEED)","metadata":{"_kg_hide-input":true,"papermill":{"duration":0.04608,"end_time":"2021-03-11T18:12:54.308988","exception":false,"start_time":"2021-03-11T18:12:54.262908","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-07-25T03:31:01.479292Z","iopub.execute_input":"2022-07-25T03:31:01.479650Z","iopub.status.idle":"2022-07-25T03:31:01.526589Z","shell.execute_reply.started":"2022-07-25T03:31:01.479613Z","shell.execute_reply":"2022-07-25T03:31:01.525489Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data\nOne important thing here is the train/val split. To avoid possible leaks resulted by a similarity of tiles from the same images, it is better to keep tiles from each image together in train or in test.","metadata":{"papermill":{"duration":0.010988,"end_time":"2021-03-11T18:12:54.330833","exception":false,"start_time":"2021-03-11T18:12:54.319845","status":"completed"},"tags":[]}},{"cell_type":"code","source":"mean_std = np.load('../input/hubmap2022256x256-v1-xjtlu/avr_std.npz')\nmean = mean_std['mean']\nstd = mean_std['std']","metadata":{"execution":{"iopub.status.busy":"2022-07-25T03:31:01.529297Z","iopub.execute_input":"2022-07-25T03:31:01.529914Z","iopub.status.idle":"2022-07-25T03:31:01.591874Z","shell.execute_reply.started":"2022-07-25T03:31:01.529873Z","shell.execute_reply":"2022-07-25T03:31:01.590960Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# https://www.kaggle.com/datasets/thedevastator/hubmap-2022-256x256\n# mean = np.array([0.7720342, 0.74582646, 0.76392896])\n# std = np.array([0.24745085, 0.26182273, 0.25782376])\n\ndef img2tensor(img,dtype:np.dtype=np.float32):\n    if img.ndim==2 : img = np.expand_dims(img,2)\n    img = np.transpose(img,(2,0,1))\n    return torch.from_numpy(img.astype(dtype, copy=False))\n\nclass HuBMAPDataset(Dataset):\n    def __init__(self, fold=fold, train=True, tfms=None):\n        ids = pd.read_csv(LABELS).id.astype(str).values\n        kf = KFold(n_splits=nfolds,random_state=SEED,shuffle=True)\n        ids = set(ids[list(kf.split(ids))[fold][0 if train else 1]])\n        self.fnames = [fname for fname in os.listdir(TRAIN) if fname.split('_')[0] in ids]\n        self.train = train\n        self.tfms = tfms\n        \n    def __len__(self):\n        return len(self.fnames)\n    \n    def __getitem__(self, idx):\n        fname = self.fnames[idx]\n        img = cv2.cvtColor(cv2.imread(os.path.join(TRAIN,fname)), cv2.COLOR_BGR2RGB)\n        mask = cv2.imread(os.path.join(MASKS,fname),cv2.IMREAD_GRAYSCALE)\n#         img = cv2.resize(img, (256, 256))\n#         mask = cv2.resize(mask, (256, 256))\n        if self.tfms is not None:\n            augmented = self.tfms(image=img,mask=mask)\n            img,mask = augmented['image'],augmented['mask']\n        return img2tensor((img/255.0 - mean)/std),img2tensor(mask)\n    \ndef get_aug(p=1.0):\n    return Compose([\n        HorizontalFlip(),\n        VerticalFlip(),\n        RandomRotate90(),\n        Transpose(),\n        ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=15, p=0.9, \n                         border_mode=cv2.BORDER_REFLECT),\n        \n        OneOf([\n            OpticalDistortion(p=0.3),\n            GridDistortion(p=.1),\n            GaussianBlur(p=.3),\n            IAAPiecewiseAffine(p=0.3),\n        ], p=0.3),\n        OneOf([\n            HueSaturationValue(10,15,10),\n            CLAHE(clip_limit=2),\n            RandomBrightnessContrast(),   \n            IAASharpen(),\n        ], p=0.3),\n    ], p=p)","metadata":{"papermill":{"duration":0.057765,"end_time":"2021-03-11T18:12:54.40978","exception":false,"start_time":"2021-03-11T18:12:54.352015","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-07-25T03:31:01.593242Z","iopub.execute_input":"2022-07-25T03:31:01.593595Z","iopub.status.idle":"2022-07-25T03:31:01.648065Z","shell.execute_reply.started":"2022-07-25T03:31:01.593560Z","shell.execute_reply":"2022-07-25T03:31:01.646961Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#example of train images with masks\nds = HuBMAPDataset(tfms=get_aug())\ndl = DataLoader(ds,batch_size=64,shuffle=False,num_workers=NUM_WORKERS)\nimgs,masks = next(iter(dl))\n\nplt.figure(figsize=(16,16))\nfor i,(img,mask) in enumerate(zip(imgs,masks)):\n    img = ((img.permute(1,2,0)*std + mean)*255.0).numpy().astype(np.uint8)\n    plt.subplot(8,8,i+1)\n    plt.imshow(img,vmin=0,vmax=255)\n    plt.imshow(mask.squeeze().numpy(), alpha=0.2)\n    plt.axis('off')\n    plt.subplots_adjust(wspace=None, hspace=None)\n    \ndel ds,dl,imgs,masks","metadata":{"papermill":{"duration":12.580982,"end_time":"2021-03-11T18:13:07.001785","exception":false,"start_time":"2021-03-11T18:12:54.420803","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-07-25T03:31:01.651508Z","iopub.execute_input":"2022-07-25T03:31:01.651821Z","iopub.status.idle":"2022-07-25T03:31:17.933179Z","shell.execute_reply.started":"2022-07-25T03:31:01.651795Z","shell.execute_reply":"2022-07-25T03:31:17.932030Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model\nThe model used in this kernel is based on a U-shape network (UneXt50, see image below), which I used in Severstal and Understanding Clouds competitions. The idea of a U-shape network is coming from a [Unet](https://arxiv.org/pdf/1505.04597.pdf) architecture proposed in 2015 for medical images: the encoder part creates a representation of features at different levels, while the decoder combines the features and generates a prediction as a segmentation mask. The skip connections between encoder and decoder allow us to utilize features from the intermediate conv layers of the encoder effectively, without a need for the information to go the full way through entire encoder and decoder. The latter is especially important to link the predicted mask to the specific pixels of the detected object. Later people realized that ImageNet pretrained computer vision models could drastically improve the quality of a segmentation model because of optimized architecture of the encoder, high encoder capacity (in contrast to one used in the original Unet), and the power of the transfer learning.\n\nThere are several important things that must be added to a Unet network, however, to make it able to reach competitive results with current state of the art approaches. First, it is **Feature Pyramid Network (FPN)**: additional skip connection between different upscaling blocks of the decoder and the output layer. So, the final prediction is produced based on the concatenation of U-net output with resized outputs of the intermediate layers. These skip-connections provide a shortcut for gradient flow improving model performance and convergence speed. Since intermediate layers have many channels, their upscaling and use as an input for the final layer would introduce a significant overhead in terms of the computational time and memory. Therefore, 3x3+3x3 convolutions are applied (factorization) before the resize to reduce the number of channels.\n\nAnother very important thing is the **Atrous Spatial Pyramid Pooling (ASPP) block** added between encoder and decoder. The flaw of the traditional U-shape networks is resulted by a small receptive field. Therefore, if a model needs to make a decision about a segmentation of a large object, especially for a large image resolution, it can get confused being able to look only into parts of the object. A way to increase the receptive field and enable interactions between different parts of the image is use of a block combining convolutions with different dilatations ([Atrous convolutions](https://arxiv.org/pdf/1606.00915.pdf) with various rates in ASPP block). While the original paper uses 6,12,18 rates, they may be customized for a particular task and a particular image resolution to maximize the performance. One more thing I added is using group convolutions in ASPP block to reduce the number of model parameters.\n\nFinally, the decoder upscaling blocks are based on [pixel shuffle](https://arxiv.org/pdf/1609.05158.pdf) rather than transposed convolution used in the first Unet models. It allows to avoid artifacts in the produced masks. And I use [semisupervised Imagenet pretrained ResNeXt50](https://github.com/facebookresearch/semi-supervised-ImageNet1K-models) model as a backbone. In Pytorch it provides the performance of EfficientNet B2-B3 with much faster convergence for the computational cost and GPU RAM requirements of EfficientNet B0 (though, in TF EfficientNet is highly optimized and may be a good thing to use).","metadata":{"papermill":{"duration":0.041669,"end_time":"2021-03-11T18:13:07.084926","exception":false,"start_time":"2021-03-11T18:13:07.043257","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"![](https://i.ibb.co/z5KxDzm/Une-Xt50-1.png)","metadata":{"papermill":{"duration":0.04012,"end_time":"2021-03-11T18:13:07.166112","exception":false,"start_time":"2021-03-11T18:13:07.125992","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class FPN(nn.Module):\n    def __init__(self, input_channels:list, output_channels:list):\n        super().__init__()\n        self.convs = nn.ModuleList(\n            [nn.Sequential(nn.Conv2d(in_ch, out_ch*2, kernel_size=3, padding=1),\n             nn.ReLU(inplace=True), nn.BatchNorm2d(out_ch*2),\n             nn.Conv2d(out_ch*2, out_ch, kernel_size=3, padding=1))\n            for in_ch, out_ch in zip(input_channels, output_channels)])\n        \n    def forward(self, xs:list, last_layer):\n        hcs = [F.interpolate(c(x),scale_factor=2**(len(self.convs)-i),mode='bilinear') \n               for i,(c,x) in enumerate(zip(self.convs, xs))]\n        hcs.append(last_layer)\n        return torch.cat(hcs, dim=1)\n\nclass UnetBlock(Module):\n    def __init__(self, up_in_c:int, x_in_c:int, nf:int=None, blur:bool=False,\n                 self_attention:bool=False, **kwargs):\n        super().__init__()\n        self.shuf = PixelShuffle_ICNR(up_in_c, up_in_c//2, blur=blur, **kwargs)\n        self.bn = nn.BatchNorm2d(x_in_c)\n        ni = up_in_c//2 + x_in_c\n        nf = nf if nf is not None else max(up_in_c//2,32)\n        self.conv1 = ConvLayer(ni, nf, norm_type=None, **kwargs)\n        self.conv2 = ConvLayer(nf, nf, norm_type=None,\n            xtra=SelfAttention(nf) if self_attention else None, **kwargs)\n        self.relu = nn.ReLU(inplace=True)\n\n    def forward(self, up_in:Tensor, left_in:Tensor) -> Tensor:\n        s = left_in\n        up_out = self.shuf(up_in)\n        cat_x = self.relu(torch.cat([up_out, self.bn(s)], dim=1))\n        return self.conv2(self.conv1(cat_x))\n        \nclass _ASPPModule(nn.Module):\n    def __init__(self, inplanes, planes, kernel_size, padding, dilation, groups=1):\n        super().__init__()\n        self.atrous_conv = nn.Conv2d(inplanes, planes, kernel_size=kernel_size,\n                stride=1, padding=padding, dilation=dilation, bias=False, groups=groups)\n        self.bn = nn.BatchNorm2d(planes)\n        self.relu = nn.ReLU()\n\n        self._init_weight()\n\n    def forward(self, x):\n        x = self.atrous_conv(x)\n        x = self.bn(x)\n\n        return self.relu(x)\n\n    def _init_weight(self):\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                torch.nn.init.kaiming_normal_(m.weight)\n            elif isinstance(m, nn.BatchNorm2d):\n                m.weight.data.fill_(1)\n                m.bias.data.zero_()\n\nclass ASPP(nn.Module):\n    def __init__(self, inplanes=512, mid_c=256, dilations=[6, 12, 18, 24], out_c=None):\n        super().__init__()\n        self.aspps = [_ASPPModule(inplanes, mid_c, 1, padding=0, dilation=1)] + \\\n            [_ASPPModule(inplanes, mid_c, 3, padding=d, dilation=d,groups=4) for d in dilations]\n        self.aspps = nn.ModuleList(self.aspps)\n        self.global_pool = nn.Sequential(nn.AdaptiveMaxPool2d((1, 1)),\n                        nn.Conv2d(inplanes, mid_c, 1, stride=1, bias=False),\n                        nn.BatchNorm2d(mid_c), nn.ReLU())\n        out_c = out_c if out_c is not None else mid_c\n        self.out_conv = nn.Sequential(nn.Conv2d(mid_c*(2+len(dilations)), out_c, 1, bias=False),\n                                    nn.BatchNorm2d(out_c), nn.ReLU(inplace=True))\n        self.conv1 = nn.Conv2d(mid_c*(2+len(dilations)), out_c, 1, bias=False)\n        self._init_weight()\n\n    def forward(self, x):\n        x0 = self.global_pool(x)\n        xs = [aspp(x) for aspp in self.aspps]\n        x0 = F.interpolate(x0, size=xs[0].size()[2:], mode='bilinear', align_corners=True)\n        x = torch.cat([x0] + xs, dim=1)\n        return self.out_conv(x)\n    \n    def _init_weight(self):\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                torch.nn.init.kaiming_normal_(m.weight)\n            elif isinstance(m, nn.BatchNorm2d):\n                m.weight.data.fill_(1)\n                m.bias.data.zero_()","metadata":{"_kg_hide-input":true,"papermill":{"duration":0.110724,"end_time":"2021-03-11T18:13:07.317006","exception":false,"start_time":"2021-03-11T18:13:07.206282","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-07-25T03:31:17.934819Z","iopub.execute_input":"2022-07-25T03:31:17.935141Z","iopub.status.idle":"2022-07-25T03:31:18.027536Z","shell.execute_reply.started":"2022-07-25T03:31:17.935110Z","shell.execute_reply":"2022-07-25T03:31:18.026454Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class UneXt50(nn.Module):\n    def __init__(self, stride=1, **kwargs):\n        super().__init__()\n        #encoder\n        m = torch.hub.load('facebookresearch/semi-supervised-ImageNet1K-models',\n                           'resnext50_32x4d_ssl')\n        self.enc0 = nn.Sequential(m.conv1, m.bn1, nn.ReLU(inplace=True))\n        self.enc1 = nn.Sequential(nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1),\n                            m.layer1) #256\n        self.enc2 = m.layer2 #512\n        self.enc3 = m.layer3 #1024\n        self.enc4 = m.layer4 #2048\n        #aspp with customized dilatations\n        self.aspp = ASPP(2048,256,out_c=512,dilations=[stride*1,stride*2,stride*3,stride*4])\n        self.drop_aspp = nn.Dropout2d(0.5)\n        #decoder\n        self.dec4 = UnetBlock(512,1024,256)\n        self.dec3 = UnetBlock(256,512,128)\n        self.dec2 = UnetBlock(128,256,64)\n        self.dec1 = UnetBlock(64,64,32)\n        self.fpn = FPN([512,256,128,64],[16]*4)\n        self.drop = nn.Dropout2d(0.1)\n        self.final_conv = ConvLayer(32+16*4, 1, ks=1, norm_type=None, act_cls=None)\n        \n    def forward(self, x):\n        enc0 = self.enc0(x)\n        enc1 = self.enc1(enc0)\n        enc2 = self.enc2(enc1)\n        enc3 = self.enc3(enc2)\n        enc4 = self.enc4(enc3)\n        enc5 = self.aspp(enc4)\n        dec3 = self.dec4(self.drop_aspp(enc5),enc3)\n        dec2 = self.dec3(dec3,enc2)\n        dec1 = self.dec2(dec2,enc1)\n        dec0 = self.dec1(dec1,enc0)\n        x = self.fpn([enc5, dec3, dec2, dec1], dec0)\n        x = self.final_conv(self.drop(x))\n        x = F.interpolate(x,scale_factor=2,mode='bilinear')\n        return x\n\n#split the model to encoder and decoder for fast.ai\nsplit_layers = lambda m: [list(m.enc0.parameters())+list(m.enc1.parameters())+\n                list(m.enc2.parameters())+list(m.enc3.parameters())+\n                list(m.enc4.parameters()),\n                list(m.aspp.parameters())+list(m.dec4.parameters())+\n                list(m.dec3.parameters())+list(m.dec2.parameters())+\n                list(m.dec1.parameters())+list(m.fpn.parameters())+\n                list(m.final_conv.parameters())]","metadata":{"papermill":{"duration":0.152684,"end_time":"2021-03-11T18:13:07.526419","exception":false,"start_time":"2021-03-11T18:13:07.373735","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-07-25T03:31:18.029267Z","iopub.execute_input":"2022-07-25T03:31:18.029962Z","iopub.status.idle":"2022-07-25T03:31:18.085595Z","shell.execute_reply.started":"2022-07-25T03:31:18.029925Z","shell.execute_reply":"2022-07-25T03:31:18.084592Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Loss and metric\nA famous loss for image segmentation is the [Lovász loss](https://arxiv.org/pdf/1705.08790.pdf), a surrogate of IoU. Following [iafoss](https://www.kaggle.com/iafoss)'s [work](https://www.kaggle.com/code/iafoss/hubmap-pytorch-fast-ai-starter):\n- **ReLU in it must be replaced by (ELU + 1)**(, like he did [here](https://www.kaggle.com/iafoss/lovasz).\n- **Symmetric Lovász loss:** consider not only a predicted segmentation and a provided mask but also the inverse prediction and the inverse mask (predict mask for negative case).","metadata":{"papermill":{"duration":0.076542,"end_time":"2021-03-11T18:13:07.688572","exception":false,"start_time":"2021-03-11T18:13:07.61203","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def symmetric_lovasz(outputs, targets):\n    return 0.5*(lovasz_hinge(outputs, targets) + lovasz_hinge(-outputs, 1.0 - targets))","metadata":{"papermill":{"duration":0.125886,"end_time":"2021-03-11T18:13:07.88963","exception":false,"start_time":"2021-03-11T18:13:07.763744","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-07-25T03:31:18.089970Z","iopub.execute_input":"2022-07-25T03:31:18.090344Z","iopub.status.idle":"2022-07-25T03:31:18.134487Z","shell.execute_reply.started":"2022-07-25T03:31:18.090312Z","shell.execute_reply":"2022-07-25T03:31:18.133506Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Dice_soft(Metric):\n    def __init__(self, axis=1): \n        self.axis = axis \n    def reset(self): self.inter,self.union = 0,0\n    def accumulate(self, learn):\n        pred,targ = flatten_check(torch.sigmoid(learn.pred), learn.y)\n        self.inter += (pred*targ).float().sum().item()\n        self.union += (pred+targ).float().sum().item()\n    @property\n    def value(self): return 2.0 * self.inter/self.union if self.union > 0 else None\n    \n# dice with automatic threshold selection\nclass Dice_th(Metric):\n    def __init__(self, ths=np.arange(0.1,0.9,0.05), axis=1): \n        self.axis = axis\n        self.ths = ths\n        \n    def reset(self): \n        self.inter = torch.zeros(len(self.ths))\n        self.union = torch.zeros(len(self.ths))\n        \n    def accumulate(self, learn):\n        pred,targ = flatten_check(torch.sigmoid(learn.pred), learn.y)\n        for i,th in enumerate(self.ths):\n            p = (pred > th).float()\n            self.inter[i] += (p*targ).float().sum().item()\n            self.union[i] += (p+targ).float().sum().item()\n\n    @property\n    def value(self):\n        dices = torch.where(self.union > 0.0, \n                2.0*self.inter/self.union, torch.zeros_like(self.union))\n        return dices.max()","metadata":{"papermill":{"duration":0.090188,"end_time":"2021-03-11T18:13:08.037967","exception":false,"start_time":"2021-03-11T18:13:07.947779","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-07-25T03:31:18.135974Z","iopub.execute_input":"2022-07-25T03:31:18.136735Z","iopub.status.idle":"2022-07-25T03:31:18.189366Z","shell.execute_reply.started":"2022-07-25T03:31:18.136697Z","shell.execute_reply":"2022-07-25T03:31:18.188354Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model evaluation","metadata":{"papermill":{"duration":0.040333,"end_time":"2021-03-11T18:13:08.119299","exception":false,"start_time":"2021-03-11T18:13:08.078966","status":"completed"},"tags":[]}},{"cell_type":"code","source":"#iterator like wrapper that returns predicted and gt masks\nclass Model_pred:\n    def __init__(self, model, dl, tta:bool=True, half:bool=False):\n        self.model = model\n        self.dl = dl\n        self.tta = tta\n        self.half = half\n        \n    def __iter__(self):\n        self.model.eval()\n        name_list = self.dl.dataset.fnames\n        count=0\n        with torch.no_grad():\n            for x,y in iter(self.dl):\n                x = x.cuda()\n                if self.half: x = x.half()\n                p = self.model(x)\n                py = torch.sigmoid(p).detach()\n                if self.tta:\n                    #x,y,xy flips as TTA\n                    flips = [[-1],[-2],[-2,-1]]\n                    for f in flips:\n                        p = self.model(torch.flip(x,f))\n                        p = torch.flip(p,f)\n                        py += torch.sigmoid(p).detach()\n                    py /= (1+len(flips))\n                if y is not None and len(y.shape)==4 and py.shape != y.shape:\n                    py = F.upsample(py, size=(y.shape[-2],y.shape[-1]), mode=\"bilinear\")\n                py = py.permute(0,2,3,1).float().cpu()\n                batch_size = len(py)\n                for i in range(batch_size):\n                    taget = y[i].detach().cpu() if y is not None else None\n                    yield py[i],taget,name_list[count]\n                    count += 1\n                    \n    def __len__(self):\n        return len(self.dl.dataset)\n    \nclass Dice_th_pred(Metric):\n    def __init__(self, ths=np.arange(0.1,0.9,0.01), axis=1): \n        self.axis = axis\n        self.ths = ths\n        self.reset()\n        \n    def reset(self): \n        self.inter = torch.zeros(len(self.ths))\n        self.union = torch.zeros(len(self.ths))\n        \n    def accumulate(self,p,t):\n        pred,targ = flatten_check(p, t)\n        for i,th in enumerate(self.ths):\n            p = (pred > th).float()\n            self.inter[i] += (p*targ).float().sum().item()\n            self.union[i] += (p+targ).float().sum().item()\n\n    @property\n    def value(self):\n        dices = torch.where(self.union > 0.0, 2.0*self.inter/self.union, \n                            torch.zeros_like(self.union))\n        return dices\n    \ndef save_img(data,name,out):\n    data = data.float().cpu().numpy()\n    img = cv2.imencode('.png',(data*255).astype(np.uint8))[1]\n    out.writestr(name, img)","metadata":{"papermill":{"duration":0.095878,"end_time":"2021-03-11T18:13:08.255835","exception":false,"start_time":"2021-03-11T18:13:08.159957","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-07-25T03:31:18.192767Z","iopub.execute_input":"2022-07-25T03:31:18.193190Z","iopub.status.idle":"2022-07-25T03:31:18.251305Z","shell.execute_reply.started":"2022-07-25T03:31:18.193141Z","shell.execute_reply":"2022-07-25T03:31:18.250382Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train","metadata":{"papermill":{"duration":0.040885,"end_time":"2021-03-11T18:13:08.336994","exception":false,"start_time":"2021-03-11T18:13:08.296109","status":"completed"},"tags":[]}},{"cell_type":"code","source":"dice = Dice_th_pred(np.arange(0.2,0.7,0.01))\npath_pre = [\n    '../input/hubmap-pytorch-fast-ai-starter/model_0.pth',\n    '../input/hubmap-pytorch-fast-ai-starter/model_1.pth',\n    '../input/hubmap-pytorch-fast-ai-starter/model_2.pth',\n    '../input/hubmap-pytorch-fast-ai-starter/model_3.pth'\n]\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nfor fold in range(nfolds):\n    ds_t = HuBMAPDataset(fold=fold, train=True, tfms=get_aug())\n    ds_v = HuBMAPDataset(fold=fold, train=False)\n    data = ImageDataLoaders.from_dsets(ds_t,ds_v,bs=bs,\n                num_workers=NUM_WORKERS,pin_memory=True).cuda()\n    state_dict = torch.load(path_pre[fold], map_location=torch.device('cpu'))\n    model = UneXt50()\n    model.load_state_dict(state_dict)\n    model.to(device)\n    learn = Learner(data, model, loss_func=symmetric_lovasz,\n                metrics=[Dice_soft(),Dice_th()], \n                splitter=split_layers).to_fp16()\n    \n    #start with training the head\n    learn.freeze_to(-1) #doesn't work\n    for param in learn.opt.param_groups[0]['params']:\n        param.requires_grad = False\n    learn.fit_one_cycle(4, lr_max=0.5e-2)\n\n    #continue training full model\n    learn.unfreeze()\n    learn.fit_one_cycle(25, lr_max=slice(2e-4,2e-3),\n        cbs=SaveModelCallback(monitor='dice_th',comp=np.greater))\n    torch.save(learn.model.state_dict(),f'model_{fold}.pth')\n    \n    #model evaluation on val and saving the masks\n    mp = Model_pred(learn.model,learn.dls.loaders[1])\n    with zipfile.ZipFile('val_masks_tta.zip', 'a') as out:\n        for p in progress_bar(mp):\n            dice.accumulate(p[0],p[1])\n            save_img(p[0],p[2],out)\n    gc.collect()","metadata":{"_kg_hide-output":true,"papermill":{"duration":29637.81713,"end_time":"2021-03-12T02:27:06.195179","exception":false,"start_time":"2021-03-11T18:13:08.378049","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-07-25T03:31:18.252626Z","iopub.execute_input":"2022-07-25T03:31:18.252966Z","iopub.status.idle":"2022-07-25T06:40:59.491227Z","shell.execute_reply.started":"2022-07-25T03:31:18.252931Z","shell.execute_reply":"2022-07-25T06:40:59.490197Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dices = dice.value\nnoise_ths = dice.ths\nbest_dice = dices.max()\nbest_thr = noise_ths[dices.argmax()]\nplt.figure(figsize=(8,4))\nplt.plot(noise_ths, dices, color='blue')\nplt.vlines(x=best_thr, ymin=dices.min(), ymax=dices.max(), colors='black')\nd = dices.max() - dices.min()\nplt.text(noise_ths[-1]-0.1, best_dice-0.1*d, f'DICE = {best_dice:.3f}', fontsize=12);\nplt.text(noise_ths[-1]-0.1, best_dice-0.2*d, f'TH = {best_thr:.3f}', fontsize=12);\nplt.show()","metadata":{"papermill":{"duration":0.252836,"end_time":"2021-03-12T02:27:06.509129","exception":false,"start_time":"2021-03-12T02:27:06.256293","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-07-25T06:40:59.492891Z","iopub.execute_input":"2022-07-25T06:40:59.493261Z","iopub.status.idle":"2022-07-25T06:40:59.726526Z","shell.execute_reply.started":"2022-07-25T06:40:59.493224Z","shell.execute_reply":"2022-07-25T06:40:59.725405Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","papermill":{"duration":0.060975,"end_time":"2021-03-12T02:27:06.631927","exception":false,"start_time":"2021-03-12T02:27:06.570952","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]}]}