{"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":"code","source":"# # Setup for offline notebook run\n!cp -r ../input/pytorch-segmentation-models-lib/ ./\n!pip config set global.disable-pip-version-check true\n!pip install -q ./pytorch-segmentation-models-lib/timm-0.4.12-py3-none-any.whl\n!cp -r ../input/einops-041-wheel/ ./\n!pip config set global.disable-pip-version-check true\n!pip install -q ../input/einops-041-wheel/einops-0.4.1-py3-none-any.whl\n!pip install /kaggle/input/staintools-offline/spams-2.6.5.4-cp37-cp37m-linux_x86_64.whl","metadata":{"_uuid":"b87d619a-1ef7-4e74-bb31-4219321f241e","_cell_guid":"78143a57-e409-4aff-9aa4-c81334fff55e","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-13T14:41:08.562651Z","iopub.execute_input":"2022-10-13T14:41:08.563576Z","iopub.status.idle":"2022-10-13T14:41:44.445715Z","shell.execute_reply.started":"2022-10-13T14:41:08.563456Z","shell.execute_reply":"2022-10-13T14:41:44.444541Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nsys.path.append('../input/hubmap-modular')\nsys.path.append('../input/cnn-vision')","metadata":{"execution":{"iopub.status.busy":"2022-10-13T14:41:44.448077Z","iopub.execute_input":"2022-10-13T14:41:44.449769Z","iopub.status.idle":"2022-10-13T14:41:44.455216Z","shell.execute_reply.started":"2022-10-13T14:41:44.449726Z","shell.execute_reply":"2022-10-13T14:41:44.454129Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install segmentation-models-pytorch","metadata":{"execution":{"iopub.status.busy":"2022-10-13T14:41:44.456924Z","iopub.execute_input":"2022-10-13T14:41:44.457329Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"## -*- coding: utf-8 -*-\n\n# ## ** Kaggle Setup for offline notebook run **\n# !cp -r ../input/pytorch-segmentation-models-lib/ ./\n# !pip config set global.disable-pip-version-check true\n# !pip install -q ./pytorch-segmentation-models-lib/timm-0.4.12-py3-none-any.whl\n# !cp -r ../input/einops-041-wheel/ ./\n# !pip config set global.disable-pip-version-check true\n# !pip install -q ../input/einops-041-wheel/einops-0.4.1-py3-none-any.whl\n# !pip install /kaggle/input/staintools-offline/spams-2.6.5.4-cp37-cp37m-linux_x86_64.whl\n\n# ## ** Imports **\nimport numpy as np\nimport pandas as pd\nimport os\nimport matplotlib.pyplot as plt\nfrom skimage import io,img_as_float\nimport cv2\nfrom sklearn.model_selection import StratifiedKFold, KFold, StratifiedGroupKFold\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport random\nimport time\nimport copy\nfrom collections import defaultdict\nimport gc\nfrom PIL import Image\nimport imagecodecs\nimport tifffile \nimport gc\n# from tqdm.notebook import tqdm # for notebook runs\nfrom tqdm import tqdm \nfrom pathlib import Path\nimport glob\nimport copy\nfrom timeit import default_timer as timer\nimport timm\nimport sys\n#SMP\n# import segmentation_models_pytorch as smp\n\n# HUGGINGFACE \nimport transformers \n\n# SWA\nfrom torch.optim.swa_utils import AveragedModel, SWALR\n\n\n# PYTORCH\nimport torch\nfrom torch.utils.data import Dataset,DataLoader\nfrom torch.optim import lr_scheduler\nimport torch.nn as nn\nfrom torchmetrics.functional import dice\nimport torch.optim as optim\nfrom torchvision.utils import make_grid\nfrom torchvision.io import read_image\nimport torchvision.transforms.functional as F\nfrom torch.nn import functional as nnF\nfrom torch.cuda import amp\n\nimport warnings\nfrom colorama import Fore, Back, Style\nfrom tabulate import tabulate\nimport pandas as pd\n\nc_  = Fore.GREEN\nlo_  = Fore.RED\nyel_ = Fore.BLUE\nsr_ = Style.RESET_ALL\n\nwarnings.filterwarnings(\"ignore\")\nplt.rcParams[\"savefig.bbox\"] = 'tight'\n\n# ## **Import local modules**\n\nfrom mit import *\nfrom pvt_v2 import *\nfrom coat import *\nfrom daformer import *\n\n# from efficientnet import *\nfrom vahadane_stain_normalization import stain_normalize\n\nprint('import ok\\n')\n\n# ## **Train config**\n\nclass CFG:\n    def __init__(self,n_fold = 5,seed = 42,batch_size = 4,debug=False,kaggle = True):\n        \n        # Directories\n        self.kaggle     = kaggle\n        self.root_dir   = \"../input/hubmap-organ-segmentation\" if kaggle else './data'\n        self.IMAGES     = '../input/hubmap-stain-norm-original/train_images_stain_normalised' if kaggle else './data/images_stain_normalized'\n        self.MASKS      = '../input/hubmap-hpa-2022-png-dataset/train_masks_png' if kaggle else './data/masks'\n        \n        self.n_fold        = n_fold\n        self.seed          = seed\n        # step2: data\n        self.batch_size    = batch_size\n        self.train_bs      = batch_size\n        self.valid_bs      = batch_size*2\n        self.debug         = debug\n        self.img_size      = [768, 768] \n        self.num_workers   = 0\n        \n        # step3: model  \n        \n        # ##  a) CoaT\n#         self.model_parameters_list = dotdict(module='model_coat_daformer',param={'encoder': coat_parallel_small_plus1(), 'decoder': daformer_conv1x1,},\n#                                             checkpoint='../input/coat-pretrained-models/coat_small_7479cf9b.pth' if kaggle else './pretrained_weights/coat_small_7479cf9b.pth')\n        #  ## b) Pvt_v2\n#         self.model_parameters_list = dotdict(module='model_pvt_v2_daformer',param={'encoder': pvt_v2_b3(), 'decoder': daformer_conv3x3, 'encoder_cfg': {'img_size': self.img_size[0]}},\n#                                               checkpoint='../input/coat-pretrained-models/pvt_v2_b3.pth' if kaggle else './pretrained_weights/pvt_v2_b3.pth')\n                \n        #  ## c) Segformer\n#         self.model_parameters_list = dotdict(module='model_mit_segformer',param={'encoder': mit_b2(), 'decoder': daformer_conv1x1,},\n#                                              checkpoint='../input/coat-pretrained-models/mit_b2.pth' if kaggle else './pretrained_weights/mit_b2.pth')\n        \n        # ##  a) CNN MODELS\n        # Normal\n#         self.model_parameters_list = dotdict(module='model_deeplabv3',encoder_name='efficientnet-b5') # resnet101 or resnext101_32x8d or efficientnet-b5\n\n        # Continue Aux_Shift Training  # resnet101 or resnext101_32x8d or efficientnet-b5\n        self.model_parameters_list = dotdict(module='model_deeplabv3',encoder_name='resnext101_32x8d',\n                                     checkpoint = '../input/cnn-vision-results/deeplab/resnext/with_aux/model_resnext_deeplabv3_fold_0_768-best_epoch-49.bin')\n       \n        self.device        = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")        \n#         self.backbone      = self.model_parameters_list.param['encoder'].__class__.__name__\n        self.model_name    = self.model_parameters_list.module\n        self.exp_name      = f'Exp__{self.model_name}__Training'\n        \n        # step4: optimizer\n        self.epochs        = 80\n        self.lr            = 3e-5\n        self.lr_drop       = 8\n        self.optimizer     = 'AdamW'\n        self.weight_decay  = 1e-5\n        \n        # step5: scheduler\n        self.scheduler     = \"ReduceLROnPlateau\" # ReduceLROnPlateau OR CosineAnnealingLR\n        self.min_lr        = 1e-6 # 0.00005\n        self.T_max         = self.epochs + 50 #int(280/self.batch_size*self.epochs)+50\n#         self.T_0           = 25\n        self.warmup_epochs = 0\n        self.wd            = 1e-5\n        self.n_accumulate  = 1 # max(1, 32//self.batch_size)       \n       \n\n        \n    def display(self):\n        print(f\"Experiment     {self.exp_name}\")\n        print(f\"Debug is       {self.debug}\") \n        print(f\"Device is      {self.device}\")\n#         print(f\"Model is       {self.model_name}_{self.backbone}\")\n        print(f'Total folds    {self.n_fold}')\n        print(f'-----------------------------')\n        print(f\"Batch size is  {self.batch_size}\")\n        print(f\"img_size is    {self.img_size}\")        \n        print(f\"epochs is      {self.epochs}\")\n        print(f\"Optimiser is   {self.optimizer}\")\n        print(f\"Scheduler is   {self.scheduler}\")\n        print(f\"LR is          {self.lr}\")\n        print(f'-----------------------------')\n\n\ndef set_seed(seed = 42):\n    '''Sets the seed of the entire notebook so results are the same every time we run.\n    This is for REPRODUCIBILITY.'''\n    print(f\"Setting seed as {seed}\")\n    np.random.seed(seed)\n    random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    # When running on the CuDNN backend, two further options must be set\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    # Set a fixed value for the hash seed\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    print('> SEEDING DONE')\n\n\ndef initialize_config(debug=False,batch_size=4):\n    cfg = CFG(batch_size = batch_size,debug = debug)    \n    set_seed(cfg.seed)\n    cfg.display()\n    return cfg\n\n# ## **Model**\n\nimport importlib\n\ndef build_model(model_parameters):    \n    \n    M = importlib.import_module(model_parameters.module)\n#     print(M)\n#     model = M.Net()\n#     model.output_type = [\"loss\",\"inference\"]\n#     model.is_train = True\n#     encoder    =  model_parameters.param['encoder']\n#     path       =  model_parameters.checkpoint\n#     checkpoint = torch.load(path, map_location=lambda storage, loc: storage)\n#     state_dict = checkpoint#['model']\n#     encoder.load_state_dict(state_dict,strict=False)\n#     model.encoder = encoder\n#     model.drop_path_rate=0.1\n#     print('Loaded pretrained weights: ',encoder.pretrain)\n    model = M.DeepLabV3Plus(encoder_name = model_parameters.encoder_name)\n    #-----------\n    model.load_state_dict(torch.load(model_parameters.checkpoint, map_location=lambda storage, loc: storage))\n    print('Loaded for continue train!')\n    #-----------\n    model.output_type = [\"loss\",\"inference\"]\n#     print(f'Weights loadded {model_parameters.encoder_name}')\n    return model\n\n\ndef load_model(model_parameters):\n    M = importlib.import_module(model_parameters.module)\n    model = M.Net(**model_parameters.param)   \n    model.load_state_dict(torch.load(model_parameters.checkpoint, map_location=lambda storage, loc: storage)) \n    model.output_type = [\"inference\"]\n    model.is_train = False\n    model.dropout = nn.Dropout(p = 0.0)\n    print('Loaded model: ',model_parameters.checkpoint)\n    model = model.eval()\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    model = model.to(device)\n    \n    return model\n\n\n# ## **Mix-Upsampling**\n\nimport torch.nn.functional as Fx\n\nclass MixUpSample_2(nn.Module):\n    '''\n    Mixed upsampling combining bilinear and nearest interpolation.\n    Convert the model prediction back to the original image size.\n    '''\n    def __init__( self, scale_factor=4, originalshape = True):\n        super().__init__()\n        self.mixing = nn.Parameter(torch.tensor(0.5))\n        self.scale_factor = scale_factor\n        self.originalshape = originalshape\n        \n    def forward(self, x, originalshape):\n        \n        if self.originalshape:\n\n            x = self.mixing *Fx.interpolate(x, size = originalshape,mode='bilinear', align_corners=False) \\\n                + (1-self.mixing )*Fx.interpolate(x, size = originalshape, mode='nearest')\n\n            return x        \n\n# ## **Util functions**\n\ndef create_folds(cfg = None):\n    \n    train_csv = pd.read_csv(os.path.join(cfg.root_dir,\"train.csv\"))\n    train_csv.set_index('id',inplace = True)\n    \n    images_path = os.listdir(cfg.IMAGES)\n    images_path = sorted(images_path)\n    masks_path =  os.listdir(cfg.MASKS)\n    masks_path = sorted(masks_path)\n    \n    ids  = [filename[:-4] for filename in masks_path]\n\n    organs      = [train_csv['organ'][int(idx)] for idx in ids]\n    img_height  = [train_csv['img_height'][int(idx)] for idx in ids]\n    img_width   = [train_csv['img_width'][int(idx)] for idx in ids]\n    data_source = [train_csv['data_source'][int(idx)] for idx in ids] \n    rle_mask    = [train_csv['rle'][int(idx)] for idx in ids] \n    images_path = [os.path.join(cfg.IMAGES,f) for f in images_path]\n    masks_path  = [os.path.join(cfg.MASKS,f) for f in masks_path] \n    \n\n    maping = {\n        'id': ids,\n        'organ':organs,\n        'image_path':images_path,\n        'mask_path': masks_path,\n        'data_source': data_source,\n#         'rle': rle_mask,\n        'img_height': img_height,\n        'img_width': img_width,\n\n    }\n    df = pd.DataFrame.from_dict(maping)    \n    skf = KFold(n_splits=cfg.n_fold, shuffle=True,random_state=cfg.seed)\n\n    df.loc[:,'fold']=-1\n    for f,(t_idx, v_idx) in enumerate(skf.split(X=df['id'], y=df['organ'])):\n        df.iloc[v_idx,-1]=f    \n\n    return df\n\ndef prepare_loaders(fold,df,cfg, debug=False):\n    \n    train_df = df.query(\"fold!=@fold\").reset_index(drop=True)\n    valid_df = df.query(\"fold==@fold\").reset_index(drop=True)\n    \n    if debug==True:\n        train_df = train_df.head(4*5)\n        train_df =  train_df.iloc[:50, :] # CHECKING TIME\n        valid_df = valid_df.head(4*5)\n    \n    \n    train_dataset = HuBMAPDataset(train_df, transforms=True,cfg=cfg,mode = 'train') \n    valid_dataset = HuBMAPDataset(valid_df, transforms=False,cfg=cfg, mode = 'train') \n    \n    train_loader = DataLoader(train_dataset, batch_size=cfg.train_bs if not cfg.debug else 20, \n                            num_workers=0, shuffle=True, pin_memory=True, drop_last=False)\n    valid_loader = DataLoader(valid_dataset, batch_size=cfg.valid_bs if not cfg.debug else 20, \n                            num_workers=0, shuffle=False, pin_memory=True)\n    \n    return train_loader, valid_loader\n\n\ndef get_optimizer(cfg,optimizer_name= 'Adam'):\n    if optimizer_name == 'Adam':\n        optimizer = optim.Adam(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)\n    \n    elif optimizer_name == 'AdamW':\n        optimizer = optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()),lr=cfg.lr)\n#         optimizer = optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)\n        \n    return optimizer\n\n\ndef get_scheduler(cfg,optimizer,df):\n    \n    if cfg.scheduler == 'CosineAnnealingLR':\n        scheduler = lr_scheduler.CosineAnnealingLR(optimizer,T_max=cfg.T_max,eta_min=cfg.min_lr)\n        \n    elif cfg.scheduler == 'ReduceLROnPlateau':\n        scheduler = lr_scheduler.ReduceLROnPlateau( optimizer  = optimizer, verbose=True, factor=0.7,mode=\"min\",patience=5, threshold=0.001,min_lr =cfg.min_lr)\n        \n    elif cfg.scheduler == 'StepLR':\n        scheduler = lr_scheduler.StepLR(optimizer,cfg.lr_drop ,gamma=0.1)\n        \n    else:\n        scheduler = None\n        \n    return scheduler\n\n# ## ** Augmentations **\n\ndef do_random_flip(image, mask):\n    if np.random.rand()>0.5:\n        image = cv2.flip(image,0)\n        mask = cv2.flip(mask,0)\n    if np.random.rand()>0.5:\n        image = cv2.flip(image,1)\n        mask = cv2.flip(mask,1)\n    if np.random.rand()>0.5:\n        image = image.transpose(1,0,2)\n        mask = mask.transpose(1,0)\n    \n    image = np.ascontiguousarray(image)\n    mask = np.ascontiguousarray(mask)\n    return image, mask\n\ndef do_random_rot90(image, mask):\n    r = np.random.choice([\n        0,\n        cv2.ROTATE_90_CLOCKWISE,\n        cv2.ROTATE_90_COUNTERCLOCKWISE,\n        cv2.ROTATE_180,\n    ])\n    if r==0:\n        return image, mask\n    else:\n        image = cv2.rotate(image, r)\n        mask = cv2.rotate(mask, r)\n        return image, mask\n    \ndef do_random_contast(image, mask, mag=0.3):\n    alpha = 1 + random.uniform(-1,1)*mag\n    image = image * alpha\n    image = np.clip(image,0,1)\n    return image, mask\n\ndef do_random_hsv(image, mask, mag=[0.15,0.25,0.25]):\n    image = (image*255).astype(np.uint8)\n    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)\n\n    h = hsv[:, :, 0].astype(np.float32)  # hue\n    s = hsv[:, :, 1].astype(np.float32)  # saturation\n    v = hsv[:, :, 2].astype(np.float32)  # value\n    h = (h*(1 + random.uniform(-1,1)*mag[0]))%180\n    s =  s*(1 + random.uniform(-1,1)*mag[1])\n    v =  v*(1 + random.uniform(-1,1)*mag[2])\n\n    hsv[:, :, 0] = np.clip(h,0,180).astype(np.uint8)\n    hsv[:, :, 1] = np.clip(s,0,255).astype(np.uint8)\n    hsv[:, :, 2] = np.clip(v,0,255).astype(np.uint8)\n    image = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)\n    image = image.astype(np.float32)/255\n    return image, mask\n\ndef do_random_noise(image, mask, mag=0.1):\n    height, width = image.shape[:2]\n    noise = np.random.uniform(-1,1, (height, width,1))*mag\n    image = image + noise\n    image = np.clip(image,0,1)\n    return image, mask\n\ndef do_random_rotate_scale(image, mask, angle=30, scale=[0.8,1.2] ):\n    angle = np.random.uniform(-angle, angle)\n    scale = np.random.uniform(*scale) if scale is not None else 1\n    \n    height, width = image.shape[:2]\n    center = (height // 2, width // 2)\n    \n    transform = cv2.getRotationMatrix2D(center, angle, scale)\n    image = cv2.warpAffine( image, transform, (width, height), flags=cv2.INTER_LINEAR,\n                            borderMode=cv2.BORDER_CONSTANT, borderValue=(0,0,0))\n    mask  = cv2.warpAffine( mask, transform, (width, height), flags=cv2.INTER_LINEAR,\n                            borderMode=cv2.BORDER_CONSTANT, borderValue=0)\n    return image, mask\n\n#----------------------Below combined aug functions--------------------\n\ndef valid_augment5(image, mask):\n  \n    return image, mask\n\ndef train_augment5b(image, mask):\n    \n    more_transform = A.Compose([        \n         A.OneOf([ \n             A.ElasticTransform(p=1.0, alpha=image.shape[1]*3, sigma=image.shape[1] * 0.07, alpha_affine=image.shape[1] * 0.09),\n             A.GridDistortion(num_steps=5, distort_limit=0.05, p=1.0),\n              ], p=0.25), \n#          A.CoarseDropout(max_holes=8, max_height=image.shape[0]//20, max_width=image.shape[1]//20,\n#                          min_holes=5, fill_value=0, mask_fill_value=0, p=0.25),\n        ], p=1.0)\n\n    \n    \n    image, mask = do_random_flip(image, mask)\n    image, mask = do_random_rot90(image, mask)      \n    for fn in np.random.choice([\n        lambda image, mask: (image, mask),\n        lambda image, mask: do_random_noise(image, mask, mag=0.1),\n        lambda image, mask: do_random_contast(image, mask, mag=0.40),\n        lambda image, mask: do_random_hsv(image, mask, mag=[0.40, 0.40, 0]) # Remove and try\n    ], 2): image, mask = fn(image, mask)\n\n    for fn in np.random.choice([\n        lambda image, mask: (image, mask),\n        lambda image, mask: do_random_rotate_scale(image, mask, angle=45, scale=[0.50, 2.0]),# scale=[0.50, 2.0] or 0.7,1.5 or 0.50,1.5\n    ], 1): image, mask = fn(image, mask)\n\n#     augmented = more_transform(image= (image*255).astype(np.uint8), mask=mask)\n#     image = augmented['image'].astype(np.float32)/255\n#     mask = augmented['mask']\n        \n    return image, mask\n\n\n# ## ** Utils 2 **\n\nclass dotdict(dict):\n\t__setattr__ = dict.__setitem__\n\t__delattr__ = dict.__delitem__\n\t\n\tdef __getattr__(self, name):\n\t\ttry:\n\t\t\treturn self[name]\n\t\texcept KeyError:\n\t\t\traise AttributeError(name)\n\n\ndef read_tiff(image_file, mode='rgb'):\n    # !pip install imagecodecs if facing COMPRESSION.lZW error\n\timage = tifffile.imread(image_file)\n\timage = image.squeeze()\n\tif image.shape[0] == 3:\n\t\timage = image.transpose(1, 2, 0)\n        \n\tif mode=='bgr':\n\t\timage = image[:,:,::-1]\n    \n\tmx = np.max(image)\n\timage = image.astype(np.float32)\n\tif mx:\n\t\timage /= mx # scale image to [0, 1]\n        \n\timage = np.ascontiguousarray(image)\n\treturn image\n\n\n# Ref: https://www.kaggle.com/bguberfain/memory-aware-rle-encoding (with transposed mask)\n\ndef rle_encode_less_memory(img):\n    # the image should be transposed\n    pixels = img.T.flatten()\n    \n    # This simplified method requires first and last pixel to be zero\n    pixels[0] = 0\n    pixels[-1] = 0\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 2\n    runs[1::2] -= runs[::2]\n    \n    return ' '.join(str(x) for x in runs)\n\ndef rle_encode(img):\n    '''\n    img: numpy array, 1 - mask, 0 - background\n    Returns run length as string formated\n    '''\n    pixels = img.T.flatten()\n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    return ' '.join(str(x) for x in runs)\n\n# Ref: https://www.kaggle.com/code/paulorzp/rle-functions-run-lenght-encode-decode/script\ndef rle_decode(mask_rle, shape):\n    '''\n    mask_rle: run-length as string formated (start length)\n    shape: (height,width) of array to return \n    Returns numpy array, 1 - mask, 0 - background\n    '''\n    s = mask_rle.split()\n    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n    starts -= 1\n    ends = starts + lengths\n    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)\n    for lo, hi in zip(starts, ends):\n        img[lo:hi] = 1\n    return img.reshape(shape).T  # Needed to align to RLE direction\n\ndef read_json_as_list(json_file):\n\twith open(json_file) as f:\n\t\tj = json.load(f)\n\treturn j\n\n\n# ## ** Dataset **\nclass HuBMAPDataset(Dataset):\n    def __init__(self, df= None, cfg = None, transforms = None, mode='train'):\n        self.df = df\n        self.mode = mode\n        self.cfg = cfg\n        self.transforms = transforms\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, index):\n        \n        try: \n            img_path   = self.df.loc[index, 'image_path'] \n            organ      = self.df.loc[index, 'organ']\n            mask_path  = self.df.loc[index,'mask_path'] \n            \n            # Read image\n            image      = read_tiff(img_path)\n            image_size = cfg.img_size[0]\n            \n            if self.mode=='train':\n                \n                # Read mask\n                mask = cv2.imread(mask_path,cv2.IMREAD_GRAYSCALE)\n                mask  = (mask/255).astype(np.uint8)\n                \n                # Resize\n                image = cv2.resize(image,dsize=(image_size,image_size),interpolation=cv2.INTER_LINEAR)\n                mask  = cv2.resize(mask, dsize=(image_size,image_size),interpolation=cv2.INTER_LINEAR)\n                \n                # Augemnt\n                if self.transforms:                   \n                    image, mask = train_augment5b(image, mask)\n                else:\n                    image, mask = valid_augment5(image, mask)\n               \n                # -------------- Sanity check------------------------\n#                 plt.imshow(image)\n                # print('Data class read image,mask: img,mask',image.shape,mask.shape,np.unique(mask))\n                # ---------------------------------------------------\n                image = np.transpose(image, (2, 0, 1)) # image.shape = (3,H,W)\n                mask = np.expand_dims(mask, axis=0)    # mask.shape = (1,H,W)\n                \n                image = image.astype(np.float32)\n                mask  = mask.astype(np.float32)\n\n                return torch.tensor(image), torch.tensor(mask) , organ\n\n            else:    # resize for infer\n                if self.transforms:\n                    data = self.transforms(image=img)\n                    img  = data['image']\n\n                img = np.transpose(img, (2, 0, 1))   #(c, h, w)      \n                return torch.tensor(img), img_height, img_width,id_,organs,sours\n            \n        except:\n            print(index,\"error occured in dataloader/dataset\")\n            return None\n\n\n# ## ** Loss **\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 = F.sigmoid(inputs)       \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# ## ** Metrics **\ndef dice_coef(y_true, y_pred, thr=0.5, dim=(2,3), epsilon=1.):\n    y_true = y_true.to(torch.float32)\n    y_pred = (y_pred>thr).to(torch.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(dim=(1,0))\n    return dice\n\ndef iou_coef(y_true, y_pred, thr=0.5, dim=(2,3), epsilon=0.001):\n    y_true = y_true.to(torch.float32)\n    y_pred = (y_pred>thr).to(torch.float32)\n    inter = (y_true*y_pred).sum(dim=dim)\n    union = (y_true + y_pred - y_true*y_pred).sum(dim=dim)\n    iou = ((inter+epsilon)/(union+epsilon)).mean(dim=(1,0))\n    return iou\n\n# ## ** Train epoch **\n\n'''\nFLOW OF TRAINING:\n\n1. Load batch of data (add to gpu if accessible)\n2. forward pass\n   |_ output = model (input)\n   |_ loss   = criterion (output,target)\n   \n3. backward pass + optimize (only in training)\n   |_ loss.backward()\n   |_ optimizer.step()\n   \n4. Calculate statistics\n   |_ running_loss += loss.item()*batch_size\n\nreturn epoch_loss = running_loss/len(dataloader)\n'''\n\ndef train_one_epoch(cfg,model, optimizer, scheduler, dataloader, device, epoch):\n    model.train()\n    scaler = amp.GradScaler()\n    \n    dataset_size = 0\n    running_loss = 0.0\n    running_dice_loss = 0.0\n    criterion = DiceLoss() \n\n    pbar = tqdm(enumerate(dataloader), total=len(dataloader), desc='Train ')\n    for step, (images, masks, organs) in pbar:  \n#         loss     = 0\n#         aux_loss = 0\n#         bce_loss = 0\n        \n        # get a batch of inputs\n        images = images.to(device, dtype=torch.float)\n        masks  = masks.to(device, dtype=torch.float)\n        batch_size = images.size(0)\n        \n        send_batch = {\n            'image': images,\n            'mask': masks,\n            'organ': organs,\n        }\n        \n        # ------------- forward ---------------\n        with amp.autocast(enabled=True):\n            \n            output    = model(send_batch)   \n            y_pred    = output['probability']\n            loss      = output['bce_loss'].mean()\n            loss1     = output['aux2_loss'].mean()\n            loss2     = output['aux1_loss'].mean()\n            dice_loss = criterion(y_pred,masks)\n\n        scaler.scale(loss+0.2*loss2).backward() #/cfg.n_accumulate loss+0.2*loss1\n    \n        if ((step+1)%cfg.n_accumulate==0 or (step+1)==len(dataloader)):\n            scaler.step(optimizer)\n            scaler.update()\n\n            # zero the parameter gradients\n            optimizer.zero_grad()\n            \n            if cfg.scheduler == 'CosineAnnealingLR' and cfg.scheduler: # 'CosineAnnealingLR' or 'StepLR' location\n                scheduler.step()\n       \n        # statistics\n        running_loss += (loss.item() * batch_size) \n        running_dice_loss += (dice_loss.item()*batch_size)\n        dataset_size += batch_size\n        \n        epoch_loss = running_loss / dataset_size\n        epoch_dice_loss = running_dice_loss / dataset_size\n        \n        mem = torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0\n        current_lr = optimizer.param_groups[0]['lr']\n        pbar.set_postfix(train_BCE_loss=f'{epoch_loss:0.4f}',dice_loss=f'{epoch_dice_loss:0.4f}', lr=f'{current_lr:0.6f}',gpu_mem=f'{mem:0.2f} GB')\n    torch.cuda.empty_cache()\n    gc.collect()\n    \n    return epoch_loss\n\n# ## ** Val epoch **\n\n@torch.no_grad()\ndef valid_one_epoch(model, dataloader, device, epoch):\n    model.eval()\n    \n    dataset_size = 0\n    running_loss = 0.0\n    running_dice_loss = 0.0\n    \n    criterion = DiceLoss()    \n    val_scores = []\n\n    pbar = tqdm(enumerate(dataloader), total=len(dataloader), desc='Valid ')\n    for step, (images, masks, organs) in pbar:       \n        \n        images  = images.to(device, dtype=torch.float)\n        masks   = masks.to(device, dtype=torch.float)\n\n        batch_size = images.size(0)\n        \n        send_batch = {\n            'image': images,\n            'mask': masks,\n            'organ': organs,\n        }\n        output   = model(send_batch)   \n        y_pred = output['probability']      \n\n        loss = output['bce_loss'].mean()\n        dice_loss = criterion(y_pred,masks)\n\n        \n        running_loss += (loss.item() * batch_size)\n        running_dice_loss += (dice_loss.item()*batch_size)\n        dataset_size += batch_size\n        \n        epoch_loss = running_loss / dataset_size\n        epoch_dice_loss = running_dice_loss / dataset_size\n\n        val_dice = dice_coef(masks, y_pred).cpu().detach().numpy()\n        val_jaccard = iou_coef(masks, y_pred).cpu().detach().numpy()\n        val_scores.append([val_dice, val_jaccard])\n        \n        mem = torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0\n        current_lr = optimizer.param_groups[0]['lr']\n        pbar.set_postfix(valid_BCE_loss=f'{epoch_loss:0.4f}',dice_loss=f'{epoch_dice_loss:0.4f}',lr=f'{current_lr:0.6f}',gpu_memory=f'{mem:0.2f} GB')\n    \n    val_scores  = np.mean(val_scores, axis=0)\n    torch.cuda.empty_cache()\n    gc.collect()\n    \n    return epoch_loss, val_scores\n\n# ** Training **\n\nimport copy\ndef run_training(cfg,fold, model, optimizer, scheduler, device, num_epochs):\n    # To automatically log gradients\n    \n    if torch.cuda.is_available():\n        print(\"cuda: {}\\n\".format(torch.cuda.get_device_name()))\n    \n    start = time.time()\n    best_model_wts = copy.deepcopy(model.state_dict())\n    best_dice      = -np.inf\n    best_epoch     = -1\n    history = defaultdict(list)\n  \n    for epoch in range(50, num_epochs + 1): \n        gc.collect()\n        print(f'Epoch {epoch}/{num_epochs}', end='')\n        train_loss = train_one_epoch(cfg,model, optimizer, scheduler,dataloader=train_loader,device=cfg.device, epoch=epoch)        \n        val_loss, val_scores = valid_one_epoch(model, valid_loader,device=cfg.device, epoch=epoch)        \n        val_dice, val_jaccard = val_scores\n        \n        #ReduceLR scheduler\n        if cfg.scheduler =='ReduceLROnPlateau': \n#             print('not none scheduler...')\n            scheduler.step(val_loss) # Monitor metric\n            \n        \n        history['Fold'].append(fold)\n        history['Epoch'].append(epoch)\n        history['Train Loss'].append(train_loss)\n        history['Valid Loss'].append(val_loss)\n        history['Valid Dice'].append(val_dice)\n        history['Valid Jaccard'].append(val_jaccard)\n        \n        pd.DataFrame(history).to_csv(f'./{cfg.model_name}_fold_{fold}_{cfg.img_size[0]}.csv', index=False)\n        \n        print(f'Valid Dice: {val_dice:0.4f} | Valid Jaccard: {val_jaccard:0.4f}')\n        print(f'{lo_}Valid Loss: {val_loss:0.4f} | Train Loss: {train_loss:0.4f}{sr_}')\n        # deep copy the model\n        if val_dice >= best_dice:\n            \n            print(f\"{c_}Valid Dice Score Improved ({best_dice:0.4f} ---> {val_dice:0.4f}){sr_}\")\n            best_dice    = val_dice\n            best_jaccard = val_jaccard\n            best_epoch   = epoch\n            \n            \n            best_model_wts = copy.deepcopy(model.state_dict())\n            PATH = f\"{cfg.model_name}_fold_{fold}_{cfg.img_size[0]}-best_epoch-{best_epoch}.bin\"          \n            \n            torch.save(model.state_dict(), PATH)\n            # Save a model file from the current directory\n            print(f\"Model Saved{sr_}\")\n        else:    \n            last_model_wts = copy.deepcopy(model.state_dict())\n            PATH = f\"{cfg.model_name}_fold_{fold}_{cfg.img_size[0]}-last_epoch.bin\" \n            torch.save(model.state_dict(), PATH)\n            \n        print(); print(f'{sr_}')\n\n    \n    end = time.time()\n    time_elapsed = end - start\n    print('Training complete in {:.0f}h {:.0f}m {:.0f}s'.format(\n        time_elapsed // 3600, (time_elapsed % 3600) // 60, (time_elapsed % 3600) % 60))\n    print(\"Best Score: {:.4f}\".format(best_dice))\n    \n    # load best model weights\n    model.load_state_dict(best_model_wts)\n    \n    return model, history\n\ndef check_data_loader(cfg =None, fold = 0, df = None):\n    \n    train_loader, valid_loader = prepare_loaders(fold=fold,df = df,cfg=cfg)\n    print(f'\\n---DataLoader info:---')\n    for i in range(4):\n        batch = next(iter(train_loader))\n        _ = random.randint(0, cfg.batch_size-1)\n        images,labels,organs = batch\n        image = images[_]\n        label = labels[_].squeeze()\n        organ = organs[_]\n        \n        plt.figure(figsize=(8,8))\n        plt.subplot(1, 2, 1)\n        plt.imshow((image.permute(1,2,0)), interpolation='none'); plt.title(f'Image ({organ})')\n        plt.subplot(1, 2, 2)\n        plt.imshow(image.permute(1,2,0), 'gray', interpolation='none');plt.title('Overlay')\n        plt.imshow(label, 'jet', interpolation='none', alpha=0.7)\n        plt.show()\n \n    del images,labels, batch,organs\n    torch.cuda.empty_cache()\n    print(f'---DataLoader Ok!---\\n')\n\n# ## ** MAIN **\n  \ncfg                        = initialize_config(debug=False,batch_size=4)\ndf                         = create_folds(cfg=cfg)\n# check_data_loader(cfg = cfg, df = df)\n\nfor fold in range(cfg.n_fold):\n    print(f'#'*15*2)\n    print(f'### Fold: {fold}')\n    print(f'#'*15*2)\n    \n    train_loader, valid_loader = prepare_loaders(fold=fold,df = df,cfg=cfg, debug = False)\n    model                      = build_model(cfg.model_parameters_list).to(cfg.device) \n    optimizer                  = get_optimizer(cfg,optimizer_name=cfg.optimizer)\n    scheduler                  = get_scheduler(cfg,optimizer,df)\n    model, history             = run_training(cfg,fold, model, optimizer, scheduler,device=cfg.device,num_epochs=cfg.epochs)\n    \n    #Plot 1\n    #plt.subplot(3, 1, 1)\n    plt.plot(history['Epoch'], history['Train Loss'], 'r--')\n    plt.plot(history['Epoch'], history['Valid Loss'], 'b-')\n    plt.legend(['Training Loss', 'Valid Loss'])\n    plt.xlabel('Epoch')\n    plt.ylabel('Loss')\n    plt.title(f'{cfg.model_name}_fold_{fold}_Epoch_vs_Loss')\n    plt.savefig(f'./{cfg.model_name}_fold_{fold}_Epoch_vs_Loss.png')\n    plt.show()\n    \n    #Plot 2\n    #plt.subplot(3,1, 2)\n    plt.plot(history['Epoch'], history['Valid Dice'])\n    plt.xlabel('Epoch')\n    plt.ylabel('Valid Dice')\n    plt.title(f'{cfg.model_name}_fold_{fold}_Epoch_vs_Valid-Dice')\n    plt.savefig(f'./{cfg.model_name}_fold_{fold}_Epoch_vs_Valid-Dice.png')\n    plt.show()\n    \n    #Plot 3\n    #plt.subplot(3,1, 3)\n    plt.plot(history['Valid Dice'], history['Train Loss'] )\n    plt.xlabel('Valid Dice') \n    plt.ylabel('Train Loss')\n    plt.title(f'{cfg.model_name}_fold_{fold}_Valid-Dice_vs_Loss')\n    plt.savefig(f'./{cfg.model_name}_fold_{fold}_Valid-Dice_vs_Loss.png')\n    plt.show()\n    \n    break\n\n\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}