U_Net_128.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0d18f380",
   "metadata": {
    "tags": []
   },
   "source": [
    "###### Importing libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dd0e5e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import rasterio as rio \n",
    "import numpy as np\n",
    "import os\n",
    "import json\n",
    "from rasterio.features import rasterize\n",
    "from osgeo import gdal\n",
    "from matplotlib import pyplot as plt\n",
    "from scipy.interpolate import NearestNDInterpolator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b8f2dc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision\n",
    "import gc\n",
    "from torch.utils.data import DataLoader, random_split, RandomSampler\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from torch import nn, optim\n",
    "from torch.optim import Adam\n",
    "from torchvision import transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc9102a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "from tqdm.autonotebook import tqdm\n",
    "from sklearn.metrics import accuracy_score\n",
    "from sklearn.metrics import jaccard_score"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d79b86bf-b8dc-4723-8f69-022c94536005",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "940be772-a271-4417-8127-c1b2014cc3b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "biobio_dataset=pd.read_csv(\"firescarbiobio128/biobio_final_128_10_03.csv\")  \n",
    "valparaiso_dataset=pd.read_csv(\"firescarvalpo128/valparaiso_128_01_02.csv\")\n",
    "dataset=pd.concat([valparaiso_dataset,biobio_dataset], axis=0, ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e013d5b7-99a5-48a6-8803-a28e9c903856",
   "metadata": {},
   "outputs": [],
   "source": [
    "LS_max=[3660.5,  4303.5,  4832.0,  6956.0,  6174.5,  6234.0, 1,  1000,  3811.0,  4504.0,  4950.5,\n",
    "7000.5,  6210.0,  6286.5, 1,  1000]\n",
    "LI_min= [0.0,  0.0,  0.0,  0.0,  0.0,  0.0,  -0.69034,  -605.32933, 0.0,  0.0,  0.0,  0.0,  0.0,\n",
    "0.0,  -0.7970004,  -573.05183]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea04899a-0a07-45f8-b3eb-70b99fdd7a71",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def preprocessing(imgdata):\n",
    "    LS_max=[3660.5,  4303.5,  4832.0,  6956.0,  6174.5,  6234.0, 1,  1000,  3811.0,  4504.0,  4950.5,\n",
    "    7000.5,  6210.0,  6286.5, 1,  1000]\n",
    "    LI_min= [0.0,  0.0,  0.0,  0.0,  0.0,  0.0,  -0.6903,  -605.3293, 0.0,  0.0,  0.0,  0.0,  0.0,\n",
    "    0.0,  -0.7970,  -573.0518]\n",
    "    mean_sprepro=[415.8392, 646.4156, 761.339, 2109.5385, 1843.5207, 1189.6269, 0.4677, 301.9405,\n",
    "    414.9891, 657.8648, 761.992, 2235.3473, 1853.7414, 1160.7475, 0.4924, 341.9075]\n",
    "    for k in range(1,17):\n",
    "        if (imgdata[k-1]>LS_max[k-1]).any():\n",
    "            if imgdata[k-1].mean()<LS_max[k-1]:\n",
    "                imgdata[k-1][imgdata[k-1]>LS_max[k-1]]=imgdata[k-1].mean()\n",
    "            else:\n",
    "                imgdata[k-1][imgdata[k-1]>LS_max[k-1]]=mean_sprepro[k-1]\n",
    "        elif (imgdata[k-1]<LI_min[k-1]).any():\n",
    "            if imgdata[k-1].mean()>LI_min[k-1]:\n",
    "                imgdata[k-1][imgdata[k-1]<LI_min[k-1]]=imgdata[k-1].mean()\n",
    "            else: \n",
    "                imgdata[k-1][imgdata[k-1]<LI_min[k-1]]=mean_sprepro[k-1]\n",
    "    return imgdata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b504663-6e8d-4ff0-bcdd-41c864fd47de",
   "metadata": {},
   "outputs": [],
   "source": [
    "class firescardataset():\n",
    "    def __init__(self,dataset, subset_size1, subset_size2, subset_size3,subset_size4,mult=1,transform=None):\n",
    "        \"\"\":param dataset: dataset with data filenames from two different regions. The first one from the index subset_size1-subset_size2, and the\n",
    "        second one from subset_size3-subset_size4. There are 3 columns with the required data filenames for each input: \n",
    "        \"ImPosF\"=The image post Fire, \"ImgPreF\"=The image pre Fire, and \"FireScar_tif\"=The label, in a raster file\"\"\"\n",
    "        self.transform = transform\n",
    "        # list of image files (pre and post fire), and labels\n",
    "        # label vector edge coordinates\n",
    "        self.imgfiles = []\n",
    "        self.imgprefiles=[]\n",
    "        self.labels = []\n",
    "        self.seglabels = []\n",
    "        imgposfiles = []\n",
    "        # read in segmentation label files\n",
    "        for i in range(subset_size1,subset_size2):\n",
    "            segdata = os.path.join(\"firescarvalpoallsizes/FireScar/\", dataset.loc[i,\"FireScar_tif\"])\n",
    "            self.seglabels.append(segdata)\n",
    "            self.imgfiles.append(os.path.join(\"firescarvalpo128/ImgPosF/\",dataset.loc[i,\"ImgPosF\"]))\n",
    "            self.imgprefiles.append(os.path.join(\"firescarvalpo128/ImgPreF/\",dataset.loc[i,\"ImgPreF\"]))\n",
    "        for i in range(subset_size3,subset_size4):\n",
    "            self.seglabels.append(os.path.join(\"firescarbiobioallsizes/FireScar/\",dataset.loc[i,\"FireScar_tif\"]))\n",
    "            self.imgfiles.append(os.path.join(\"firescarbiobio128/ImgPosF/\",dataset.loc[i,\"ImgPosF\"]))\n",
    "            self.imgprefiles.append(os.path.join(\"firescarbiobio128/ImgPreF/\",dataset.loc[i,\"ImgPreF\"]))\n",
    "        self.imgfiles = np.array(self.imgfiles)\n",
    "        self.imgprefiles=np.array(self.imgprefiles)\n",
    "        self.labels = np.array(self.labels)\n",
    "        \n",
    "        if mult > 1:\n",
    "            self.imgfiles = np.array([*self.imgfiles] * mult)\n",
    "            self.imgprefiles = np.array([*self.imgprefiles] * mult)\n",
    "            self.labels = np.array([*self.labels] * mult)\n",
    "            self.seglabels = self.seglabels * mult\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.imgfiles)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        idx=idx-1\n",
    "        imgfile = rio.open(self.imgfiles[idx])\n",
    "        imgpre=rio.open(self.imgprefiles[idx])\n",
    "        imgdata1 = np.array([imgfile.read(i) for i in [1,2,3,4,5,6,7,8]])\n",
    "        imgdatapre=np.array([imgpre.read(i) for i in [1,2,3,4,5,6,7,8]])\n",
    "        imgdata=np.concatenate((imgdata1, imgdatapre), axis=0)\n",
    "        imgdata[imgdata==0]=np.nan\n",
    "\n",
    "        if (np.isfinite(imgdata)==False).any():                      #Replace nan for the neighbours mean values\n",
    "            mask=np.where(np.isfinite(imgdata))\n",
    "            interp=NearestNDInterpolator(np.transpose(mask), imgdata[mask])\n",
    "            imgdata=interp(*np.indices(imgdata.shape))\n",
    "\n",
    "        ds = gdal.Open(self.seglabels[idx])\n",
    "        myarray = np.array(ds.GetRasterBand(1).ReadAsArray())\n",
    "\n",
    "        x=imgdata1.shape[1]\n",
    "        y=imgdata1.shape[2]\n",
    "\n",
    "      #FireScar padding to 128 in case is not that size\n",
    "        x,y=myarray.shape\n",
    "                                                                    #only to equalize to 128x128 images or it could be to image size \n",
    "        ulx_i, lry_i, lrx_i, uly_i=imgfile.bounds\n",
    "        ulx, xres, xskew, uly, yskew, yres  = ds.GetGeoTransform()\n",
    "        lrx = ulx + (ds.RasterXSize * xres)\n",
    "        lry = uly + (ds.RasterYSize * yres)\n",
    "        left=round((ulx-ulx_i)/xres)    #np.pad(a, up, down, left, right)\n",
    "        right=round((lrx_i-lrx)/xres)\n",
    "        up=round((uly-uly_i)/yres)\n",
    "        down=round((lry_i-lry)/yres)\n",
    "        myarray=np.pad(myarray,((up, down),(left,right)),\"constant\")\n",
    "\n",
    "        imgdata=preprocessing(imgdata)                               #preprocessing to the data when there are values off range (i.e outliers)\n",
    "\n",
    "        sample = {'idx': idx,\n",
    "              'img': imgdata,\n",
    "              'fpt': myarray,\n",
    "              'imgfile': self.imgfiles[idx]}\n",
    "        if self.transform:\n",
    "            sample = self.transform(sample)\n",
    "        return sample\n",
    "    \n",
    "class ToTensor(object):\n",
    "#     \"\"\"Convert ndarrays in sample to Tensors.\"\"\"\n",
    "    def __call__(self, sample):\n",
    "#         \"\"\"\n",
    "#         :param sample: sample to be converted to Tensor\n",
    "#         :return: converted Tensor sample\n",
    "#         \"\"\"\n",
    "        out = {'idx': sample['idx'],\n",
    "        'img': torch.from_numpy(sample['img'].copy()),\n",
    "        'fpt': torch.from_numpy(sample['fpt'].copy()),\n",
    "        'imgfile': sample['imgfile']}\n",
    "\n",
    "        return out\n",
    "class Randomize(object):\n",
    "    \"\"\"Randomize image orientation including rotations by integer multiples of\n",
    "       90 deg, (horizontal) mirroring, and (vertical) flipping.\"\"\"\n",
    "\n",
    "    def __call__(self, sample):\n",
    "        \"\"\"\n",
    "        :param sample: sample to be randomized\n",
    "        :return: randomized sample\n",
    "        \"\"\"\n",
    "        imgdata = sample['img']\n",
    "        fptdata = sample['fpt']\n",
    "        idx=sample[\"idx\"]\n",
    "        # mirror horizontally\n",
    "        mirror = np.random.randint(0, 2)\n",
    "        if mirror:\n",
    "            imgdata = np.flip(imgdata, 2)\n",
    "            fptdata = np.flip(fptdata, 1)\n",
    "        # flip vertically\n",
    "        flip = np.random.randint(0, 2)\n",
    "        if flip:\n",
    "            imgdata = np.flip(imgdata, 1)\n",
    "            fptdata = np.flip(fptdata, 0)\n",
    "        # rotate by [0,1,2,3]*90 deg\n",
    "        rot = np.random.randint(0, 4)\n",
    "        if rot:\n",
    "            imgdata = np.rot90(imgdata, rot, axes=(1,2))\n",
    "            fptdata = np.rot90(fptdata, rot, axes=(0,1))\n",
    "\n",
    "        return {'idx': sample['idx'],\n",
    "                'img': imgdata.copy(),\n",
    "                'fpt': fptdata.copy(),\n",
    "                'imgfile': sample['imgfile']}\n",
    "\n",
    "class Normalize(object):\n",
    "    \"\"\"Normalize pixel values to the range [0, 1] measured using minmax-scaling\"\"\"\n",
    "    def __init__(self):\n",
    "\n",
    "        self.channel_min=np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.6737967729568481, -605.4163208007812,\n",
    "                                   0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.7714285850524902, -572.1392822265625])\n",
    "        \n",
    "        self.channel_max=np.array([3560.0, 4147.0, 4780.0, 6832.0, 6143.0, 6261.0, 1.0, 1000.0,\n",
    "                                   3765.0, 4398.0, 4866.0, 6930.0, 6157.0, 6167.0, 1.0, 1000.0])\n",
    "            \n",
    "    def __call__(self, sample):\n",
    "        \"\"\"\n",
    "        :param sample: sample to be normalized\n",
    "        :return: normalized sample\"\"\"\n",
    "        sample['img'] = (sample['img']-self.channel_min.reshape(\n",
    "            sample['img'].shape[0], 1, 1))/(self.channel_max.reshape(\n",
    "            sample['img'].shape[0], 1, 1)-self.channel_min.reshape(\n",
    "            sample['img'].shape[0], 1, 1))\n",
    "        return sample \n",
    "def create_dataset(*args, apply_transforms=True, **kwargs):\n",
    "    \"\"\"Create a dataset; uses same input parameters as PowerPlantDataset.\n",
    "    :param apply_transforms: if `True`, apply available transformations\n",
    "    :return: data set\"\"\"\n",
    "    if apply_transforms:\n",
    "        data_transforms = transforms.Compose([\n",
    "            Normalize(),\n",
    "            Randomize(),\n",
    "            ToTensor()\n",
    "           ])\n",
    "    else:\n",
    "        data_transforms = None\n",
    "\n",
    "    data = firescardataset(*args, **kwargs,\n",
    "                                         transform=data_transforms)\n",
    "    \n",
    "    return data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd9c2e81",
   "metadata": {},
   "source": [
    "##### Modelo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbd97f05",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a478234",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57163895",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "class DoubleConv(nn.Module):\n",
    "#     \"\"\"(convolution => [BN] => ReLU) * 2\"\"\"\n",
    "\n",
    "    def __init__(self, in_channels, out_channels, mid_channels=None):\n",
    "        super().__init__()\n",
    "        if not mid_channels:\n",
    "            mid_channels = out_channels\n",
    "        self.double_conv = nn.Sequential(\n",
    "            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(mid_channels),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(out_channels),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.double_conv(x)\n",
    "\n",
    "class Down(nn.Module):\n",
    "#     \"\"\"Downscaling with maxpool then double conv\"\"\"\n",
    "\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super().__init__()\n",
    "        self.maxpool_conv = nn.Sequential(\n",
    "            nn.MaxPool2d(2),\n",
    "            DoubleConv(in_channels, out_channels)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.maxpool_conv(x)\n",
    "\n",
    "\n",
    "class Up(nn.Module):\n",
    "#     \"\"\"Upscaling then double conv\"\"\"\n",
    "\n",
    "    def __init__(self, in_channels, out_channels, bilinear=True):\n",
    "        super().__init__()\n",
    "\n",
    "        # if bilinear, use the normal convolutions to reduce the number of channels\n",
    "        if bilinear:\n",
    "            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n",
    "            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)\n",
    "        else:\n",
    "            self.up = nn.ConvTranspose2d(in_channels , in_channels // 2, kernel_size=2, stride=2)\n",
    "            self.conv = DoubleConv(in_channels, out_channels)\n",
    "\n",
    "\n",
    "    def forward(self, x1, x2):\n",
    "        x1 = self.up(x1)\n",
    "        # input is CHW\n",
    "        diffY = x2.size()[2] - x1.size()[2]\n",
    "        diffX = x2.size()[3] - x1.size()[3]\n",
    "\n",
    "        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,\n",
    "                        diffY // 2, diffY - diffY // 2])\n",
    "        # if you have padding issues, see\n",
    "        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a\n",
    "        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd\n",
    "        x = torch.cat([x2, x1], dim=1)\n",
    "        return self.conv(x)\n",
    "\n",
    "class OutConv(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super(OutConv, self).__init__()\n",
    "        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.conv(x)\n",
    "\n",
    "class UNet(nn.Module):\n",
    "    def __init__(self, n_channels, n_classes, bilinear=True):\n",
    "        super(UNet, self).__init__()\n",
    "        self.n_channels = n_channels\n",
    "        self.n_classes = n_classes\n",
    "        self.bilinear = bilinear\n",
    "        \n",
    "        self.inc = DoubleConv(n_channels, 128)\n",
    "        self.down1 = Down(128, 256)\n",
    "        self.down2 = Down(256, 512)\n",
    "        self.down3 = Down(512, 1024)\n",
    "        factor = 2 if bilinear else 1\n",
    "        self.down4 = Down(1024, 2048 // factor)\n",
    "        self.up1 = Up(2048, 1024 // factor, bilinear)\n",
    "        self.up2 = Up(1024, 512 // factor, bilinear)\n",
    "        self.up3 = Up(512, 256 // factor, bilinear)\n",
    "        self.up4 = Up(256, 128, bilinear)\n",
    "        self.outc = OutConv(128, n_classes)\n",
    "    def forward(self, x):\n",
    "        x1 = self.inc(x)\n",
    "        x2 = self.down1(x1)\n",
    "        x3 = self.down2(x2)\n",
    "        x4 = self.down3(x3)\n",
    "        x5 = self.down4(x4)\n",
    "        x = self.up1(x5, x4)\n",
    "        x = self.up2(x, x3)\n",
    "        x = self.up3(x, x2)\n",
    "        x = self.up4(x, x1)\n",
    "        logits = self.outc(x)\n",
    "        return logits\n",
    "\n",
    "\n",
    "model = UNet(n_channels=16, n_classes=1)\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ace38da1",
   "metadata": {},
   "outputs": [],
   "source": [
    "   # start training process\n",
    "print('running on...', device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f666cd16",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a07f4fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dice2d(pred, targs):  \n",
    "    pred = pred.squeeze()\n",
    "    targs = targs.squeeze()\n",
    "    return 2. * (pred*targs).sum() / (pred+targs).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "234cebe9-2c16-482e-9038-1c16c10cf3c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, epochs, opt, loss, batch_size):  \n",
    "    \"\"\"\n",
    "    :param model: model instance\n",
    "    :param epochs: (int) number of epochs to be trained\n",
    "    :param opt: optimizer instance\n",
    "    :param loss: loss function instance\n",
    "    :param batch_size: (int) batch size\n",
    "    :param mult: (int) augmentation factor that amplifies the data * mult\"\"\"\n",
    "\n",
    "    data_train = create_dataset(dataset=dataset,subset_size1=0,subset_size2=780,subset_size3=975,subset_size4=1815, mult=3)\n",
    "\n",
    "    data_val = create_dataset(dataset=dataset,subset_size1=780,subset_size2=975,subset_size3=1815,subset_size4=2024, mult=3)\n",
    "    train_dl = DataLoader(data_train, batch_size=16, num_workers=0,\n",
    "                        pin_memory=True) #drop_last=True)\n",
    "    val_dl = DataLoader(data_val, batch_size=16, num_workers=0,\n",
    "                        pin_memory=True) # drop_last=True)  \n",
    "    filename=\"\"   # ending of the model filename\n",
    "    best_model={}\n",
    "    best_model[\"val_loss_total\"]=100\n",
    "    best_dc={}\n",
    "    best_dc[\"val_DC\"]=0\n",
    "    # start training\n",
    "    for epoch in range(epochs):\n",
    "        model.train()\n",
    "        #metrics \n",
    "        dicec_train_acc=[]\n",
    "        FN_train=[]\n",
    "        TP_train=[]\n",
    "        FP_train=[]\n",
    "#         train_acc_total=0\n",
    "        train_loss_total = 0\n",
    "        train_ious = []\n",
    "        progress = tqdm(enumerate(train_dl), desc=\"Train Loss: \",\n",
    "                        total=len(train_dl))\n",
    "        for i, batch in progress:\n",
    "            # try:\n",
    "            x = batch['img'].float().to(device)\n",
    "            y = batch['fpt'].float().to(device)\n",
    "                                                                            \n",
    "            output = model(x)\n",
    "\n",
    "            # derive binary segmentation map from prediction\n",
    "            output_binary = np.zeros(output.shape)\n",
    "            output_binary[output.cpu().detach().numpy() >= 0] = 1\n",
    "            \n",
    "            # derive IoU values            \n",
    "            for j in range(y.shape[0]):                                       \n",
    "                z = jaccard_score(y[j].flatten().cpu().detach().numpy(),        \n",
    "                          output_binary[j][0].flatten())\n",
    "                if (np.sum(output_binary[j][0]) != 0 and\n",
    "                    np.sum(y[j].cpu().detach().numpy()) != 0):\n",
    "                    train_ious.append(z)\n",
    "                    TP_train.append((output_binary.squeeze()*y.cpu().detach().numpy().squeeze()).sum())\n",
    "                    FN_train.append(((output_binary.squeeze()==0) & (y.cpu().detach().numpy().squeeze()==1)).sum())\n",
    "                    FP_train.append(((output_binary.squeeze()==1) & (y.cpu().detach().numpy().squeeze()==0)).sum())\n",
    "                    dicec_train_acc.append(dice2d(output_binary,y.cpu().detach().numpy()))\n",
    "\n",
    "            # derive scalar binary labels on a per-image basis\n",
    "            y_bin = np.array(np.sum(y.cpu().detach().numpy(),\n",
    "                                    axis=(1,2)) != 0).astype(int)\n",
    "            pred_bin = np.array(np.sum(output_binary,\n",
    "                                      axis=(1,2,3)) != 0).astype(int)\n",
    "\n",
    "            # derive image-wise accuracy for this batch\n",
    "#             train_acc_total += accuracy_score(y_bin, pred_bin)\n",
    "            # derive loss                                                       \n",
    "            loss_epoch = loss(output, y.unsqueeze(dim=1))\n",
    "            train_loss_total += loss_epoch.item()\n",
    "            progress.set_description(\"Train Loss: {:.4f}\".format(\n",
    "                train_loss_total/(i+1)))\n",
    "\n",
    "            # learning\n",
    "            opt.zero_grad()\n",
    "            loss_epoch.backward()\n",
    "            opt.step()\n",
    "            \n",
    "        # logging\n",
    "        writer.add_scalar(\"training DC\", np.average(dicec_train_acc),epoch)\n",
    "        writer.add_scalar(\"training CE\",  np.mean(FP_train)/(np.mean(TP_train)+np.mean(FP_train)), epoch)\n",
    "        writer.add_scalar(\"training OE\",  np.mean(FN_train)/(np.mean(TP_train)+np.mean(FN_train)), epoch)                         \n",
    "        writer.add_scalar(\"training loss\", train_loss_total/(i+1), epoch)\n",
    "        writer.add_scalar(\"training iou\", np.average(train_ious), epoch)\n",
    "#         writer.add_scalar(\"training acc\", train_acc_total/(i+1), epoch)\n",
    "        writer.add_scalar('learning_rate', opt.param_groups[0]['lr'], epoch)\n",
    "        torch.cuda.empty_cache()\n",
    "\n",
    "        # evaluation\n",
    "        model.eval()\n",
    "        val_loss_total = 0\n",
    "        val_ious = []\n",
    "#         val_acc_total = 0\n",
    "        \n",
    "        dicec_eval_acc=[]\n",
    "        FN_eval=[]\n",
    "        TP_eval=[]\n",
    "        FP_eval=[]\n",
    "        \n",
    "        progress = tqdm(enumerate(val_dl), desc=\"val Loss: \",\n",
    "                        total=len(val_dl))\n",
    "                          \n",
    "        for j, batch in progress:\n",
    "            x = batch['img'].float().to(device)\n",
    "            y = batch['fpt'].float().to(device)\n",
    "            output = model(x)\n",
    "\n",
    "          # derive loss\n",
    "            loss_epoch = loss(output, y.unsqueeze(dim=1))\n",
    "            val_loss_total += loss_epoch.item()\n",
    "\n",
    "          # derive binary segmentation map from prediction\n",
    "            output_binary = np.zeros(output.shape)\n",
    "            output_binary[output.cpu().detach().numpy() >= 0] = 1\n",
    "\n",
    "          # derive IoU values\n",
    "            ious = []\n",
    "            for k in range(y.shape[0]):\n",
    "                z = jaccard_score(y[k].flatten().cpu().detach().numpy(),\n",
    "                        output_binary[k][0].flatten())\n",
    "                if (np.sum(output_binary[k][0]) != 0 and \n",
    "                    np.sum(y[k].cpu().detach().numpy()) != 0):\n",
    "                    val_ious.append(z)\n",
    "                    TP_eval.append((output_binary.squeeze()*y.cpu().detach().numpy().squeeze()).sum())\n",
    "                    FN_eval.append(((output_binary.squeeze()==0) & (y.cpu().detach().numpy().squeeze()==1)).sum())\n",
    "                    FP_eval.append(((output_binary.squeeze()==1) & (y.cpu().detach().numpy().squeeze()==0)).sum())\n",
    "                    dicec_eval_acc.append(dice2d(output_binary,y.cpu().detach().numpy()))\n",
    "                   \n",
    "          # derive scalar binary labels on a per-image basis\n",
    "            y_bin = np.array(np.sum(y.cpu().detach().numpy(),\n",
    "                                  axis=(1,2)) != 0).astype(int)\n",
    "            pred_bin = np.array(np.sum(output_binary,\n",
    "                                      axis=(1,2,3)) != 0).astype(int)\n",
    "\n",
    "          # derive image-wise accuracy for this batch\n",
    "#             val_acc_total += accuracy_score(y_bin, pred_bin)\n",
    "            \n",
    "            progress.set_description(\"val Loss: {:.4f}\".format(\n",
    "             val_loss_total/(j+1)))\n",
    "\n",
    "        # logging\n",
    "        writer.add_scalar(\"val DC\", np.average(dicec_eval_acc),epoch)\n",
    "        writer.add_scalar(\"val CE\",  np.mean(FP_eval)/(np.mean(TP_eval)+np.mean(FP_eval)), epoch)\n",
    "        writer.add_scalar(\"val OE\",  np.mean(FN_eval)/(np.mean(TP_eval)+np.mean(FN_eval)), epoch)\n",
    "        writer.add_scalar(\"val loss\", val_loss_total/(j+1), epoch)\n",
    "        writer.add_scalar(\"val iou\", np.average(val_ious), epoch)\n",
    "#         writer.add_scalar(\"val acc\", val_acc_total/(j+1), epoch)        \n",
    "        \n",
    "        print((\"Epoch {:d}: train loss={:.3f}, val loss={:.3f}, \"\n",
    "               \"train iou={:.3f}, val iou={:.3f}, \"\n",
    "                \"DC training={:.3f}, val DC={:.3f}\").format(\n",
    "                   epoch+1, train_loss_total/(i+1), val_loss_total/(j+1),\n",
    "                   np.average(train_ious), np.average(val_ious),np.average(dicec_train_acc),\n",
    "                    np.average(dicec_eval_acc)))\n",
    "\n",
    "        if (val_loss_total/(j+1))<best_model[\"val_loss_total\"]:\n",
    "            best_model[\"val_loss_total\"]=(val_loss_total/(j+1))\n",
    "            best_model[\"epoch\"]=epoch\n",
    "        if (np.average(dicec_eval_acc))>best_dc[\"val_DC\"]:\n",
    "            best_dc[\"val_DC\"]=np.average(dicec_eval_acc)\n",
    "            best_dc[\"epoch\"]=epoch\n",
    "            \n",
    "#         if epoch % 1 == 0:                                  #uncomment to save the model files\n",
    "#             torch.save(model.state_dict(),\n",
    "#             'U_Net/runs/ep{:0d}_lr{:.0e}_bs{:02d}_{:03d}_{}.model'.format(\n",
    "#                 args.ep, args.lr, args.bs, epoch, filename))\n",
    "\n",
    "        writer.flush()\n",
    "        scheduler.step(val_loss_total/(j+1))\n",
    "        torch.cuda.empty_cache()\n",
    "    print(\"best model: epoch (file): {}, val loss: {}\".format(best_model[\"epoch\"], best_model[\"val_loss_total\"]))\n",
    "    print(\"best model_dc: epoch (file): {}, val dc: {}\".format(best_dc[\"epoch\"], best_dc[\"val_DC\"]))\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a4c3320",
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup argument parser\n",
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument('-f')\n",
    "\n",
    "parser.add_argument('-ep', type=int, default=25,    \n",
    "                    help='Number of epochs')\n",
    "parser.add_argument('-bs', type=int, nargs='?',             \n",
    "                    default=16, help='Batch size')\n",
    "parser.add_argument('-lr', type=float,\n",
    "                    nargs='?', default=0.0001, help='Learning rate')\n",
    "# parser.add_argument('-mo', type=float,\n",
    "#                     nargs='?', default=0.7, help='Momentum')    #for SGD optimizer\n",
    "args = parser.parse_args()\n",
    "\n",
    "\n",
    "# setup tensorboard writer\n",
    "writer = SummaryWriter('U_Net/runs/'+\"ep{:0d}_lr{:.0e}_bs{:03d}/\".format(\n",
    "    args.ep, args.lr, args.bs))\n",
    "\n",
    "# initialize loss function\n",
    "loss = nn.BCEWithLogitsLoss()\n",
    "\n",
    "# initialize optimizer\n",
    "# opt = optim.SGD(model.parameters(), lr=args.lr, momentum=args.mo)\n",
    "opt = optim.Adam(model.parameters(), lr=args.lr)\n",
    "\n",
    "# initialize scheduler\n",
    "scheduler = optim.lr_scheduler.ReduceLROnPlateau(opt, 'min',\n",
    "                                                 factor=0.5, threshold=1e-4,\n",
    "                                                 min_lr=1e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e353adc0",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# # run training\n",
    "# model.load_state_dict(torch.load(\n",
    "#  \"U_Net/runs/ep25_lr1e-03_bs10_mo0.7.model\" , map_location=torch.device('cpu')))\n",
    "train_model(model, args.ep, opt, loss, args.bs)\n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec15d5d5",
   "metadata": {},
   "source": [
    "Eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfe299d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # model.load_state_dict(model)\n",
    "# model.load_state_dict(torch.load(\n",
    "#     \"U_net/runs/128_11_03_fixed_fixed/ep25_lr1e-04_bs16_015.model\", map_location=torch.device('cpu')))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "927c5fc1",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3a1a95a",
   "metadata": {},
   "outputs": [],
   "source": [
    "evalb=pd.read_csv(\"firescarbiobioallsizes/biobio_testf_01_02.csv\")\n",
    "evalv=pd.read_csv(\"firescarvalpoallsizes/valpo_testf_01_02.csv\")\n",
    "\n",
    "evald=pd.concat([evalv,evalb],axis=0,ignore_index=True)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "587ab902",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "np.random.seed(3)\n",
    "torch.manual_seed(3)\n",
    "\n",
    "# load data\n",
    "data_val = create_dataset(dataset=evald,subset_size1=0,subset_size2=50,subset_size3=50,subset_size4=100, mult=1)\n",
    "\n",
    "batch_size = 1 # 1 to create diagnostic images, any value otherwise\n",
    "all_dl = DataLoader(data_val, batch_size=batch_size)#, shuffle=True)\n",
    "progress = tqdm(enumerate(all_dl), total=len(all_dl))\n",
    "\n",
    "dicec_eval_acc=[]\n",
    "FN_eval=[]\n",
    "TP_eval=[]\n",
    "FP_eval=[]\n",
    "comission=[]\n",
    "omission=[]\n",
    "cont=0\n",
    "model.eval()\n",
    "\n",
    "# define loss function\n",
    "loss_fn = nn.BCEWithLogitsLoss()\n",
    "\n",
    "# run through test data\n",
    "all_ious = []\n",
    "# all_accs = []\n",
    "test_df=pd.DataFrame(columns=[\"ImgPosF\",\"iou\",\"DC\",\"CE\",\"OE\"])\n",
    "for i, batch in progress:\n",
    "    x, y = batch['img'].float().to(device), batch['fpt'].float().to(device)\n",
    "    idx = batch['idx']\n",
    "\n",
    "    output = model(x).cpu()\n",
    "\n",
    "    # obtain binary prediction map\n",
    "    pred = np.zeros(output.shape)\n",
    "    pred[output >= 0] = 1\n",
    "\n",
    "    # derive Iou score\n",
    "    cropped_iou = []\n",
    "    for j in range(y.shape[0]):\n",
    "        z = jaccard_score(y[j].flatten().cpu().detach().numpy(),\n",
    "                          pred[j][0].flatten())\n",
    "        if (np.sum(pred[j][0]) != 0 and\n",
    "            np.sum(y[j].cpu().detach().numpy()) != 0):\n",
    "            cropped_iou.append(z)       \n",
    "            \n",
    "    all_ious = [*all_ious, *cropped_iou]\n",
    "        \n",
    "    # derive scalar binary labels on a per-image basis\n",
    "    y_bin = np.array(np.sum(y.cpu().detach().numpy(),\n",
    "                            axis=(1,2)) != 0).astype(int)\n",
    "    prediction = np.array(np.sum(pred,\n",
    "                               axis=(1,2,3)) != 0).astype(int)\n",
    "    # derive image-wise accuracy for this batch\n",
    "#     all_accs.append(accuracy_score(y_bin, prediction))\n",
    "\n",
    "    # derive binary segmentation map from prediction\n",
    "    output_binary = np.zeros(output.shape)\n",
    "    output_binary[output.cpu().detach().numpy() >= 0] = 1\n",
    "\n",
    "    if batch_size == 1:\n",
    "\n",
    "        if prediction == 1 and y_bin == 1:\n",
    "            res = 'true_pos'\n",
    "        elif prediction == 0 and y_bin == 0:\n",
    "            res = 'true_neg'\n",
    "        elif prediction == 0 and y_bin == 1:\n",
    "            res = 'false_neg'\n",
    "        elif prediction == 1 and y_bin == 0:\n",
    "            res = 'false_pos'    \n",
    "        #scores fix\n",
    "        TP_eval.append((output_binary.squeeze()*y.cpu().detach().numpy().squeeze()).sum())\n",
    "        FN_eval.append(((output_binary.squeeze()==0) & (y.cpu().detach().numpy().squeeze()==1)).sum())\n",
    "        FP_eval.append(((output_binary.squeeze()==1) & (y.cpu().detach().numpy().squeeze()==0)).sum())\n",
    "        dicec_eval_acc.append(dice2d(output_binary,y.cpu().detach().numpy()))\n",
    "        test_df.loc[cont,\"OE\"]=FN_eval[cont]/(TP_eval[cont]+FN_eval[cont])\n",
    "        test_df.loc[cont,\"CE\"]=FP_eval[cont]/(TP_eval[cont]+FP_eval[cont])\n",
    "        test_df.loc[cont,\"DC\"]=dice2d(output_binary,y.cpu().detach().numpy()) \n",
    "        test_df.loc[cont,\"ImgPosF\"]=(batch['imgfile'][0].split(\"/\")[2])\n",
    "        OE=FN_eval[cont]/(TP_eval[cont]+FN_eval[cont])\n",
    "        this_iou = jaccard_score(y[0].flatten().cpu().detach().numpy(),\n",
    "                                 pred[0][0].flatten())\n",
    "        test_df.loc[i,\"iou\"]=this_iou        \n",
    "\n",
    "\n",
    "         # create plot\n",
    "        f, (ax1, ax2, ax3,ax4) = plt.subplots(1, 4, figsize=(20,20))\n",
    "        x=x.cpu()\n",
    "        y=y.cpu()\n",
    "        \n",
    "        # false color plot Image prefire\n",
    "        ax1.imshow(0.2+1.5*(np.dstack([x[0][12], x[0][11], x[0][10]])-np.min([x[0][12].numpy(),\n",
    "                            x[0][11].numpy(), x[0][10].numpy()]))/(np.max([x[0][12].numpy(),\n",
    "                            x[0][11].numpy(), x[0][10].numpy()])-np.min([x[0][12].numpy(),\n",
    "                            x[0][11].numpy(), x[0][10].numpy()])), origin='upper')\n",
    "        \n",
    "        ax1.set_title(\"ImgPreF\",fontsize=12)\n",
    "        ax1.set_xticks([])\n",
    "        ax1.set_yticks([])\n",
    "        #Image Pos-Fire\n",
    "        ax2.imshow(0.2+1.5*(np.dstack([x[0][4], x[0][3], x[0][2]])-np.min([x[0][4].numpy(), \n",
    "                            x[0][3].numpy(), x[0][2].numpy()]))/(np.max([x[0][4].numpy(),\n",
    "                            x[0][3].numpy(), x[0][2].numpy()])-np.min([x[0][4].numpy(),\n",
    "                            x[0][3].numpy(), x[0][2].numpy()])), origin='upper')\n",
    "        \n",
    "        ax2.set_title(\"ImgPosF\",fontsize=12)\n",
    "        ax2.set_xticks([])\n",
    "        ax2.set_yticks([])\n",
    "\n",
    "        # segmentation ground-truth and prediction\n",
    "        ax3.imshow(y[0], cmap='Greys_r', alpha=1)\n",
    "        ax4.imshow(pred[0][0], cmap='Greys_r', alpha=1)\n",
    "        ax3.set_title(\"Original Scar\",fontsize=12)\n",
    "        ax3.set_xticks([])\n",
    "        ax3.set_yticks([])\n",
    "        ax3.annotate(\"IoU={:.2f}\".format(this_iou), xy=(5,15), fontsize=15)\n",
    "\n",
    "        ax4.set_title({'true_pos': 'Scar Prediction: True Positive \\n  -IoU={:.2f},' \n",
    "                       '-OE={:.2f}, -CE={:.2f}, -DC={:.2F}'.format(this_iou, test_df.loc[cont,\"OE\"],test_df.loc[cont,\"CE\"],test_df.loc[cont,\"DC\"]),\n",
    "               'true_neg': 'Scar Prediction: True Negative \\n  -IoU={:.2f},' \n",
    "            '-OE={:.2f}, -CE={:.2f}, -DC={:.2F}'.format(this_iou,test_df.loc[cont,\"OE\"],test_df.loc[cont,\"CE\"],test_df.loc[cont,\"DC\"]),\n",
    "               'false_pos': 'Scar Prediction: False Positive   -IoU={:.2f},'\n",
    "             '-OE={:.2f}, -CE={:.2f}, -DC={:.2F}'.format(this_iou, test_df.loc[cont,\"OE\"],test_df.loc[cont,\"CE\"],test_df.loc[cont,\"DC\"]),\n",
    "                'false_neg': 'Scar Prediction: False Negative \\n  -IoU={:.2f},'\n",
    "            '-OE={:.2f}, -CE={:.2f}, -DC={:.2F}'.format(this_iou,test_df.loc[cont,\"OE\"], 0,test_df.loc[cont,\"DC\"])}[res],\n",
    "                      fontsize=12)\n",
    "        \n",
    "        cont+=1      \n",
    "\n",
    "        f.subplots_adjust(0.05, 0.02, 0.95, 0.9, 0.05, 0.05)\n",
    "\n",
    "        # plt.savefig(\"U_Net/output/\"+(os.path.split(batch['imgfile'][0])[1]).\\\n",
    "        #             replace('.tif', '.png').replace(':', '_'),\n",
    "        #              dpi=200)   \n",
    "        \n",
    "        plt.close()     #comment to display\n",
    "\n",
    "print('iou:', len(all_ious), np.average(all_ious))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df06b17c",
   "metadata": {
    "tags": []
   },
   "source": [
    "###### Test analysis "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "651577eb-cc2a-4238-8e34-d7e0272dd1dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# test_df.to_csv(\"U_Net/runs/test_df_128_final1103.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44d48e16-e5e8-491a-9837-d80dde55bfc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_df[\"DC\"].mean(), test_df[\"OE\"].mean(),test_df[\"CE\"].mean()"
   ]
  }
 ],
 "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}