ESFPNet_Endoscope_Generalization_Capacity_and_Power_Balance.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bab7d951",
   "metadata": {},
   "outputs": [],
   "source": [
    "# loading in and transforming data\n",
    "import os\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset,DataLoader,ConcatDataset\n",
    "from torch.autograd import Variable\n",
    "\n",
    "import torchvision\n",
    "import torchvision.datasets as datasets\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "#from skimage import io, transform\n",
    "from PIL import Image\n",
    "\n",
    "# visualizing data\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import warnings\n",
    "\n",
    "# load dataset information\n",
    "import yaml\n",
    "\n",
    "# image writing\n",
    "import imageio\n",
    "from skimage import img_as_ubyte\n",
    "\n",
    "# Clear GPU cache\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e19c5c3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# configuration\n",
    "datasetTrain = ['Kvasir','CVC-ClinicDB']\n",
    "datasetTest = ['Kvasir', 'CVC-ColonDB', 'CVC-ClinicDB', 'ETIS-LaribPolypDB', 'CVC-300']\n",
    "\n",
    "_model_name = 'ESFP_B0_Endo'\n",
    "config = open('Configure.yaml')\n",
    "config = yaml.safe_load(config)\n",
    "model_type = 'B0'\n",
    "\n",
    "repeats = 2\n",
    "n_epochs = 50\n",
    "init_trainsize = 352\n",
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f7b2d99",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PolypDataset(Dataset):\n",
    "    \"\"\"\n",
    "    dataloader for polyp segmentation tasks\n",
    "    \"\"\"\n",
    "    def __init__(self, image_root, gt_root, trainsize, augmentations):\n",
    "        self.trainsize = trainsize\n",
    "        self.augmentations = augmentations\n",
    "        print(self.augmentations)\n",
    "        self.images = [image_root + f for f in os.listdir(image_root) if f.endswith('.jpg') or f.endswith('.png')]\n",
    "        #print(image_root)\n",
    "        self.gts = [gt_root + f for f in os.listdir(gt_root) if f.endswith('.png') or f.endswith('.jpg')]\n",
    "        self.images = sorted(self.images)\n",
    "        self.gts = sorted(self.gts)\n",
    "        self.filter_files()\n",
    "        self.size = len(self.images)\n",
    "        if self.augmentations == True:\n",
    "            print('Using RandomRotation, RandomFlip')\n",
    "            self.img_transform = transforms.Compose([\n",
    "                transforms.RandomRotation(90, resample=False, expand=False, center=None),\n",
    "                transforms.RandomVerticalFlip(p=0.5),\n",
    "                transforms.RandomHorizontalFlip(p=0.5),\n",
    "                transforms.Resize((self.trainsize, self.trainsize)),\n",
    "                transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0, hue=0),\n",
    "                transforms.ToTensor(),\n",
    "                transforms.Normalize([0.485, 0.456, 0.406],\n",
    "                                     [0.229, 0.224, 0.225])])\n",
    "            self.gt_transform = transforms.Compose([\n",
    "                transforms.RandomRotation(90, resample=False, expand=False, center=None),\n",
    "                transforms.RandomVerticalFlip(p=0.5),\n",
    "                transforms.RandomHorizontalFlip(p=0.5),\n",
    "                transforms.Resize((self.trainsize, self.trainsize)),\n",
    "                transforms.ToTensor()])\n",
    "            \n",
    "        else:\n",
    "            print('no augmentation')\n",
    "            self.img_transform = transforms.Compose([\n",
    "                transforms.Resize((self.trainsize, self.trainsize)),\n",
    "                transforms.ToTensor(),\n",
    "                transforms.Normalize([0.485, 0.456, 0.406],\n",
    "                                     [0.229, 0.224, 0.225])])\n",
    "            \n",
    "            self.gt_transform = transforms.Compose([\n",
    "                transforms.Resize((self.trainsize, self.trainsize)),\n",
    "                transforms.ToTensor()])\n",
    "            \n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \n",
    "        image = self.rgb_loader(self.images[index])\n",
    "        gt = self.binary_loader(self.gts[index])\n",
    "        \n",
    "        seed = np.random.randint(2147483647) # make a seed with numpy generator \n",
    "        np.random.seed(seed) # apply this seed to img tranfsorms\n",
    "        torch.manual_seed(seed) # needed for torchvision 0.7\n",
    "        if self.img_transform is not None:\n",
    "            image = self.img_transform(image)\n",
    "            \n",
    "        np.random.seed(seed) # apply this seed to img tranfsorms\n",
    "        torch.manual_seed(seed) # needed for torchvision 0.7\n",
    "        if self.gt_transform is not None:\n",
    "            gt = self.gt_transform(gt)\n",
    "        return image, gt\n",
    "\n",
    "    def filter_files(self):\n",
    "        assert len(self.images) == len(self.gts)\n",
    "        images = []\n",
    "        gts = []\n",
    "        for img_path, gt_path in zip(self.images, self.gts):\n",
    "            img = Image.open(img_path)\n",
    "            gt = Image.open(gt_path)\n",
    "            if img.size == gt.size:\n",
    "                images.append(img_path)\n",
    "                gts.append(gt_path)\n",
    "        self.images = images\n",
    "        self.gts = gts\n",
    "\n",
    "    def rgb_loader(self, path):\n",
    "        with open(path, 'rb') as f:\n",
    "            img = Image.open(f)\n",
    "            return img.convert('RGB')\n",
    "\n",
    "    def binary_loader(self, path):\n",
    "        with open(path, 'rb') as f:\n",
    "            img = Image.open(f)\n",
    "            # return img.convert('1')\n",
    "            return img.convert('L')\n",
    "\n",
    "    def resize(self, img, gt):\n",
    "        assert img.size == gt.size\n",
    "        w, h = img.size\n",
    "        if h < self.trainsize or w < self.trainsize:\n",
    "            h = max(h, self.trainsize)\n",
    "            w = max(w, self.trainsize)\n",
    "            return img.resize((w, h), Image.BILINEAR), gt.resize((w, h), Image.NEAREST)\n",
    "        else:\n",
    "            return img, gt\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.size\n",
    "\n",
    "class test_dataset:\n",
    "    def __init__(self, image_root, gt_root, testsize):\n",
    "        self.testsize = testsize\n",
    "        self.images = [image_root + f for f in os.listdir(image_root) if f.endswith('.jpg') or f.endswith('.png')]\n",
    "        self.gts = [gt_root + f for f in os.listdir(gt_root) if f.endswith('.tif') or f.endswith('.png') or f.endswith('.jpg')]\n",
    "        self.images = sorted(self.images)\n",
    "        self.gts = sorted(self.gts)\n",
    "        self.transform = transforms.Compose([\n",
    "            transforms.Resize((self.testsize, self.testsize)),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize([0.485, 0.456, 0.406],\n",
    "                                 [0.229, 0.224, 0.225])])\n",
    "        self.gt_transform = transforms.ToTensor()\n",
    "        self.size = len(self.images)\n",
    "        self.index = 0\n",
    "\n",
    "    def load_data(self):\n",
    "        image = self.rgb_loader(self.images[self.index])\n",
    "        image = self.transform(image).unsqueeze(0)\n",
    "        gt = self.binary_loader(self.gts[self.index])\n",
    "        name = self.images[self.index].split('/')[-1]\n",
    "        if name.endswith('.jpg'):\n",
    "            name = name.split('.jpg')[0] + '.png'\n",
    "        self.index += 1\n",
    "        return image, gt, name\n",
    "\n",
    "    def rgb_loader(self, path):\n",
    "        with open(path, 'rb') as f:\n",
    "            img = Image.open(f)\n",
    "            return img.convert('RGB')\n",
    "\n",
    "    def binary_loader(self, path):\n",
    "        with open(path, 'rb') as f:\n",
    "            img = Image.open(f)\n",
    "            return img.convert('L')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7aedd79b",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainDatasetList = []\n",
    "for train_id in datasetTrain:\n",
    "    images_train_path = config['dataset']['train_' + str(train_id) + '_dataset'] + '/images/'\n",
    "    masks_train_path = config['dataset']['train_' + str(train_id) + '_dataset'] + '/masks/'\n",
    "    trainDataset = PolypDataset(images_train_path, masks_train_path, trainsize=init_trainsize, augmentations = True)\n",
    "    trainDatasetList.append(trainDataset)\n",
    "trainWholeDataset = ConcatDataset([trainDatasetList[0], trainDatasetList[1]])\n",
    "train_loader = DataLoader(dataset=trainDataset,batch_size=batch_size,shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13f7a8ce",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# imshow function\n",
    "def imshow(img):\n",
    "    npimg = img.numpy()\n",
    "    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
    "    \n",
    "# set rows and columns\n",
    "rows = 1\n",
    "columns = 16\n",
    "\n",
    "dataiter = iter(train_loader)\n",
    "# the \"_\" is a placeholder for no labels\n",
    "images, masks= dataiter.next()\n",
    "\n",
    "# show images\n",
    "fig = plt.figure(figsize=(17,5))\n",
    "for i in range(1,17):\n",
    "    fig.add_subplot(rows,columns,i)\n",
    "    imshow(torchvision.utils.make_grid(images[i-1,:,:,:]))\n",
    "    \n",
    "# show images\n",
    "fig = plt.figure(figsize=(17,5))\n",
    "for i in range(1,17):\n",
    "    fig.add_subplot(rows,columns,i)\n",
    "    imshow(torchvision.utils.make_grid(masks[i-1,:,:,:]))\n",
    "    \n",
    "fig = plt.figure()   \n",
    "#plt.hist(masks.reshape(720*720*5).detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc993709",
   "metadata": {},
   "outputs": [],
   "source": [
    "from Encoder import mit\n",
    "from Decoder import mlp\n",
    "from mmcv.cnn import ConvModule\n",
    "\n",
    "class ESFPNetStructure(nn.Module):\n",
    "\n",
    "    def __init__(self, embedding_dim = 160):\n",
    "        super(ESFPNetStructure, self).__init__()\n",
    "        \n",
    "        # Backbone\n",
    "        if model_type == 'B0':\n",
    "            self.backbone = mit.mit_b0()\n",
    "        if model_type == 'B1':\n",
    "            self.backbone = mit.mit_b1()\n",
    "        if model_type == 'B2':\n",
    "            self.backbone = mit.mit_b2()\n",
    "        if model_type == 'B3':\n",
    "            self.backbone = mit.mit_b3()\n",
    "        if model_type == 'B4':\n",
    "            self.backbone = mit.mit_b4()\n",
    "        if model_type == 'B5':\n",
    "            self.backbone = mit.mit_b5()\n",
    "        \n",
    "        self._init_weights()  # load pretrain\n",
    "        \n",
    "        # LP Header\n",
    "        self.LP_1 = mlp.LP(input_dim = self.backbone.embed_dims[0], embed_dim = self.backbone.embed_dims[0])\n",
    "        self.LP_2 = mlp.LP(input_dim = self.backbone.embed_dims[1], embed_dim = self.backbone.embed_dims[1])\n",
    "        self.LP_3 = mlp.LP(input_dim = self.backbone.embed_dims[2], embed_dim = self.backbone.embed_dims[2])\n",
    "        self.LP_4 = mlp.LP(input_dim = self.backbone.embed_dims[3], embed_dim = self.backbone.embed_dims[3])\n",
    "        \n",
    "        # Linear Fuse\n",
    "        self.linear_fuse34 = ConvModule(in_channels=(self.backbone.embed_dims[2] + self.backbone.embed_dims[3]), out_channels=self.backbone.embed_dims[2], kernel_size=1,norm_cfg=dict(type='BN', requires_grad=True))\n",
    "        self.linear_fuse23 = ConvModule(in_channels=(self.backbone.embed_dims[1] + self.backbone.embed_dims[2]), out_channels=self.backbone.embed_dims[1], kernel_size=1,norm_cfg=dict(type='BN', requires_grad=True))\n",
    "        self.linear_fuse12 = ConvModule(in_channels=(self.backbone.embed_dims[0] + self.backbone.embed_dims[1]), out_channels=self.backbone.embed_dims[0], kernel_size=1,norm_cfg=dict(type='BN', requires_grad=True))\n",
    "        \n",
    "        # Fused LP Header\n",
    "        self.LP_12 = mlp.LP(input_dim = self.backbone.embed_dims[0], embed_dim = self.backbone.embed_dims[0])\n",
    "        self.LP_23 = mlp.LP(input_dim = self.backbone.embed_dims[1], embed_dim = self.backbone.embed_dims[1])\n",
    "        self.LP_34 = mlp.LP(input_dim = self.backbone.embed_dims[2], embed_dim = self.backbone.embed_dims[2])\n",
    "        \n",
    "        # Final Linear Prediction\n",
    "        self.linear_pred = nn.Conv2d((self.backbone.embed_dims[0] + self.backbone.embed_dims[1] + self.backbone.embed_dims[2] + self.backbone.embed_dims[3]), 1, kernel_size=1)\n",
    "        \n",
    "    def _init_weights(self):\n",
    "        \n",
    "        if model_type == 'B0':\n",
    "            pretrained_dict = torch.load('./Pretrained/mit_b0.pth')\n",
    "        if model_type == 'B1':\n",
    "            pretrained_dict = torch.load('./Pretrained/mit_b1.pth')\n",
    "        if model_type == 'B2':\n",
    "            pretrained_dict = torch.load('./Pretrained/mit_b2.pth')\n",
    "        if model_type == 'B3':\n",
    "            pretrained_dict = torch.load('./Pretrained/mit_b3.pth')\n",
    "        if model_type == 'B4':\n",
    "            pretrained_dict = torch.load('./Pretrained/mit_b4.pth')\n",
    "        if model_type == 'B5':\n",
    "            pretrained_dict = torch.load('./Pretrained/mit_b5.pth')\n",
    "            \n",
    "            \n",
    "        model_dict = self.backbone.state_dict()\n",
    "        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}\n",
    "        model_dict.update(pretrained_dict)\n",
    "        self.backbone.load_state_dict(model_dict)\n",
    "        print(\"successfully loaded!!!!\")\n",
    "        \n",
    "        \n",
    "    def forward(self, x):\n",
    "        \n",
    "        ##################  Go through backbone ###################\n",
    "        \n",
    "        B = x.shape[0]\n",
    "        \n",
    "        #stage 1\n",
    "        out_1, H, W = self.backbone.patch_embed1(x)\n",
    "        for i, blk in enumerate(self.backbone.block1):\n",
    "            out_1 = blk(out_1, H, W)\n",
    "        out_1 = self.backbone.norm1(out_1)\n",
    "        out_1 = out_1.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()  #(Batch_Size, self.backbone.embed_dims[0], 88, 88)\n",
    "        \n",
    "        # stage 2\n",
    "        out_2, H, W = self.backbone.patch_embed2(out_1)\n",
    "        for i, blk in enumerate(self.backbone.block2):\n",
    "            out_2 = blk(out_2, H, W)\n",
    "        out_2 = self.backbone.norm2(out_2)\n",
    "        out_2 = out_2.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()  #(Batch_Size, self.backbone.embed_dims[1], 44, 44)\n",
    "        \n",
    "        # stage 3\n",
    "        out_3, H, W = self.backbone.patch_embed3(out_2)\n",
    "        for i, blk in enumerate(self.backbone.block3):\n",
    "            out_3 = blk(out_3, H, W)\n",
    "        out_3 = self.backbone.norm3(out_3)\n",
    "        out_3 = out_3.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()  #(Batch_Size, self.backbone.embed_dims[2], 22, 22)\n",
    "        \n",
    "        # stage 4\n",
    "        out_4, H, W = self.backbone.patch_embed4(out_3)\n",
    "        for i, blk in enumerate(self.backbone.block4):\n",
    "            out_4 = blk(out_4, H, W)\n",
    "        out_4 = self.backbone.norm4(out_4)\n",
    "        out_4 = out_4.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()  #(Batch_Size, self.backbone.embed_dims[3], 11, 11)\n",
    "        \n",
    "        # go through LP Header\n",
    "        lp_1 = self.LP_1(out_1)\n",
    "        lp_2 = self.LP_2(out_2)  \n",
    "        lp_3 = self.LP_3(out_3)  \n",
    "        lp_4 = self.LP_4(out_4)\n",
    "        \n",
    "        # linear fuse and go pass LP Header\n",
    "        lp_34 = self.LP_34(self.linear_fuse34(torch.cat([lp_3, F.interpolate(lp_4,scale_factor=2,mode='bilinear', align_corners=False)], dim=1)))\n",
    "        lp_23 = self.LP_23(self.linear_fuse23(torch.cat([lp_2, F.interpolate(lp_34,scale_factor=2,mode='bilinear', align_corners=False)], dim=1)))\n",
    "        lp_12 = self.LP_12(self.linear_fuse12(torch.cat([lp_1, F.interpolate(lp_23,scale_factor=2,mode='bilinear', align_corners=False)], dim=1)))\n",
    "        \n",
    "        # get the final output\n",
    "        lp4_resized = F.interpolate(lp_4,scale_factor=8,mode='bilinear', align_corners=False)\n",
    "        lp3_resized = F.interpolate(lp_34,scale_factor=4,mode='bilinear', align_corners=False)\n",
    "        lp2_resized = F.interpolate(lp_23,scale_factor=2,mode='bilinear', align_corners=False)\n",
    "        lp1_resized = lp_12\n",
    "        \n",
    "        out = self.linear_pred(torch.cat([lp1_resized, lp2_resized, lp3_resized, lp4_resized], dim=1))\n",
    "        out_resized = F.interpolate(out,scale_factor=4,mode='bilinear', align_corners=True)\n",
    "        \n",
    "        return out_resized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e858b5ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ange_structure_loss(pred, mask, smooth=1):\n",
    "    \n",
    "    weit = 1 + 5*torch.abs(F.avg_pool2d(mask, kernel_size=15, stride=1, padding=7) - mask)\n",
    "    wbce = F.binary_cross_entropy_with_logits(pred, mask, reduction='mean')\n",
    "    wbce = (weit*wbce).sum(dim=(2, 3)) / weit.sum(dim=(2, 3))\n",
    "\n",
    "    pred = torch.sigmoid(pred)\n",
    "    inter = ((pred * mask)*weit).sum(dim=(2, 3))\n",
    "    union = ((pred + mask)*weit).sum(dim=(2, 3))\n",
    "    wiou = 1 - (inter + smooth)/(union - inter + smooth)\n",
    "    \n",
    "    return (wbce + wiou).mean()\n",
    "\n",
    "def dice_loss_coff(pred, target, smooth = 0.0001):\n",
    "    \n",
    "    num = target.size(0)\n",
    "    pred = pred.contiguous()\n",
    "    target = target.contiguous()    \n",
    "\n",
    "    intersection = (pred * target).sum(dim=2).sum(dim=2)\n",
    "    loss = (2. * intersection + smooth) / (pred.sum(dim=2).sum(dim=2) + target.sum(dim=2).sum(dim=2) + smooth)\n",
    "    \n",
    "    return loss.sum()/num"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68c112d6",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from torch.autograd import Variable\n",
    "\n",
    "def evaluate(): \n",
    "    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "    ESFPNet.eval()\n",
    "    val = 0\n",
    "    dataset_count = 0\n",
    "    datasetValidation = []\n",
    "    smooth = 1e-4\n",
    "    for _data_name in datasetTest:\n",
    "        test_loader = test_dataset(config['dataset']['test_' + str(_data_name) + '_img'], config['dataset']['test_' + str(_data_name) + '_label'], 352)\n",
    "        Thresholds = np.linspace(1, 0, 256)\n",
    "        count = 0\n",
    "        total_meanDic = 0\n",
    "            \n",
    "        for i in range(test_loader.size):\n",
    "            image, gt, name = test_loader.load_data()\n",
    "            gt = np.asarray(gt, np.float32)\n",
    "            gt /= (gt.max() + 1e-8)\n",
    "            target = np.array(gt)\n",
    "\n",
    "            image = image.cuda()\n",
    "\n",
    "            pred = ESFPNet(image)\n",
    "                \n",
    "            pred = F.upsample(pred, size=gt.shape, mode='bilinear', align_corners=False)\n",
    "            pred = pred.sigmoid()\n",
    "            threshold = torch.tensor([0.5]).to(device)\n",
    "            pred = (pred > threshold).float() * 1\n",
    "            pred = pred.data.cpu().numpy().squeeze()\n",
    "            pred = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)\n",
    "\n",
    "            input_flat = np.reshape(pred,(-1))\n",
    "            target_flat = np.reshape(target,(-1))\n",
    " \n",
    "            intersection = (input_flat*target_flat)\n",
    "        \n",
    "            loss =  (2 * intersection.sum() + smooth) / (pred.sum() + target.sum() + smooth)\n",
    "\n",
    "            Dice =  '{:.4f}'.format(loss)\n",
    "            Dice = float(Dice)\n",
    "            total_meanDic = total_meanDic + Dice\n",
    "            count = count + 1\n",
    "        datasetValidation.append(total_meanDic/count)\n",
    "        val = val + total_meanDic/count\n",
    "        dataset_count = dataset_count + 1\n",
    "        \n",
    "    ESFPNet.train()\n",
    "    \n",
    "    return val/dataset_count, datasetValidation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ce57ffc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# train the network\n",
    "def training_loop(dataloader_X, n_epochs, coeff_max_global, datasetValidation_max, ESFPNet_optimizer, numIters):\n",
    "    \n",
    "    # keep track of losses over time\n",
    "    losses = []\n",
    "    coeff_max = 0;\n",
    "    \n",
    "    validation_coeff, datasetValidation = evaluate();\n",
    "    if coeff_max<validation_coeff:\n",
    "        coeff_max = validation_coeff\n",
    "    \n",
    "    # set up data and then train\n",
    "    iter_X = iter(dataloader_X)\n",
    "    steps_per_epoch = len(iter_X)\n",
    "    num_epoch = 0\n",
    "    total_steps = (n_epochs+1)*steps_per_epoch\n",
    "    \n",
    "    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "    \n",
    "    for step in range(1, total_steps):\n",
    "\n",
    "        # Reset iterators for each epoch\n",
    "        if step % steps_per_epoch == 0:\n",
    "            iter_X = iter(dataloader_X)\n",
    "            num_epoch = num_epoch + 1\n",
    "        \n",
    "        # make sure to scale to a range -1 to 1\n",
    "        images, masks = iter_X.next()\n",
    "        \n",
    "        # move images to GPU if available (otherwise stay on CPU)\n",
    "        images = images.to(device)\n",
    "        masks = masks.to(device)\n",
    "\n",
    "\n",
    "        # ============================================\n",
    "        #            TRAIN THE NETWORKS\n",
    "        # ============================================\n",
    "       \n",
    "        ESFPNet_optimizer.zero_grad()\n",
    "        \n",
    "        # 1. Compute the losses from the network\n",
    "        \n",
    "        out = ESFPNet(images)\n",
    "        \n",
    "        loss = ange_structure_loss(out, masks)\n",
    "        \n",
    "        loss.backward()\n",
    "        ESFPNet_optimizer.step() \n",
    "        \n",
    "        # ============================================\n",
    "        #            TRAIN THE NETWORKS\n",
    "        # ============================================\n",
    "        # Print the log info\n",
    "        if step % steps_per_epoch == 0:\n",
    "            \n",
    "            losses.append(loss.item())\n",
    "            print('Epoch [{:5d}/{:5d}] | preliminary loss: {:6.6f} '.format(num_epoch, n_epochs, loss.item()))\n",
    "            \n",
    "        if step % steps_per_epoch == 0:\n",
    "            \n",
    "            validation_coeff, datasetValidation = evaluate()\n",
    "            print('Epoch [{:5d}/{:5d}] | validation_coeffient: {:6.6f} '.format(\n",
    "                    num_epoch, n_epochs, validation_coeff))\n",
    "            \n",
    "            if coeff_max < validation_coeff:\n",
    "                coeff_max = validation_coeff\n",
    "            \n",
    "            if coeff_max_global < validation_coeff:\n",
    "                coeff_max_global = validation_coeff\n",
    "                save_model_path = './SaveModel/{}_AverageOptimize_{:1d}'.format(_model_name,numIters)\n",
    "                os.makedirs(save_model_path, exist_ok=True)\n",
    "                torch.save(ESFPNet, save_model_path + '/ESFPNet.pt')\n",
    "                print('Save Average Optimized Model at Epoch [{:5d}/{:5d}]'.format(num_epoch, n_epochs))\n",
    "            \n",
    "            count_dataset = 0\n",
    "            for _data_name in datasetTest:\n",
    "                if datasetValidation_max[count_dataset]<datasetValidation[count_dataset]:\n",
    "                    datasetValidation_max[count_dataset]=datasetValidation[count_dataset]\n",
    "                    save_model_path = './SaveModel/{}_DatasetBest_{:1d}/{}'.format(_model_name,numIters,_data_name)\n",
    "                    os.makedirs(save_model_path, exist_ok=True)\n",
    "                    torch.save(ESFPNet, save_model_path + '/ESFPNet.pt')\n",
    "                    print('Save Optimized {} Model at Epoch [{:5d}/{:5d}, with coefficient: {:6.6f}]'.format(_data_name, num_epoch, n_epochs, datasetValidation[count_dataset]))\n",
    "                count_dataset = count_dataset + 1\n",
    "                \n",
    "            \n",
    "    return losses, coeff_max_global, datasetValidation_max"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74d06570",
   "metadata": {},
   "outputs": [],
   "source": [
    "def saveResult(numIters):\n",
    "    for _data_name in datasetTest:\n",
    "    \n",
    "        save_path = './results/{}_DatasetBest_{:1d}/{}/'.format(_model_name,numIters,_data_name)\n",
    "        os.makedirs(save_path, exist_ok=True)\n",
    "        \n",
    "        test_loader = test_dataset(config['dataset']['test_' + str(_data_name) + '_img'], config['dataset']['test_' + str(_data_name) + '_label'], 352)\n",
    "    \n",
    "        model_path = './SaveModel/{}_DatasetBest_{:1d}/{}'.format(_model_name,numIters,_data_name)\n",
    "        TransSegBest = torch.load(model_path + '/ESFPNet.pt')\n",
    "\n",
    "        for i in range(test_loader.size):\n",
    "            image, gt, name = test_loader.load_data()\n",
    "            gt = np.asarray(gt, np.float32)\n",
    "            gt /= (gt.max() + 1e-8)\n",
    "            image = image.cuda()\n",
    "\n",
    "            pred = TransSegBest(image)\n",
    "            pred = F.upsample(pred, size=gt.shape, mode='bilinear', align_corners=False)\n",
    "            pred = pred.sigmoid()\n",
    "            threshold = torch.tensor([0.5]).to(device)\n",
    "            pred = (pred > threshold).float() * 1\n",
    "            pred = pred.data.cpu().numpy().squeeze()\n",
    "            pred = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)\n",
    "        \n",
    "            imageio.imwrite(save_path+name,img_as_ubyte(pred))\n",
    "\n",
    "    for _data_name in datasetTest:\n",
    "    \n",
    "        save_path = './results/{}_AverageOptimal_{:1d}/{}/'.format(_model_name,numIters,_data_name)\n",
    "        os.makedirs(save_path, exist_ok=True)\n",
    "        \n",
    "        test_loader = test_dataset(config['dataset']['test_' + str(_data_name) + '_img'], config['dataset']['test_' + str(_data_name) + '_label'], 352)\n",
    "    \n",
    "        model_path = './SaveModel/{}_AverageOptimize_{:1d}'.format(_model_name,numIters)\n",
    "        TransSegBest = torch.load(model_path + '/ESFPNet.pt')\n",
    "\n",
    "        for i in range(test_loader.size):\n",
    "            image, gt, name = test_loader.load_data()\n",
    "            gt = np.asarray(gt, np.float32)\n",
    "            gt /= (gt.max() + 1e-8)\n",
    "            image = image.cuda()\n",
    "\n",
    "            pred = TransSegBest(image)\n",
    "            pred = F.upsample(pred, size=gt.shape, mode='bilinear', align_corners=False)\n",
    "            pred = pred.sigmoid()\n",
    "            threshold = torch.tensor([0.5]).to(device)\n",
    "            pred = (pred > threshold).float() * 1\n",
    "            pred = pred.data.cpu().numpy().squeeze()\n",
    "            pred = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)\n",
    "        \n",
    "            imageio.imwrite(save_path+name,img_as_ubyte(pred))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ced1cb9e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "for i in range(repeats):\n",
    "    \n",
    "    coeff_max_global = 0\n",
    "    datasetValidation_max = [0,0,0,0,0]\n",
    "    \n",
    "    print('#####################################################################################')  \n",
    "    print('[{:5d}/{:5d}]'.format(i+1, repeats))\n",
    "    \n",
    "    # Clear GPU cache\n",
    "    torch.cuda.empty_cache()\n",
    "    \n",
    "    ESFPNet = ESFPNetStructure()\n",
    "    if torch.cuda.is_available():\n",
    "        device = torch.device(\"cuda:0\")\n",
    "        ESFPNet.to(device)\n",
    "        print('Models moved to GPU.')\n",
    "    else:\n",
    "        print('Only CPU available.')\n",
    "    print('#####################################################################################')  \n",
    "        \n",
    "    # hyperparams for Adam optimizer\n",
    "    lr=0.0001 #0.0001\n",
    "\n",
    "    ESFPNet_optimizer = optim.AdamW(ESFPNet.parameters(), lr=lr)\n",
    "\n",
    "    # scheduler_warmup is chained with schduler_steplr\n",
    "    losses, coeff_max_global, datasetValidation_max = training_loop(train_loader, n_epochs, coeff_max_global, datasetValidation_max, ESFPNet_optimizer, i+1)\n",
    "    \n",
    "    plt.plot(losses)\n",
    "    \n",
    "    print('#####################################################################################')  \n",
    "    print('optimize_average_m_dice: {:6.6f}'.format(coeff_max_global))\n",
    "    print('Kvasir: {:6.6f} | CVC-ColonDB: {:6.6f} | CVC-ClinicDB: {:6.6f} | ETIS-LaribPolypDB: {:6.6f}'.format(datasetValidation_max[0],datasetValidation_max[1],datasetValidation_max[2],datasetValidation_max[3]))\n",
    "    print('#####################################################################################')  \n",
    "    \n",
    "    saveResult(i+1)\n",
    "    print('#####################################################################################')  \n",
    "    print('saved the results')\n",
    "    print('#####################################################################################')  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97cf48ac",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfa9e630",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}