ESFPNet_Endoscope_Learning_Ability.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",
    "WholeDatasetName = 'CVC-ClinicDB'\n",
    "#WholeDatasetName = 'Kvasir'\n",
    "\n",
    "model_type = 'B4'\n",
    "_model_name = 'ESFP_{}_Endo_{}'.format(model_type,WholeDatasetName)\n",
    "config = open('Configure.yaml')\n",
    "config = yaml.safe_load(config)\n",
    "\n",
    "init_trainsize = 352\n",
    "batch_size = 8\n",
    "\n",
    "repeats = 1\n",
    "n_epochs = 200\n",
    "if_renew = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f7b2d99",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SplittingDataset(Dataset):\n",
    "    \"\"\"\n",
    "    dataloader for polyp segmentation tasks\n",
    "    \"\"\"\n",
    "    def __init__(self, image_root, gt_root):\n",
    "\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('.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",
    "        self.transform = transforms.ToTensor()\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \n",
    "        image = self.rgb_loader(self.images[index])\n",
    "        gt = self.binary_loader(self.gts[index])\n",
    "        name = self.images[index].split('/')[-1]\n",
    "        return self.transform(image), self.transform(gt), name\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('L')\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c3dcbd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def splitDataset(renew):\n",
    "    \n",
    "    split_train_images_save_path = './Endoscope-WL/{}_Splited/trainSplited/images'.format(WholeDatasetName)\n",
    "    os.makedirs(split_train_images_save_path, exist_ok=True)\n",
    "    split_train_masks_save_path = './Endoscope-WL/{}_Splited/trainSplited/masks'.format(WholeDatasetName)\n",
    "    os.makedirs(split_train_masks_save_path, exist_ok=True)\n",
    "    \n",
    "    split_validation_images_save_path = './Endoscope-WL/{}_Splited/validationSplited/images'.format(WholeDatasetName)\n",
    "    os.makedirs(split_validation_images_save_path, exist_ok=True)\n",
    "    split_validation_masks_save_path = './Endoscope-WL/{}_Splited/validationSplited/masks'.format(WholeDatasetName)\n",
    "    os.makedirs(split_validation_masks_save_path, exist_ok=True)\n",
    "    \n",
    "    split_test_images_save_path = './Endoscope-WL/{}_Splited/testSplited/images'.format(WholeDatasetName)\n",
    "    os.makedirs(split_test_images_save_path, exist_ok=True)\n",
    "    split_test_masks_save_path = './Endoscope-WL/{}_Splited/testSplited/masks'.format(WholeDatasetName)\n",
    "    os.makedirs(split_test_masks_save_path, exist_ok=True)\n",
    "    \n",
    "    if renew == True:\n",
    "    \n",
    "        DatasetList = []\n",
    "\n",
    "        images_train_path = config['dataset']['train_' + str(WholeDatasetName) + '_dataset'] + '/images/'\n",
    "        masks_train_path = config['dataset']['train_' + str(WholeDatasetName) + '_dataset'] + '/masks/'\n",
    "        Dataset_part_train = SplittingDataset(images_train_path, masks_train_path)\n",
    "        DatasetList.append(Dataset_part_train)\n",
    "\n",
    "        images_test_path = config['dataset']['test_' + str(WholeDatasetName) + '_img']\n",
    "        masks_test_path = config['dataset']['test_' + str(WholeDatasetName) + '_label']\n",
    "        Dataset_part_test = SplittingDataset(images_test_path, masks_test_path)\n",
    "        DatasetList.append(Dataset_part_test)\n",
    "                                    \n",
    "        wholeDataset = ConcatDataset([DatasetList[0], DatasetList[1]])\n",
    "\n",
    "        val_num = int(0.1*len(wholeDataset))\n",
    "        test_num = int(0.1*len(wholeDataset))\n",
    "                                    \n",
    "        train_num = len(wholeDataset) - val_num - test_num\n",
    "        train_dataset, val_dataset, test_dataset = torch.utils.data.random_split(wholeDataset, [train_num, val_num, test_num])\n",
    "                                    \n",
    "        train_loader = DataLoader(dataset=train_dataset,batch_size=1,shuffle=False)\n",
    "        val_loader = DataLoader(dataset=val_dataset,batch_size=1,shuffle=False)\n",
    "        test_loader = DataLoader(dataset=test_dataset,batch_size=1,shuffle=False)\n",
    "    \n",
    "        iter_train = iter(train_loader)\n",
    "        for i in range(len(train_loader)):\n",
    "            image, gt, name = iter_train.next()\n",
    "            image_data = image.data.cpu().numpy().squeeze().transpose(1,2,0)\n",
    "            gt_data = gt.data.cpu().numpy().squeeze()\n",
    "            imageio.imwrite(split_train_images_save_path + '/' + name[0],img_as_ubyte(image_data))\n",
    "            imageio.imwrite(split_train_masks_save_path + '/' + name[0],img_as_ubyte(gt_data))\n",
    "            \n",
    "        iter_val = iter(val_loader)\n",
    "        for i in range(len(val_loader)):\n",
    "            image, gt, name = iter_val.next()\n",
    "            image_data = image.data.cpu().numpy().squeeze().transpose(1,2,0)\n",
    "            gt_data = gt.data.cpu().numpy().squeeze()\n",
    "            imageio.imwrite(split_validation_images_save_path + '/' + name[0],img_as_ubyte(image_data))\n",
    "            imageio.imwrite(split_validation_masks_save_path + '/' + name[0],img_as_ubyte(gt_data))\n",
    "            \n",
    "        \n",
    "        iter_test = iter(test_loader)\n",
    "        for i in range(len(test_loader)):\n",
    "            image, gt, name = iter_test.next()\n",
    "            image_data = image.data.cpu().numpy().squeeze().transpose(1,2,0)\n",
    "            gt_data = gt.data.cpu().numpy().squeeze()\n",
    "            imageio.imwrite(split_test_images_save_path + '/' + name[0],img_as_ubyte(image_data))\n",
    "            imageio.imwrite(split_test_masks_save_path + '/' + name[0],img_as_ubyte(gt_data))\n",
    "            \n",
    "    \n",
    "    return split_train_images_save_path, split_train_masks_save_path, split_validation_images_save_path, split_validation_masks_save_path, split_test_images_save_path, split_test_masks_save_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60f6ba8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_images_path, train_masks_path, val_images_path, val_masks_path, test_images_path, test_masks_path = splitDataset(if_renew)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15e5a003",
   "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": "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",
    "        \n",
    "        return out"
   ]
  },
  {
   "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",
    "    \n",
    "    val = 0\n",
    "    count = 0\n",
    "\n",
    "    smooth = 1e-4\n",
    "    \n",
    "    val_loader = test_dataset(val_images_path + '/',val_masks_path + '/', init_trainsize)\n",
    "    for i in range(val_loader.size):\n",
    "        image, gt, name = val_loader.load_data()\n",
    "        gt = np.asarray(gt, np.float32)\n",
    "        gt /= (gt.max() + 1e-8)\n",
    "\n",
    "        image = image.cuda()\n",
    "        \n",
    "        pred= ESFPNet(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",
    "\n",
    "        pred = pred.data.cpu().numpy().squeeze()\n",
    "        pred = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)\n",
    "        \n",
    "        target = np.array(gt)\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",
    "        a =  '{:.4f}'.format(loss)\n",
    "        a = float(a)\n",
    "        \n",
    "        val = val + a\n",
    "        count = count + 1\n",
    "        \n",
    "    ESFPNet.train()\n",
    "    \n",
    "    return val/count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ce57ffc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# train the network\n",
    "def training_loop(n_epochs, ESFPNet_optimizer, numIters):\n",
    "    \n",
    "    # keep track of losses over time\n",
    "    losses = []\n",
    "    coeff_max = 0;\n",
    "    \n",
    "    # set up data and then train\n",
    "    trainDataset = PolypDataset(train_images_path + '/', train_masks_path + '/', trainsize=init_trainsize, augmentations = True)\n",
    "    train_loader = DataLoader(dataset=trainDataset,batch_size=batch_size,shuffle=True)\n",
    "    \n",
    "    iter_X = iter(train_loader)\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(train_loader)\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",
    "        out = F.interpolate(out, scale_factor=4, mode='bilinear', align_corners=False)\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 = 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",
    "                save_model_path = './SaveModel/{}_LA_{:1d}'.format(_model_name,numIters)\n",
    "                os.makedirs(save_model_path, exist_ok=True)\n",
    "                print(save_model_path)\n",
    "                torch.save(ESFPNet, save_model_path + '/ESFPNet.pt')\n",
    "                print('Save Learning Ability Optimized Model at Epoch [{:5d}/{:5d}]'.format(num_epoch, n_epochs))\n",
    "                \n",
    "    return losses, coeff_max"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74d06570",
   "metadata": {},
   "outputs": [],
   "source": [
    "def saveResult(numIters):\n",
    "    \n",
    "    save_path = './results/{}_LA_{:1d}/{}_Splited/'.format(_model_name,numIters,str(WholeDatasetName))\n",
    "    os.makedirs(save_path, exist_ok=True)\n",
    "    print(save_path)\n",
    "        \n",
    "    model_path =  './SaveModel/{}_LA_{:1d}'.format(_model_name,numIters)\n",
    "    TransSegBest = torch.load(model_path + '/ESFPNet.pt')\n",
    "    \n",
    "    test_loader = test_dataset(test_images_path + '/', test_masks_path + '/', init_trainsize)\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": "114565ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "for i in range(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",
    "    losses, coeff_max = training_loop(n_epochs, ESFPNet_optimizer, i+1)\n",
    "    \n",
    "    plt.plot(losses)\n",
    "    \n",
    "    print('#####################################################################################')  \n",
    "    print('optimize_m_dice: {:6.6f}'.format(coeff_max))\n",
    "\n",
    "    saveResult(i+1)\n",
    "    print('#####################################################################################')  \n",
    "    print('saved the results')\n",
    "    print('#####################################################################################')  "
   ]
  }
 ],
 "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
}