{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":61446,"databundleVersionId":6962461,"sourceType":"competition"},{"sourceId":1807973,"sourceType":"datasetVersion","datasetId":1074109},{"sourceId":150248402,"sourceType":"kernelVersion"}],"dockerImageVersionId":30588,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Segment vasculature in 3D scans of human kidney [Training Notebook]\nThis was my first try at both computer vision and image segmentation, and I jumped in a bit late (two weeks before the deadline!). Thankfully, awesome guys in the discussion shared some kernels that I learned a lot from (give credit!).\n\nThis notebook used `segmentation_models_pytorch` package to build CNN (convolution neural network) models using `Unet` architecture that build a fully convolution neural network for image semantic segmentation.\n**References**\n- [training 6](https://www.kaggle.com/code/jeanlucvanlite/training-6/notebook)\n- [se_resnext50_32x4d](https://smp.readthedocs.io/en/latest/encoders.html)\n- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://arxiv.org/abs/1505.04597)\n- [2.5d Cutting model baseline [training]](https://www.kaggle.com/code/yoyobar/2-5d-cutting-model-baseline-training) \n- [2.5d segmentaion baseline [training]](https://www.kaggle.com/code/tanakar/2-5d-segmentaion-baseline-training)","metadata":{}},{"cell_type":"markdown","source":"# Offline installation and Copy imagenet-weights \nCopying imagenet-wieghts to local folder to avoid long-awiting downloads","metadata":{}},{"cell_type":"code","source":"!mkdir -p /root/.cache/torch/hub/checkpoints/\n!cp /kaggle/input/se-net-pretrained-imagenet-weights/* /root/.cache/torch/hub/checkpoints/","metadata":{"_kg_hide-output":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!python -m pip install --no-index --find-links=/kaggle/input/pip-download-for-segmentation-models-pytorch segmentation-models-pytorch","metadata":{"_kg_hide-output":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Import","metadata":{}},{"cell_type":"code","source":"import os, sys, cv2\nfrom glob import glob\nimport numpy as np\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\n# PyTorch\nimport torch as tc \nimport torch.nn as nn  \nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.nn.parallel import DataParallel\nfrom torch.cuda.amp import autocast\n# segmentation_models_pytorch\nimport segmentation_models_pytorch as smp\nfrom albumentations.pytorch import ToTensorV2\nimport albumentations as A","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-01-11T20:23:17.294087Z","iopub.execute_input":"2024-01-11T20:23:17.294857Z","iopub.status.idle":"2024-01-11T20:24:00.516912Z","shell.execute_reply.started":"2024-01-11T20:23:17.294799Z","shell.execute_reply":"2024-01-11T20:24:00.515893Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import random, torch\n# Seed the same seed to all \ndef seed_everything(seed=42):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n\nSEED = 42\nseed_everything(SEED)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import ctypes, gc, torch\nlibc = ctypes.CDLL(\"libc.so.6\")\ndef clear_memory():\n    libc.malloc_trim(0)\n    torch.cuda.empty_cache()\n    gc.collect()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# config","metadata":{}},{"cell_type":"code","source":"p_augm = 0.05 #0.5\n#add rotate.  less p_augm\n\nclass CFG:\n    # ============== pred target =============\n    num_classes = 1\n\n    # ============== model CFG =============\n    model_name = 'Unet'\n    backbone = 'se_resnext50_32x4d'\n\n    in_chans = 1   #5 # 65\n    weights=\"imagenet\"\n    # ============== training CFG =============\n    image_size = 512 # 512 # 512\n    input_size = 512 # 512 #=512\n\n    train_batch_size = 1 #4 #16 # Reduce the batch size = 1 to avoid OOM errors\n    valid_batch_size = train_batch_size * 2\n\n    epochs = 30 #27 #30 #25\n    lr = 6e-5\n    chopping_percentile=1e-3\n    # ============== fold =============\n    valid_id = 1\n    # ============== augmentation =============\n    train_aug_list = [\n        A.Rotate(limit=270, p= 0.5),\n        A.RandomScale(scale_limit=(0.8,1.25), interpolation=cv2.INTER_CUBIC, p=p_augm),\n        A.RandomCrop(input_size, input_size, p=1),\n        A.RandomGamma(p=p_augm*2/3),\n        A.RandomBrightnessContrast(p=p_augm,),\n        A.GaussianBlur(p=p_augm),\n        A.MotionBlur(p=p_augm),\n        A.GridDistortion(num_steps=5, distort_limit=0.3, p=p_augm),\n        ToTensorV2(transpose_mask=True),\n    ]\n    train_aug = A.Compose(train_aug_list)\n    valid_aug_list = [\n        ToTensorV2(transpose_mask=True),\n    ]\n    valid_aug = A.Compose(valid_aug_list)","metadata":{"execution":{"iopub.status.busy":"2024-01-11T20:24:00.519017Z","iopub.execute_input":"2024-01-11T20:24:00.519332Z","iopub.status.idle":"2024-01-11T20:24:00.529838Z","shell.execute_reply.started":"2024-01-11T20:24:00.519303Z","shell.execute_reply":"2024-01-11T20:24:00.528767Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"class SegModel(nn.Module):\n    def __init__(self, backbone, in_chans, num_classes, weights):\n        super().__init__()\n        self.model = smp.Unet(  #FPN Unet\n            encoder_name=backbone, \n            encoder_weights=weights,\n            in_channels=in_chans,\n            classes=num_classes,\n            activation=None,\n        )\n\n    def forward(self, image):\n        output = self.model(image)\n        # output = output.squeeze(-1)\n        return output[:,0]#.sigmoid()\n\n\ndef build_model():\n    from dotenv import load_dotenv\n    load_dotenv()\n\n    print('model_name', CFG.model_name)\n    print('backbone', CFG.backbone)\n    \n    model = SegModel(backbone=CFG.backbone, in_chans=CFG.in_chans,\n                     num_classes=CFG.num_classes,  weights=CFG.weights) # Create an model of empty weights\n\n    return model.cuda()","metadata":{"execution":{"iopub.status.busy":"2024-01-11T20:24:00.53115Z","iopub.execute_input":"2024-01-11T20:24:00.531473Z","iopub.status.idle":"2024-01-11T20:24:00.544468Z","shell.execute_reply.started":"2024-01-11T20:24:00.531448Z","shell.execute_reply":"2024-01-11T20:24:00.543548Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Image processing (change image size)","metadata":{}},{"cell_type":"code","source":"def to_1024(img , image_size=1024):\n    if image_size > img.shape[1]:\n        img = np.rot90(img)\n        start1 = (CFG.image_size - img.shape[0])//2 \n        top =     img[0                    : start1,   0: img.shape[1] ]\n        bottom  = img[img.shape[0] -start1 : img.shape[0],   0 : img.shape[1] ]\n        img_result = np.concatenate((top,img,bottom ),axis=0)\n        img_result = np.rot90(img_result)\n        img_result = np.rot90(img_result)\n        img_result = np.rot90(img_result)\n    else :\n        img_result = img\n    return img_result\n\ndef to_1024_no_rot(img, image_size=1024):\n    if image_size > img.shape[0]:  \n        start1 = ( image_size - img.shape[0])//2\n        top =     img[0                    : start1,   0: img.shape[1] ]\n        bottom  = img[img.shape[0] -start1 : img.shape[0],   0 : img.shape[1] ]\n        img_result = np.concatenate((top,img,bottom ),axis=0)\n    else: \n        img_result = img\n    return img_result\n\n#  add border\ndef to_1024_1024(img, image_size=1024):\n    img_result = to_1024(img, image_size)\n    return img_result\n    \n#  drop border\ndef to_original(im_after, img, image_size=1024):\n    top_ = 0\n    left_ = 0\n    if (im_after.shape[0] > img.shape[0]):\n        top_  = (image_size - img.shape[0])//2 \n    if (im_after.shape[1] > img.shape[1]) :\n        left_  = (image_size - img.shape[1])//2  \n    if (top_>0) or (left_>0):\n        img_result = im_after[top_ : img.shape[0] + top_,  left_: img.shape[1] + left_]\n    else:\n        img_result = im_after\n    return img_result  ","metadata":{"execution":{"iopub.status.busy":"2024-01-11T20:24:00.546908Z","iopub.execute_input":"2024-01-11T20:24:00.547193Z","iopub.status.idle":"2024-01-11T20:24:00.561517Z","shell.execute_reply.started":"2024-01-11T20:24:00.547169Z","shell.execute_reply":"2024-01-11T20:24:00.560608Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Utilty Functions","metadata":{}},{"cell_type":"code","source":"def min_max_normalization(x:tc.Tensor)->tc.Tensor:\n    \"\"\"input.shape=(batch,f1,...)\"\"\"\n    shape=x.shape\n    if x.ndim>2:\n        x=x.reshape(x.shape[0],-1)\n    \n    min_=x.min(dim=-1,keepdim=True)[0]\n    max_=x.max(dim=-1,keepdim=True)[0]\n    if min_.mean()==0 and max_.mean()==1:\n        return x.reshape(shape)\n    \n    x=(x-min_)/(max_-min_+1e-9)\n    return x.reshape(shape)\n\ndef norm_with_clip(x:tc.Tensor,smooth=1e-5):\n    dim=list(range(1,x.ndim))\n    mean=x.mean(dim=dim,keepdim=True)\n    std=x.std(dim=dim,keepdim=True)\n    x=(x-mean)/(std+smooth)\n    x[x>5]=(x[x>5]-5)*1e-3 +5\n    x[x<-3]=(x[x<-3]+3)*1e-3-3\n    return x\n\ndef add_noise(x:tc.Tensor,max_randn_rate=0.1,randn_rate=None,x_already_normed=False):\n    \"\"\"input.shape=(batch,f1,f2,...) output's var will be normalizate  \"\"\"\n    ndim=x.ndim-1\n    if x_already_normed:\n        x_std=tc.ones([x.shape[0]]+[1]*ndim,device=x.device,dtype=x.dtype)\n        x_mean=tc.zeros([x.shape[0]]+[1]*ndim,device=x.device,dtype=x.dtype)\n    else: \n        dim=list(range(1,x.ndim))\n        x_std=x.std(dim=dim,keepdim=True)\n        x_mean=x.mean(dim=dim,keepdim=True)\n    if randn_rate is None:\n        randn_rate=max_randn_rate*np.random.rand()*tc.rand(x_mean.shape,device=x.device,dtype=x.dtype)\n    cache=(x_std**2+(x_std*randn_rate)**2)**0.5\n    #https://blog.csdn.net/chaosir1991/article/details/106960408\n    \n    return (x-x_mean+tc.randn(size=x.shape,device=x.device,dtype=x.dtype)*randn_rate*x_std)/(cache+1e-7)\n \n","metadata":{"execution":{"iopub.status.busy":"2024-01-11T20:24:00.563036Z","iopub.execute_input":"2024-01-11T20:24:00.563387Z","iopub.status.idle":"2024-01-11T20:24:00.601994Z","shell.execute_reply.started":"2024-01-11T20:24:00.563355Z","shell.execute_reply":"2024-01-11T20:24:00.600635Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data loader","metadata":{}},{"cell_type":"code","source":"class Data_loader(Dataset):\n     \n    def __init__(self,paths,is_label):\n        self.paths=paths\n        self.paths.sort()\n        self.is_label=is_label\n    \n    def __len__(self):\n        return len(self.paths)\n    \n    def __getitem__(self,index):\n         \n        img = cv2.imread(self.paths[index],cv2.IMREAD_GRAYSCALE)\n        \n        img = to_1024_1024(img , image_size = CFG.image_size ) #  to_original( im_after, img_save, image_size = 1024)\n\n        img = tc.from_numpy(img.copy())\n        if self.is_label:\n            img=(img!=0).to(tc.uint8)*255\n        else:\n            img=img.to(tc.uint8)\n        return img\n\ndef load_data(paths,is_label=False):\n    data_loader=Data_loader(paths,is_label)\n    data_loader=DataLoader(data_loader, batch_size=16, num_workers=2)  \n    data=[]\n    for x in tqdm(data_loader):\n        data.append(x)\n    x=tc.cat(data,dim=0)\n    del data\n    if not is_label:\n        ########################################################################\n        TH=x.reshape(-1).numpy()\n        index = -int(len(TH) * CFG.chopping_percentile)\n        TH:int = np.partition(TH, index)[index]\n        x[x>TH]=int(TH)\n        ########################################################################\n        TH=x.reshape(-1).numpy()\n        index = -int(len(TH) * CFG.chopping_percentile)\n        TH:int = np.partition(TH, -index)[-index]\n        x[x<TH]=int(TH)\n        ########################################################################\n        x=(min_max_normalization(x.to(tc.float16)[None])[0]*255).to(tc.uint8)\n    return x","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Loss Function","metadata":{}},{"cell_type":"code","source":"#https://www.kaggle.com/code/kashiwaba/sennet-hoa-train-unet-simple-baseline\ndef dice_coef(y_pred:tc.Tensor,y_true:tc.Tensor, thr=0.5, dim=(-1,-2), epsilon=0.001):\n    y_pred=y_pred.sigmoid()\n    y_true = y_true.to(tc.float32)\n    y_pred = (y_pred>thr).to(tc.float32)\n    inter = (y_true*y_pred).sum(dim=dim)\n    den = y_true.sum(dim=dim) + y_pred.sum(dim=dim)\n    dice = ((2*inter+epsilon)/(den+epsilon)).mean()\n    return dice\n\nclass DiceLoss(nn.Module):\n    def __init__(self, weight=None, size_average=True):\n        super(DiceLoss, self).__init__()\n\n    def forward(self, inputs, targets, smooth=1):\n        \n        #comment out if your model contains a sigmoid or equivalent activation layer\n        inputs = inputs.sigmoid()   \n        \n        #flatten label and prediction tensors\n        inputs = inputs.view(-1)\n        targets = targets.view(-1)\n        \n        intersection = (inputs * targets).sum()                            \n        dice = (2.*intersection + smooth)/(inputs.sum() + targets.sum() + smooth)  \n        \n        return 1 - dice\n\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# DataSet","metadata":{}},{"cell_type":"code","source":"class Pipeline_Dataset(Dataset):\n    def __init__(self,x:list,y:list,arg=False):\n        super(Dataset,self).__init__()\n        self.x=x#list[(C,H,W),...]\n        self.y=y#list[(C,H,W),...]\n        self.image_size=CFG.image_size\n        self.in_chans=CFG.in_chans\n        self.arg=arg\n        if arg:\n            self.transform=CFG.train_aug\n        else: \n            self.transform=CFG.valid_aug\n\n    def __len__(self) -> int:\n        return sum([y.shape[0]-self.in_chans for y in self.y])\n    \n    def __getitem__(self,index):\n        i=0\n        for x in self.x:\n            if index>x.shape[0]-self.in_chans:\n                index-=x.shape[0]-self.in_chans\n                i+=1\n            else:\n                break\n        x=self.x[i]\n        y=self.y[i]\n        \n        print (f'x.shape[1] ={x.shape[1]}    x.shape[2]={x.shape[2]}')\n        \n        x_index= (x.shape[1]-self.image_size)//2 #np.random.randint(0,x.shape[1]-self.image_size)\n        y_index= (x.shape[2]-self.image_size)//2 # np.random.randint(0,x.shape[2]-self.image_size)\n        # i i+5 \n        x=x[index:index+self.in_chans, \n            x_index:x_index+self.image_size, \n            y_index:y_index+self.image_size]\n        # i+2\n        y=y[index+self.in_chans//2,\n            x_index:x_index+self.image_size,\n            y_index:y_index+self.image_size]\n\n        data = self.transform(image=x.numpy().transpose(1,2,0), mask=y.numpy())\n        x = data['image']\n        y = data['mask']>=127\n        if self.arg:\n            i=np.random.randint(4)\n            x=x.rot90(i,dims=(1,2))\n            y=y.rot90(i,dims=(0,1))\n            for i in range(3):\n                if np.random.randint(2):\n                    x=x.flip(dims=(i,))\n                    if i>=1:\n                        y=y.flip(dims=(i-1,))\n        return x,y#(uint8,uint8)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load data ","metadata":{}},{"cell_type":"code","source":"train_x=[]\ntrain_y=[]\n\nroot_path=\"/kaggle/input/blood-vessel-segmentation/\"\npaths=[\"/kaggle/input/blood-vessel-segmentation/train/kidney_1_dense\"]\nfor i,path in enumerate(paths):\n    if path==\"/kaggle/input/blood-vessel-segmentation/train/kidney_3_dense\":\n        continue\n    x=load_data(glob(f\"{path}/images/*\"),is_label=False)\n    #print(x.shape)\n    y=load_data(glob(f\"{path}/labels/*\"),is_label=True)\n    #print(y.shape)\n    train_x.append(x)\n    train_y.append(y)\n\n    #(C,H,W)\n\n    #aug\n    train_x.append(x.permute(1,2,0))\n    train_y.append(y.permute(1,2,0))\n    train_x.append(x.permute(2,0,1))\n    train_y.append(y.permute(2,0,1))\npath1=\"/kaggle/input/blood-vessel-segmentation/train/kidney_3_sparse\"\npath2=\"/kaggle/input/blood-vessel-segmentation/train/kidney_3_dense\"\npaths_y=glob(f\"{path2}/labels/*\")\npaths_x=[x.replace(\"labels\",\"images\").replace(\"dense\",\"sparse\") for x in paths_y]\n\nval_x=load_data(paths_x,is_label=False)\n#print(val_x.shape)\nval_y=load_data(paths_y,is_label=True)\n#print(val_y.shape)\nclear_memory()","metadata":{"execution":{"iopub.status.busy":"2024-01-11T20:24:00.603522Z","iopub.execute_input":"2024-01-11T20:24:00.603949Z","iopub.status.idle":"2024-01-11T20:27:45.285906Z","shell.execute_reply.started":"2024-01-11T20:24:00.603915Z","shell.execute_reply":"2024-01-11T20:27:45.284648Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"import copy, torch\ntc.backends.cudnn.enabled = True\ntc.backends.cudnn.benchmark = True\n    \ntrain_dataset= Pipeline_Dataset(train_x,train_y,arg=True)\ntrain_dataset = DataLoader(train_dataset, batch_size=CFG.train_batch_size ,num_workers=2, shuffle=True, pin_memory=True)\nval_dataset=Pipeline_Dataset([val_x],[val_y])\nval_dataset = DataLoader(val_dataset, batch_size=CFG.valid_batch_size, num_workers=2, shuffle=False, pin_memory=True)\n\nmodel=build_model()\nmodel=DataParallel(model)\n\nloss_fc=DiceLoss()\n#loss_fn=nn.BCEWithLogitsLoss()\noptimizer=tc.optim.AdamW(model.parameters(),lr=CFG.lr)\nscaler=tc.cuda.amp.GradScaler()\nscheduler = tc.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=CFG.lr,\n                                                steps_per_epoch=len(train_dataset), epochs=CFG.epochs+1,\n                                                pct_start=0.1,)\n\nbest_loss = np.inf\nbest_epoch = -1\nfor epoch in range(CFG.epochs):\n    model.train()\n    time=tqdm(range(len(train_dataset)))\n    losss=0\n    scores=0\n    for i,(x,y) in enumerate(train_dataset):\n        x=x.cuda().to(tc.float32)\n        y=y.cuda().to(tc.float32)\n        x=norm_with_clip(x.reshape(-1,*x.shape[2:])).reshape(x.shape)\n        x=add_noise(x,max_randn_rate=0.5,x_already_normed=True)\n        \n        with autocast():\n            pred=model(x)\n            loss=loss_fc(pred,y)\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n        optimizer.zero_grad()\n        scheduler.step()\n        score=dice_coef(pred.detach(),y)\n        losss=(losss*i+loss.item())/(i+1)\n        scores=(scores*i+score)/(i+1)\n        time.set_description(f\"epoch:{epoch},loss:{losss:.4f},score:{scores:.4f},lr{optimizer.param_groups[0]['lr']:.4e}\")\n        time.update()\n        del loss,pred\n    time.close()\n    \n    model.eval()\n    time=tqdm(range(len(val_dataset)))\n    val_loss=0\n    val_score=0\n    for i,(x,y) in enumerate(val_dataset):\n        x=x.cuda().to(tc.float32)\n        y=y.cuda().to(tc.float32)\n        x=norm_with_clip(x.reshape(-1,*x.shape[2:])).reshape(x.shape)\n\n        with autocast():\n            with tc.no_grad():\n                pred=model(x)\n                loss=loss_fc(pred,y)\n        score=dice_coef(pred.detach(),y)\n        val_loss=(val_loss*i+loss.item())/(i+1)\n        val_score=(val_score*i+score)/(i+1)\n        time.set_description(f\"val-->loss:{val_loss:.4f},score:{val_score:.4f}\")\n        time.update()\n     # Save the model of the best epoch\n    if val_loss <= best_loss:\n        print(f\"Valid loss Improved ({best_loss} ---> {val_loss})\")\n        best_loss = val_loss\n        best_epoch = epoch\n        # Save the model weights\n        best_model_wts = copy.deepcopy(model.state_dict())\n        PATH = \"best_epoch_1024x1024.bin\"\n        torch.save(model.state_dict(), PATH)\n        print(f\"Best epoch model is saved to {PATH}\")\n\n    # Save the model of last epoch\n    last_model_wts = copy.deepcopy(model.state_dict())\n    PATH = \"last_epoch_1024x1024.bin\"\n    torch.save(model.state_dict(), PATH)\n    print(\"Best Loss: {:.4f}\".format(best_loss))\n    time.close()\n    clear_memory()\n       \ntime.close()","metadata":{"execution":{"iopub.status.busy":"2024-01-11T20:27:45.288132Z","iopub.execute_input":"2024-01-11T20:27:45.28857Z"},"trusted":true},"execution_count":null,"outputs":[]}]}