sentinel2deeplab.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "777ebc72",
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "execution": {
     "iopub.execute_input": "2022-05-26T16:40:10.733160Z",
     "iopub.status.busy": "2022-05-26T16:40:10.732213Z",
     "iopub.status.idle": "2022-05-26T16:40:13.963941Z",
     "shell.execute_reply": "2022-05-26T16:40:13.963189Z"
    },
    "papermill": {
     "duration": 3.254819,
     "end_time": "2022-05-26T16:40:13.966238",
     "exception": false,
     "start_time": "2022-05-26T16:40:10.711419",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import torch.utils.data\n",
    "import os\n",
    "import sys\n",
    "import rasterio\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "import random\n",
    "import re\n",
    "\n",
    "from math import cos,pi\n",
    "from sklearn.metrics import f1_score, precision_score, recall_score, jaccard_score, accuracy_score, confusion_matrix\n",
    "from scipy.ndimage import morphology\n",
    "from scipy.ndimage.filters import maximum_filter1d\n",
    "from torch.nn import Module, Sequential\n",
    "from torch.nn import Conv3d, ConvTranspose3d, BatchNorm3d, MaxPool3d, AvgPool1d, Dropout3d\n",
    "from torch.nn import ReLU, Sigmoid\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.autograd import Variable\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6af29af5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-26T16:40:14.003779Z",
     "iopub.status.busy": "2022-05-26T16:40:14.003541Z",
     "iopub.status.idle": "2022-05-26T16:40:14.081472Z",
     "shell.execute_reply": "2022-05-26T16:40:14.080654Z"
    },
    "papermill": {
     "duration": 0.098845,
     "end_time": "2022-05-26T16:40:14.083380",
     "exception": false,
     "start_time": "2022-05-26T16:40:13.984535",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "root_path = '/kaggle/input/sentinel2-munich480/munich480'  # root dir dataset\n",
    "result_path = '/kaggle/working/'\n",
    "resume_path = '/kaggle/input/pretrained-deeplab/best_modelDeepLab.pth'\n",
    "result_train = 'train_resultsDeepLab.txt'\n",
    "result_validation = 'validation_resultsDeepLab.txt'\n",
    "no_of_classes = 18 #5\n",
    "workers = 8\n",
    "batch_size = 2\n",
    "h = w = 7\n",
    "n_classes = 18 #400\n",
    "LABEL_FILENAME = \"y.tif\"\n",
    "best_test_acc = 0\n",
    "loss = 'batch'\n",
    "ottimizzatore = 'sgd'\n",
    "learning_rate = 0.01\n",
    "weight_decay = 1e-6\n",
    "momentum = 0.9\n",
    "loss_weights = 'store_true'\n",
    "ignore_index = 0\n",
    "test_only = False\n",
    "sample_duration = 30\n",
    "n_epochs = 40\n",
    "\n",
    "num_folds = 1\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2bec318",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-26T16:40:14.120840Z",
     "iopub.status.busy": "2022-05-26T16:40:14.120603Z",
     "iopub.status.idle": "2022-05-26T16:40:14.128075Z",
     "shell.execute_reply": "2022-05-26T16:40:14.127284Z"
    },
    "papermill": {
     "duration": 0.028516,
     "end_time": "2022-05-26T16:40:14.130012",
     "exception": false,
     "start_time": "2022-05-26T16:40:14.101496",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def save(path, model, optimizer, **kwargs):\n",
    "    model_state = None\n",
    "    optimizer_state = None\n",
    "    if model is not None:\n",
    "        model_state = model.state_dict()\n",
    "    if optimizer is not None:\n",
    "        optimizer_state = optimizer.state_dict()\n",
    "    torch.save(\n",
    "        dict(model_state=model_state,\n",
    "             optimizer_state=optimizer_state,\n",
    "             **kwargs),\n",
    "        path\n",
    "    )\n",
    "\n",
    "\n",
    "def resume(path, model, optimizer):\n",
    "    if torch.cuda.is_available():\n",
    "        snapshot = torch.load(path)\n",
    "    else:\n",
    "        snapshot = torch.load(path, map_location=\"cpu\")\n",
    "    print(\"Loaded snapshot from\", path)\n",
    "\n",
    "    model_state = snapshot.pop('model_state', snapshot)\n",
    "    optimizer_state = snapshot.pop('optimizer_state', None)\n",
    "\n",
    "    if model is not None and model_state is not None:\n",
    "        print(\"loading model...\")\n",
    "        model.load_state_dict(model_state)\n",
    "\n",
    "    if optimizer is not None and optimizer_state is not None:\n",
    "        optimizer.load_state_dict(optimizer_state)\n",
    "    return snapshot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26507184",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-26T16:40:14.167240Z",
     "iopub.status.busy": "2022-05-26T16:40:14.167013Z",
     "iopub.status.idle": "2022-05-26T16:40:14.210446Z",
     "shell.execute_reply": "2022-05-26T16:40:14.209583Z"
    },
    "papermill": {
     "duration": 0.064755,
     "end_time": "2022-05-26T16:40:14.212554",
     "exception": false,
     "start_time": "2022-05-26T16:40:14.147799",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "LABEL_FILENAME = \"y.tif\"\n",
    "class SentinelDataset(torch.utils.data.Dataset):\n",
    "    '''\n",
    "    If the first label is for example \"1|unknown\" then this will be replaced with a 0 (zero).\n",
    "    If you want to ignore other labels, then remove them from the classes.txt file and\n",
    "    this class will assigne label 0 (zero).\n",
    "    Warning: this tecnique is not stable!\n",
    "    '''\n",
    "\n",
    "    def __init__(self, root_dir, seqlength=30, tileids=None):\n",
    "        self.root_dir = root_dir\n",
    "        self.name = os.path.basename(root_dir)\n",
    "        self.data_dirs = [d for d in os.listdir(self.root_dir) if d.startswith(\"data\")]\n",
    "        self.seqlength = seqlength\n",
    "        self.munich_format = None\n",
    "        self.src_labels = None\n",
    "        self.dst_labels = None\n",
    "        # labels read from groudtruth files (y.tif)\n",
    "        # useful field to check the available labels\n",
    "        self.unique_labels = np.array([], dtype=float)\n",
    "\n",
    "        self.b8_index = 3  # munich dataset\n",
    "        self.b4_index = 2  # munich dataset\n",
    "\n",
    "        stats = dict(\n",
    "            rejected_nopath=0,\n",
    "            rejected_length=0,\n",
    "            total_samples=0)\n",
    "\n",
    "        # statistics\n",
    "        self.samples = list()\n",
    "\n",
    "        self.ndates = list()\n",
    "\n",
    "        dirs = []\n",
    "        if tileids is None:\n",
    "            # files = os.listdir(self.data_dirs)\n",
    "            for d in self.data_dirs:\n",
    "                dirs_name = os.listdir(os.path.join(self.root_dir, d))\n",
    "                dirs_path = [os.path.join(self.root_dir, d, f) for f in dirs_name]\n",
    "                dirs.extend(dirs_path)\n",
    "        else:\n",
    "            # tileids e.g. \"tileids/train_fold0.tileids\" path of line separated tileids specifying\n",
    "            with open(os.path.join(self.root_dir, tileids), 'r') as f:\n",
    "                files = [el.replace(\"\\n\", \"\") for el in f.readlines()]\n",
    "            for d in self.data_dirs:\n",
    "                dirs_path = [os.path.join(self.root_dir, d, f) for f in files]\n",
    "                dirs.extend(dirs_path)\n",
    "\n",
    "        self.classids, self.classes = self.read_classes(os.path.join(self.root_dir, \"classes.txt\"))\n",
    "\n",
    "        for path in dirs:\n",
    "            if not os.path.exists(path):\n",
    "                stats[\"rejected_nopath\"] += 1\n",
    "                continue\n",
    "            if not os.path.exists(os.path.join(path, LABEL_FILENAME)):\n",
    "                stats[\"rejected_nopath\"] += 1\n",
    "                continue\n",
    "\n",
    "            ndates = len(get_dates(path))\n",
    "\n",
    "            if ndates < self.seqlength:\n",
    "                stats[\"rejected_length\"] += 1\n",
    "                continue  # skip shorter sequence lengths\n",
    "\n",
    "            stats[\"total_samples\"] += 1\n",
    "            self.samples.append(path)\n",
    "            self.ndates.append(ndates)\n",
    "\n",
    "        print_stats(stats)\n",
    "\n",
    "    def read_classes(self, csv):\n",
    "        with open(csv, 'r') as f:\n",
    "            classes = f.readlines()\n",
    "\n",
    "        ids = list()\n",
    "        names = list()\n",
    "        for row in classes:\n",
    "            row = row.replace(\"\\n\", \"\")\n",
    "            if '|' in row:\n",
    "                id, cl = row.split('|')\n",
    "                ids.append(int(id))\n",
    "                names.append(cl)\n",
    "\n",
    "        return ids, names\n",
    "\n",
    "    def get_image_h_w(self):\n",
    "        label, profile = read(os.path.join(self.samples[0], LABEL_FILENAME))\n",
    "        return label.shape[-2], label.shape[-1]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.samples)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "\n",
    "        # path = os.path.join(self.data_dir, self.samples[idx])\n",
    "        path = self.samples[idx]\n",
    "        if path.endswith(os.sep):\n",
    "            path = path[:-1]\n",
    "        patch_id = os.path.basename(path)\n",
    "\n",
    "        label, profile = read(os.path.join(path, LABEL_FILENAME))\n",
    "\n",
    "        profile[\"name\"] = self.samples[idx]\n",
    "\n",
    "        # unique dates sorted ascending\n",
    "        dates = get_dates(path, n=self.seqlength)\n",
    "\n",
    "        x10 = list()\n",
    "        x20 = list()\n",
    "        x60 = list()\n",
    "\n",
    "        for date in dates:\n",
    "            if self.munich_format is None:\n",
    "                self.munich_format = os.path.exists(os.path.join(path, date + \"_10m.tif\"))\n",
    "                if self.munich_format:  # munich dataset\n",
    "                    self.b8_index = 3\n",
    "                    self.b4_index = 2\n",
    "                else:  # IREA dataset\n",
    "                    self.b8_index = 6\n",
    "                    self.b4_index = 2\n",
    "            if self.munich_format:\n",
    "                x10.append(read(os.path.join(path, date + \"_10m.tif\"))[0])\n",
    "                x20.append(read(os.path.join(path, date + \"_20m.tif\"))[0])\n",
    "                x60.append(read(os.path.join(path, date + \"_60m.tif\"))[0])\n",
    "            else:\n",
    "                x10.append(read(os.path.join(path, date + \".tif\"))[0])\n",
    "\n",
    "        x10 = np.array(x10) * 1e-4\n",
    "        if self.munich_format:\n",
    "            x20 = np.array(x20) * 1e-4\n",
    "            x60 = np.array(x60) * 1e-4\n",
    "\n",
    "        # augmentation\n",
    "        # if np.random.rand() < self.augmentrate:\n",
    "        #     x10 = np.fliplr(x10)\n",
    "        #     x20 = np.fliplr(x20)\n",
    "        #     x60 = np.fliplr(x60)\n",
    "        #     label = np.fliplr(label)\n",
    "        # if np.random.rand() < self.augmentrate:\n",
    "        #     x10 = np.flipud(x10)\n",
    "        #     x20 = np.flipud(x20)\n",
    "        #     x60 = np.flipud(x60)\n",
    "        #     label = np.flipud(label)\n",
    "        # if np.random.rand() < self.augmentrate:\n",
    "        #     angle = np.random.choice([1, 2, 3])\n",
    "        #     x10 = np.rot90(x10, angle, axes=(2, 3))\n",
    "        #     x20 = np.rot90(x20, angle, axes=(2, 3))\n",
    "        #     x60 = np.rot90(x60, angle, axes=(2, 3))\n",
    "        #     label = np.rot90(label, angle, axes=(0, 1))\n",
    "\n",
    "        # replace stored ids with index in classes csv\n",
    "        label = label[0]\n",
    "        self.unique_labels = np.unique(np.concatenate([label.flatten(), self.unique_labels]))\n",
    "        new = np.zeros(label.shape, np.int)\n",
    "        for cl, i in zip(self.classids, range(len(self.classids))):\n",
    "            new[label == cl] = i\n",
    "\n",
    "        label = new\n",
    "\n",
    "        label = torch.from_numpy(label)\n",
    "        x10 = torch.from_numpy(x10)\n",
    "        if self.munich_format:\n",
    "            x20 = torch.from_numpy(x20)\n",
    "            x60 = torch.from_numpy(x60)\n",
    "\n",
    "            x20 = F.interpolate(x20, size=x10.shape[2:4])\n",
    "            x60 = F.interpolate(x60, size=x10.shape[2:4])\n",
    "\n",
    "            x = torch.cat((x10, x20, x60), 1)\n",
    "        else:\n",
    "            x = x10\n",
    "\n",
    "        # permute channels with time_series (t x c x h x w) -> (c x t x h x w)\n",
    "        x = x.permute(1, 0, 2, 3)\n",
    "\n",
    "        x = x.float()\n",
    "        label = label.long()\n",
    "\n",
    "        target_ndvi = get_all_signatures(x, label, len(self.classids), self.b4_index, self.b8_index)\n",
    "\n",
    "        return x, label, target_ndvi.float(), dates, patch_id\n",
    "\n",
    "\n",
    "def get_all_signatures(inp, target, num_cls, b4_index, b8_index):\n",
    "    \"\"\"\n",
    "    expected input having shape  (c, t, h, w) and target (h, w)\n",
    "        c = number of channels for each sentinel-2 image\n",
    "        t = number of images in the time series\n",
    "        hxw = image size\n",
    "    \"\"\"\n",
    "    c, t, h, w = inp.shape\n",
    "    output_ndvi = np.zeros((t, h, w), dtype=np.float)\n",
    "\n",
    "    # xin = torch.linspace(1, t, t)\n",
    "\n",
    "    for cls_index_ in range(0, num_cls):\n",
    "        pts = (target == cls_index_).numpy()\n",
    "        all_ndvi_x_cls = []\n",
    "        for row, yr in enumerate(pts):\n",
    "            for col, xc in enumerate(yr):\n",
    "                if xc:  # is True\n",
    "                    # if target[batch_index_, row, col].item() != cls_index_:\n",
    "                    #     print(\"error\")\n",
    "                    b8 = inp[b8_index, :, row, col]\n",
    "                    b4 = inp[b4_index, :, row, col]\n",
    "                    ndvi = (b8 - b4) / (b8 + b4)\n",
    "                    ndvi = np.nan_to_num(ndvi.numpy())\n",
    "                    # if np.isnan(ndvi).any():\n",
    "                    #     print(\"NAN in ndvi!\")\n",
    "                    all_ndvi_x_cls.append(ndvi)\n",
    "        mean_ndvi = np.zeros((t,), dtype=float)\n",
    "        if len(all_ndvi_x_cls) > 1:\n",
    "            mean_ndvi = np.mean(all_ndvi_x_cls, axis=0)\n",
    "        if len(all_ndvi_x_cls) == 1:\n",
    "            mean_ndvi = all_ndvi_x_cls[0]\n",
    "        mmax_ndvi = __max_filter1d_valid(mean_ndvi, 5)  # moving max x class\n",
    "\n",
    "        # print(\"batch\", batch_index_, \", cls\", cls_index_, \", ndvi\", mmax_ndvi)\n",
    "        # plt.plot(xin, mmax_ndvi)\n",
    "\n",
    "        output_ndvi[:, pts] = mmax_ndvi.reshape(t, 1)\n",
    "    # plt.show()\n",
    "    return torch.from_numpy(output_ndvi).float()\n",
    "\n",
    "\n",
    "def __max_filter1d_valid(a, w):\n",
    "    b = a.clip(min=0)  # transform negative elements to zero\n",
    "    return maximum_filter1d(b, size=w)\n",
    "\n",
    "\n",
    "def read(file):\n",
    "    with rasterio.open(file) as src:\n",
    "        return src.read(), src.profile\n",
    "\n",
    "\n",
    "def get_dates(path, n=None):\n",
    "    \"\"\"\n",
    "    extracts a list of unique dates from dataset sample\n",
    "\n",
    "    :param path: to dataset sample folder\n",
    "    :param n: choose n random samples from all available dates\n",
    "    :return: list of unique dates in YYYYMMDD format\n",
    "    \"\"\"\n",
    "\n",
    "    files = os.listdir(path)\n",
    "    dates = list()\n",
    "    for f in files:\n",
    "        f = f.split(\"_\")[0]\n",
    "        if len(f) == 8:  # 20160101\n",
    "            dates.append(f)\n",
    "\n",
    "    dates = set(dates)\n",
    "\n",
    "    if n is not None:\n",
    "        dates = random.sample(dates, n)\n",
    "\n",
    "    dates = list(dates)\n",
    "    dates.sort()\n",
    "    return dates\n",
    "\n",
    "\n",
    "def print_stats(stats):\n",
    "    print_lst = list()\n",
    "    for k, v in zip(stats.keys(), stats.values()):\n",
    "        print_lst.append(\"{}:{}\".format(k, v))\n",
    "    print('\\n', \", \".join(print_lst))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4377b75b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-26T16:40:14.250188Z",
     "iopub.status.busy": "2022-05-26T16:40:14.249787Z",
     "iopub.status.idle": "2022-05-26T16:49:43.454969Z",
     "shell.execute_reply": "2022-05-26T16:49:43.454083Z"
    },
    "papermill": {
     "duration": 569.245557,
     "end_time": "2022-05-26T16:49:43.476641",
     "exception": false,
     "start_time": "2022-05-26T16:40:14.231084",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"Starting loading Dataset...\")\n",
    "\n",
    "traindataset = SentinelDataset(root_path, tileids=\"tileids/train_fold0.tileids\", seqlength=sample_duration)\n",
    "traindataloader = torch.utils.data.DataLoader(\n",
    "    traindataset, batch_size=batch_size, shuffle=True, num_workers=workers, drop_last = True)\n",
    "# How to iterate on a dataloader\n",
    "for iteration, data in enumerate(traindataloader):\n",
    "    input, target, target_ndvi, _, _ = data\n",
    "    print('input temporal series with 30 images of size 13x48x48:', input.shape)\n",
    "    print('target segmentation image (batchx48x48):', target.shape)\n",
    "    print('target_ndvi containing 30 channels of size 48x48:', target_ndvi.shape)\n",
    "    break\n",
    "\n",
    "# Load test set\n",
    "testdataset = SentinelDataset(root_path, tileids=\"tileids/test_fold0.tileids\", seqlength=sample_duration)\n",
    "testdataloader = torch.utils.data.DataLoader(\n",
    "    testdataset, batch_size=batch_size, shuffle=False, num_workers=workers, drop_last = True)\n",
    "# Load validation set\n",
    "validationdataset = SentinelDataset(root_path, tileids=\"tileids/eval.tileids\", seqlength=sample_duration)\n",
    "validationdataloader = torch.utils.data.DataLoader(\n",
    "    validationdataset, batch_size=batch_size, shuffle=False, num_workers=workers, drop_last = True)\n",
    "\n",
    "numclasses = len(traindataset.classes)\n",
    "labels = list(range(numclasses))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70cc0076",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-26T16:49:43.522556Z",
     "iopub.status.busy": "2022-05-26T16:49:43.522289Z",
     "iopub.status.idle": "2022-05-26T16:49:43.532487Z",
     "shell.execute_reply": "2022-05-26T16:49:43.531697Z"
    },
    "papermill": {
     "duration": 0.036255,
     "end_time": "2022-05-26T16:49:43.534381",
     "exception": false,
     "start_time": "2022-05-26T16:49:43.498126",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ProgressBar(object):\n",
    "    DEFAULT = 'Progress: %(bar)s %(percent)3d%%'\n",
    "    FULL = '%(bar)s %(current)d/%(total)d (%(percent)3d%%) %(remaining)d to go'\n",
    "\n",
    "    def __init__(self, total, width=40, fmt=DEFAULT, symbol='=',\n",
    "                 output=sys.stderr):\n",
    "        assert len(symbol) == 1\n",
    "\n",
    "        self.total = total\n",
    "        self.width = width\n",
    "        self.symbol = symbol\n",
    "        self.output = output\n",
    "        self.fmt = re.sub(r'(?P<name>%\\(.+?\\))d',\n",
    "            r'\\g<name>%dd' % len(str(total)), fmt)\n",
    "\n",
    "        self.current = 0\n",
    "\n",
    "    def __call__(self):\n",
    "        percent = self.current / float(self.total)\n",
    "        size = int(self.width * percent)\n",
    "        remaining = self.total - self.current\n",
    "        bar = '[' + self.symbol * size + ' ' * (self.width - size) + ']'\n",
    "\n",
    "        args = {\n",
    "            'total': self.total,\n",
    "            'bar': bar,\n",
    "            'current': self.current,\n",
    "            'percent': percent * 100,\n",
    "            'remaining': remaining\n",
    "        }\n",
    "        print('\\r' + self.fmt % args, file=self.output, end='')\n",
    "\n",
    "    def done(self):\n",
    "        self.current = self.total\n",
    "        self()\n",
    "        print('', file=self.output)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f668670",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-26T16:49:43.581211Z",
     "iopub.status.busy": "2022-05-26T16:49:43.580986Z",
     "iopub.status.idle": "2022-05-26T16:49:43.606239Z",
     "shell.execute_reply": "2022-05-26T16:49:43.605322Z"
    },
    "papermill": {
     "duration": 0.051649,
     "end_time": "2022-05-26T16:49:43.608545",
     "exception": false,
     "start_time": "2022-05-26T16:49:43.556896",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class SegmentationLosses(object):\n",
    "    def __init__(self, weight=None, size_average=True, batch_average=True, ignore_index=-100, cuda=False):\n",
    "        self.ignore_index = ignore_index\n",
    "        self.weight = weight\n",
    "        self.size_average = size_average\n",
    "        self.batch_average = batch_average\n",
    "        self.cuda = cuda\n",
    "        self.criterion_xent = None\n",
    "        self.criterion_mse = None\n",
    "        self.beta = 0.9999  # for class balanced loss\n",
    "\n",
    "    def build_loss(self, mode='ce'):\n",
    "        loss_func = None\n",
    "        \"\"\"Choices: ['ce' | 'focal' | 'ndvi' | 'batch']\"\"\"\n",
    "        if mode == 'ce':\n",
    "            self.criterion_xent = nn.CrossEntropyLoss(weight=self.weight, ignore_index=self.ignore_index)\n",
    "            loss_func = self.cross_entropy_loss\n",
    "        elif mode == 'focal':\n",
    "            self.criterion_xent = nn.CrossEntropyLoss(weight=self.weight, ignore_index=self.ignore_index)\n",
    "            loss_func = self.focal_loss\n",
    "        elif mode == 'ndvi':\n",
    "            self.criterion_xent = nn.CrossEntropyLoss(weight=self.weight, ignore_index=self.ignore_index)\n",
    "            self.criterion_mse = nn.MSELoss()  # Default reduction: 'mean' (reduction='sum')\n",
    "            # self.criterion_mse = nn.L1Loss()\n",
    "            loss_func = self.ndvi_loss\n",
    "        elif mode == 'batch':\n",
    "            # weights are computed inside a batch:\n",
    "            # https://medium.com/gumgum-tech/handling-class-imbalance-by-introducing-sample-weighting-in-the-loss-function-3bdebd8203b4\n",
    "            self.criterion_xent = nn.CrossEntropyLoss(weight=self.weight, ignore_index=self.ignore_index)\n",
    "            loss_func = self.class_balanced_loss\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "\n",
    "        if self.cuda:\n",
    "            self.criterion_xent = self.criterion_xent.cuda()\n",
    "            if self.criterion_mse is not None:\n",
    "                self.criterion_mse = self.criterion_mse.cuda()\n",
    "\n",
    "        return loss_func\n",
    "\n",
    "    def cross_entropy_loss(self, logit, target):\n",
    "        n, c, h, w = logit.size()\n",
    "\n",
    "        loss = self.criterion_xent(logit, target.long())\n",
    "\n",
    "        if self.batch_average:\n",
    "            loss /= n\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def focal_loss(self, logit, target, gamma=2, alpha=0.5):\n",
    "        n, c, h, w = logit.size()\n",
    "\n",
    "        logpt = -self.criterion_xent(logit, target.long())\n",
    "        pt = torch.exp(logpt)\n",
    "        if alpha is not None:\n",
    "            logpt *= alpha\n",
    "        loss = -((1 - pt) ** gamma) * logpt\n",
    "\n",
    "        if self.batch_average:\n",
    "            loss /= n\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def ndvi_loss(self, ndvi_features, logit, target_cls, ndvi_target, samples_per_cls, weight_mse=1.0):\n",
    "        n, c, h, w = logit.size()\n",
    "        assert (not torch.isnan(ndvi_target).any())\n",
    "\n",
    "        # Effective Number of Samples (ENS)\n",
    "        if not samples_per_cls == None:\n",
    "            assert (samples_per_cls.shape[0] == c)\n",
    "            weights = (1.0 - self.beta) / (1.0 - torch.pow(self.beta, samples_per_cls.float()))\n",
    "            weights[weights == float('inf')] = 0\n",
    "            weights = weights / torch.sum(weights) * c  # wights in the range [0, c]\n",
    "            self.criterion_xent.weight = weights\n",
    "\n",
    "        loss_xent = self.criterion_xent(logit, target_cls.long())\n",
    "        loss_mse = self.criterion_mse(ndvi_features, ndvi_target)\n",
    "        loss_mse *= weight_mse\n",
    "        loss = loss_xent + loss_mse\n",
    "\n",
    "        if self.batch_average:\n",
    "            loss /= n\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def class_balanced_loss(self, logit, target, samples_per_cls, weight_type='ENS'):\n",
    "        \"\"\"Compute the Class Balanced Loss between `logits` and the ground truth `labels`.\n",
    "        Class Balanced Loss: ((1-beta)/(1-beta^n))*Loss(labels, logits)\n",
    "        where Loss is one of the standard losses used for Neural Networks.\"\"\"\n",
    "        # Starting point:\n",
    "        # https://medium.com/gumgum-tech/handling-class-imbalance-by-introducing-sample-weighting-in-the-loss-function-3bdebd8203b4\n",
    "        n, c, h, w = logit.size()\n",
    "        # beta = (self.total_num_samples - 1) / self.total_num_samples\n",
    "\n",
    "        assert(samples_per_cls.shape[0] == c)\n",
    "\n",
    "        # compute weights for each minibatch\n",
    "        if weight_type == 'ENS':\n",
    "            # Effective Number of Samples (ENS)\n",
    "            weights = (1.0 - self.beta) / (1.0 - torch.pow(self.beta, samples_per_cls.float()))\n",
    "            weights[weights == float('inf')] = 0\n",
    "        elif weight_type == 'ISNS':\n",
    "            # Inverse of Square Root of Number of Samples (ISNS)\n",
    "            weights = 1.0 / torch.sqrt(torch.tensor([2, 1000, 1, 20000, 500]).float())\n",
    "        else:\n",
    "            # Inverse of Number of Samples (INS)\n",
    "            weights = 1.0 / torch.tensor([2, 1000, 1, 20000, 500]).float()\n",
    "\n",
    "        weights = weights / torch.sum(weights) * c  # wights in the range [0, c]\n",
    "\n",
    "        self.criterion_xent.weight = weights\n",
    "        loss = self.criterion_xent(logit, target.long())\n",
    "        # print(\"loss weights:\", self.criterion_xent.weight)\n",
    "\n",
    "        if self.batch_average:\n",
    "            loss /= n\n",
    "\n",
    "        return loss\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84e987ba",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-26T16:49:43.654196Z",
     "iopub.status.busy": "2022-05-26T16:49:43.653973Z",
     "iopub.status.idle": "2022-05-26T16:49:43.665544Z",
     "shell.execute_reply": "2022-05-26T16:49:43.664840Z"
    },
    "papermill": {
     "duration": 0.036283,
     "end_time": "2022-05-26T16:49:43.667387",
     "exception": false,
     "start_time": "2022-05-26T16:49:43.631104",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Meter(object):\n",
    "    '''Meters provide a way to keep track of important statistics in an online manner.\n",
    "    This class is abstract, but provides a standard interface for all meters to follow.\n",
    "    '''\n",
    "\n",
    "    def reset(self):\n",
    "        '''Resets the meter to default settings.'''\n",
    "        pass\n",
    "\n",
    "    def add(self, value):\n",
    "        '''Log a new value to the meter\n",
    "        Args:\n",
    "            value: Next restult to include.\n",
    "        '''\n",
    "        pass\n",
    "\n",
    "    def value(self):\n",
    "        '''Get the value of the meter in the current state.'''\n",
    "        pass\n",
    "\n",
    "\n",
    "class AverageValueMeter(Meter):\n",
    "    def __init__(self):\n",
    "        super(AverageValueMeter, self).__init__()\n",
    "        self.reset()\n",
    "        self.val = 0\n",
    "\n",
    "    def add(self, value, n=1):\n",
    "        self.val = value\n",
    "        self.sum += value\n",
    "        self.var += value * value\n",
    "        self.n += n\n",
    "\n",
    "        if self.n == 0:\n",
    "            self.mean, self.std = np.nan, np.nan\n",
    "        elif self.n == 1:\n",
    "            self.mean = 0.0 + self.sum  # This is to force a copy in torch/numpy\n",
    "            self.std = np.inf\n",
    "            self.mean_old = self.mean\n",
    "            self.m_s = 0.0\n",
    "        else:\n",
    "            self.mean = self.mean_old + (value - n * self.mean_old) / float(self.n)\n",
    "            self.m_s += (value - self.mean_old) * (value - self.mean)\n",
    "            self.mean_old = self.mean\n",
    "            self.std = np.sqrt(self.m_s / (self.n - 1.0))\n",
    "\n",
    "    def value(self):\n",
    "        return self.mean, self.std\n",
    "\n",
    "    def reset(self):\n",
    "        self.n = 0\n",
    "        self.sum = 0.0\n",
    "        self.var = 0.0\n",
    "        self.val = 0.0\n",
    "        self.mean = np.nan\n",
    "        self.mean_old = 0.0\n",
    "        self.m_s = 0.0\n",
    "        self.std = np.nan\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6691846e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-26T16:49:43.711526Z",
     "iopub.status.busy": "2022-05-26T16:49:43.711089Z",
     "iopub.status.idle": "2022-05-26T16:49:43.732715Z",
     "shell.execute_reply": "2022-05-26T16:49:43.731841Z"
    },
    "papermill": {
     "duration": 0.046254,
     "end_time": "2022-05-26T16:49:43.734865",
     "exception": false,
     "start_time": "2022-05-26T16:49:43.688611",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class BaseObject(nn.Module):\n",
    "\n",
    "    def __init__(self, name=None):\n",
    "        super().__init__()\n",
    "        self._name = name\n",
    "\n",
    "    @property\n",
    "    def __name__(self):\n",
    "        if self._name is None:\n",
    "            name = self.__class__.__name__\n",
    "            s1 = re.sub('(.)([A-Z][a-z]+)', r'\\1_\\2', name)\n",
    "            return re.sub('([a-z0-9])([A-Z])', r'\\1_\\2', s1).lower()\n",
    "        else:\n",
    "            return self._name\n",
    "\n",
    "\n",
    "class Metric(BaseObject):\n",
    "    pass\n",
    "\n",
    "class ConfusionMatrix(Metric):\n",
    "\n",
    "    def __init__(self, labels, ignore_class=None, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.ignore_class = ignore_class  # the class index to be removed\n",
    "        self.labels = labels\n",
    "        self.n_classes = len(self.labels)\n",
    "        if self.ignore_class is not None:\n",
    "            self.matrix = np.zeros((self.n_classes-1, self.n_classes-1))\n",
    "        else:\n",
    "            self.matrix = np.zeros((self.n_classes, self.n_classes))\n",
    "\n",
    "        # self.pred = []\n",
    "        # self.targ = []\n",
    "\n",
    "    def get_labels(self):\n",
    "        if self.ignore_class is not None:\n",
    "            return np.delete(self.labels, self.ignore_class)\n",
    "        return self.labels\n",
    "\n",
    "    def forward(self, y_pr, y_gt):\n",
    "        # sklearn.metrics\n",
    "        pred = y_pr.view(-1).cpu().detach().tolist()\n",
    "        targ = y_gt.view(-1).cpu().detach().tolist()\n",
    "\n",
    "        # To format the matrix\n",
    "        # https://medium.com/@dtuk81/confusion-matrix-visualization-fc31e3f30fea\n",
    "        # confusion_matrix(y_true, y_pred)\n",
    "        # array([[2, 0, 0],  # two zeros were predicted as zeros\n",
    "        #        [0, 0, 1],  # one 1 was predicted as 2\n",
    "        #        [1, 0, 2]])  # two 2s were predicted as 2, and one 2 was 0\n",
    "        matrix = confusion_matrix(targ, pred, labels=self.labels)\n",
    "\n",
    "        if self.ignore_class is not None:\n",
    "            matrix = np.delete(matrix, self.ignore_class, 0)  # remove the row\n",
    "            matrix = np.delete(matrix, self.ignore_class, 1)  # remove the column\n",
    "\n",
    "        self.matrix = np.add(self.matrix, matrix)\n",
    "\n",
    "        results_vec = {\"labels\": self.get_labels(), \"confusion matrix\": self.matrix}\n",
    "\n",
    "        total = np.sum(self.matrix)\n",
    "        true_positive = np.diag(self.matrix)\n",
    "        sum_rows = np.sum(self.matrix, axis=0)\n",
    "        sum_cols = np.sum(self.matrix, axis=1)\n",
    "        false_positive = sum_rows - true_positive\n",
    "        false_negative = sum_cols - true_positive\n",
    "        # calculate accuracy\n",
    "        overall_accuracy = np.sum(true_positive) / total\n",
    "        results_scalar = {\"OA\": overall_accuracy}\n",
    "\n",
    "        # calculate Cohen Kappa (https://en.wikipedia.org/wiki/Cohen%27s_kappa)\n",
    "        p0 = np.sum(true_positive) / total\n",
    "        pc = np.sum(sum_rows * sum_cols) / total ** 2\n",
    "        kappa = (p0 - pc) / (1 - pc)\n",
    "        results_scalar[\"Kappa\"] = kappa\n",
    "\n",
    "        # Per class recall, prec and F1\n",
    "        recall = true_positive / (sum_cols + 1e-12)\n",
    "        results_vec[\"R\"] = recall\n",
    "        precision = true_positive / (sum_rows + 1e-12)\n",
    "        results_vec[\"P\"] = precision\n",
    "        f1 = (2 * precision * recall) / ((precision + recall) + 1e-12)\n",
    "        results_vec[\"F1\"] = f1\n",
    "\n",
    "        # Just in case we get a division by 0, ignore/hide the error\n",
    "        with np.errstate(divide='ignore', invalid='ignore'):\n",
    "            iou = true_positive / (true_positive + false_positive + false_negative)\n",
    "            results_vec[\"IoU\"] = iou\n",
    "            results_scalar[\"mIoU\"] = np.nanmean(iou)\n",
    "\n",
    "        # Per class accuracy\n",
    "        cl_acc = true_positive / (sum_cols + 1e-12)\n",
    "        results_vec[\"Acc\"] = cl_acc\n",
    "\n",
    "        # weighted measures\n",
    "        prob_c = sum_rows / total\n",
    "        prob_r = sum_cols / total\n",
    "        recall_weighted = np.inner(recall, prob_r)\n",
    "        results_scalar[\"wR\"] = recall_weighted\n",
    "        precision_weighted = np.inner(precision, prob_r)\n",
    "        results_scalar[\"wP\"] = precision_weighted\n",
    "        f1_weighted = 2 * (recall_weighted * precision_weighted) / (recall_weighted + precision_weighted)\n",
    "        results_scalar[\"wF1\"] = f1_weighted\n",
    "        random_accuracy = np.inner(prob_c, prob_r)\n",
    "        results_scalar[\"RAcc\"] = random_accuracy\n",
    "\n",
    "        return results_vec, results_scalar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86f45adc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-26T16:49:43.777564Z",
     "iopub.status.busy": "2022-05-26T16:49:43.777095Z",
     "iopub.status.idle": "2022-05-26T16:49:43.796032Z",
     "shell.execute_reply": "2022-05-26T16:49:43.795387Z"
    },
    "papermill": {
     "duration": 0.042389,
     "end_time": "2022-05-26T16:49:43.797778",
     "exception": false,
     "start_time": "2022-05-26T16:49:43.755389",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def write_signatures(fcn, target, output, target_ndvi, out_ndvi, dates, patch_id, set_name):\n",
    "    # target_ndvi = torch.randint(0, nc, size=(b, s, h, w))\n",
    "    # patch_id = ('11', '12')\n",
    "    # dates = [('20200101', '20200301'), ('20200102', '20200302'), ('20200103', '20200303'), ('20200104', '20200304'),\n",
    "    #          ('20200105', '20200305')]\n",
    "    nb, nt, h, w = target_ndvi.shape\n",
    "    assert(len(dates) == nt)\n",
    "    assert(len(dates[0]) == nb)\n",
    "\n",
    "    # winner class for each pixel\n",
    "    winners = torch.softmax(output, dim=1).argmax(dim=1)\n",
    "\n",
    "    # bias = fcn.conv5out.bias\n",
    "    weight = fcn.conv5out.weight\n",
    "\n",
    "    for idx, patch_name in enumerate(patch_id):\n",
    "        # print('patch_name:', patch_name)\n",
    "        with open(os.path.join(result_path, set_name + '_patch_' + patch_name + \".txt\"), 'w') as f:\n",
    "            f.write('class, output, type_ndvi, x, y, ')\n",
    "            f.write(', '.join(map(str, [e[idx] for e in dates])))\n",
    "            f.write(\"\\n\")\n",
    "            for y in range(0, h):\n",
    "                for x in range(0, w):\n",
    "                    target_idx = target[idx, y, x].item()\n",
    "                    output_idx = winners[idx, y, x].item()\n",
    "                    f.write('%d, %d, target, %d, %d, ' % (target_idx, output_idx, x, y))\n",
    "                    f.write(', '.join(map(str, target_ndvi[idx, :, y, x].data.tolist())))\n",
    "                    f.write(\"\\n\")\n",
    "                    f.write('%d, %d, predic, %d, %d, ' % (target_idx, output_idx, x, y))\n",
    "                    f.write(', '.join(map(str, out_ndvi[idx, :, y, x].data.tolist())))\n",
    "                    f.write(\"\\n\")\n",
    "                    f.write('%d, %d, cls_ai, output, %d, %d, ' % (target_idx, output_idx, x, y))\n",
    "                    cai = class_activations(out_ndvi, weight, target_idx, idx, y, x)\n",
    "                    f.write(', '.join(map(str, cai.data.tolist())))\n",
    "                    f.write(\"\\n\")\n",
    "            # f.write(\"\\n\")\n",
    "            \n",
    "def get_all_signatures(inp, target, num_cls, b4_index, b8_index):\n",
    "    \"\"\"\n",
    "    expected input having shape  (c, t, h, w) and target (h, w)\n",
    "        c = number of channels for each sentinel-2 image\n",
    "        t = number of images in the time series\n",
    "        hxw = image size\n",
    "    \"\"\"\n",
    "    c, t, h, w = inp.shape\n",
    "    output_ndvi = np.zeros((t, h, w), dtype=np.float)\n",
    "\n",
    "    # xin = torch.linspace(1, t, t)\n",
    "\n",
    "    for cls_index_ in range(0, num_cls):\n",
    "        pts = (target == cls_index_).numpy()\n",
    "        all_ndvi_x_cls = []\n",
    "        for row, yr in enumerate(pts):\n",
    "            for col, xc in enumerate(yr):\n",
    "                if xc:  # is True\n",
    "                    # if target[batch_index_, row, col].item() != cls_index_:\n",
    "                    #     print(\"error\")\n",
    "                    b8 = inp[b8_index, :, row, col]\n",
    "                    b4 = inp[b4_index, :, row, col]\n",
    "                    ndvi = (b8 - b4) / (b8 + b4)\n",
    "                    ndvi = np.nan_to_num(ndvi.numpy())\n",
    "                    # if np.isnan(ndvi).any():\n",
    "                    #     print(\"NAN in ndvi!\")\n",
    "                    all_ndvi_x_cls.append(ndvi)\n",
    "        mean_ndvi = np.zeros((t,), dtype=float)\n",
    "        if len(all_ndvi_x_cls) > 1:\n",
    "            mean_ndvi = np.mean(all_ndvi_x_cls, axis=0)\n",
    "        if len(all_ndvi_x_cls) == 1:\n",
    "            mean_ndvi = all_ndvi_x_cls[0]\n",
    "        mmax_ndvi = __max_filter1d_valid(mean_ndvi, 5)  # moving max x class\n",
    "\n",
    "        # print(\"batch\", batch_index_, \", cls\", cls_index_, \", ndvi\", mmax_ndvi)\n",
    "        # plt.plot(xin, mmax_ndvi)\n",
    "\n",
    "        output_ndvi[:, pts] = mmax_ndvi.reshape(t, 1)\n",
    "    # plt.show()\n",
    "    return torch.from_numpy(output_ndvi).float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa80d539",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-26T16:49:43.840359Z",
     "iopub.status.busy": "2022-05-26T16:49:43.839902Z",
     "iopub.status.idle": "2022-05-26T16:49:43.885725Z",
     "shell.execute_reply": "2022-05-26T16:49:43.884987Z"
    },
    "papermill": {
     "duration": 0.070276,
     "end_time": "2022-05-26T16:49:43.888412",
     "exception": false,
     "start_time": "2022-05-26T16:49:43.818136",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def adjust_classes_weights(cur_epoch, curr_iter, num_iter_x_epoch, tot_epochs, start_w, descending=True):\n",
    "    current_iter = curr_iter + cur_epoch * num_iter_x_epoch\n",
    "    max_iter = tot_epochs * num_iter_x_epoch\n",
    "    # a = 0.75 - current_iter / (2 * max_iter)  # from 0.75 to 0.25\n",
    "    a = 1 - current_iter / (max_iter)  # from 0 to 1\n",
    "    if not descending:\n",
    "        a = 1 - a    # from 0.25 to 0.75\n",
    "\n",
    "    w = a * start_w + 1 - a\n",
    "    return w\n",
    "def adjust_classes_weights(cur_epoch, curr_iter, num_iter_x_epoch, tot_epochs, start_w, descending=True):\n",
    "    current_iter = curr_iter + cur_epoch * num_iter_x_epoch\n",
    "    max_iter = tot_epochs * num_iter_x_epoch\n",
    "    # a = 0.75 - current_iter / (2 * max_iter)  # from 0.75 to 0.25\n",
    "    a = 1 - current_iter / (max_iter)  # from 0 to 1\n",
    "    if not descending:\n",
    "        a = 1 - a    # from 0.25 to 0.75\n",
    "\n",
    "    w = a * start_w + 1 - a\n",
    "    return w\n",
    "\n",
    "\n",
    "def adjust_learning_rate(cur_epoch, curr_iter, num_iter_x_epoch, tot_epochs, start_lr, lr_decay='cos'):\n",
    "    current_iter = curr_iter + cur_epoch * num_iter_x_epoch\n",
    "    max_iter = tot_epochs * num_iter_x_epoch\n",
    "\n",
    "    if lr_decay == 'cos':\n",
    "        lr = start_lr * (1 + cos(pi * (current_iter) / (max_iter))) / 2\n",
    "    # elif lr_decay == 'step':\n",
    "    #     lr = start_lr * (0.1 ** (cur_epoch // args.step))\n",
    "    elif lr_decay == 'linear':\n",
    "        lr = start_lr * (1 - (current_iter) / (max_iter))\n",
    "    # elif lr_decay == 'schedule':\n",
    "    #     count = sum([1 for s in args.schedule if s <= cur_epoch])\n",
    "    #     lr = start_lr * pow(args.gamma, count)\n",
    "    else:\n",
    "        raise ValueError('Unknown lr mode {}'.format(lr_decay))\n",
    "\n",
    "    return lr\n",
    "    # for param_group in optimizer.param_groups:\n",
    "    #     param_group['lr'] = lr\n",
    "\n",
    "\n",
    "def _take_channels(*xs, ignore_channels=None):\n",
    "    if ignore_channels is None:\n",
    "        return xs\n",
    "    else:\n",
    "        channels = [channel for channel in range(xs[0].shape[1]) if channel not in ignore_channels]\n",
    "        xs = [torch.index_select(x, dim=1, index=torch.tensor(channels).to(x.device)) for x in xs]\n",
    "        return xs\n",
    "\n",
    "\n",
    "def _remove_index(np_matrix, ignore_index=-100):\n",
    "    if ignore_index == -100:\n",
    "        return np_matrix\n",
    "    else:\n",
    "        m = np.delete(np_matrix, ignore_index, 0) # axes=0 -> delete row index\n",
    "        m = np.delete(m, ignore_index, 1) # axes=1 -> delete col index\n",
    "        return m\n",
    "\n",
    "\n",
    "def _threshold(x, threshold=None):\n",
    "    if threshold is not None:\n",
    "        return (x > threshold).type(x.dtype)\n",
    "    else:\n",
    "        return x\n",
    "\n",
    "\n",
    "def iou(pr, gt, labels, threshold=None, ignore_channels=None):\n",
    "    \"\"\"Computes the IoU and mean IoU.\n",
    "    The mean computation ignores NaN elements of the IoU array.\n",
    "    Returns:\n",
    "        Tuple: (IoU, mIoU). The first output is the per class IoU,\n",
    "        for K classes it's numpy.ndarray with K elements. The second output,\n",
    "        is the mean IoU.\n",
    "    \"\"\"\n",
    "    # Dimensions check\n",
    "    assert pr.size(0) == gt.size(0), \\\n",
    "        'number of targets and predicted outputs do not match'\n",
    "    assert pr.dim() == gt.dim(), \\\n",
    "        \"predictions and targets must be of dimension (N, H, W)\"\n",
    "\n",
    "    pr = _threshold(pr, threshold=threshold)\n",
    "    pr, gt = _take_channels(pr, gt, ignore_channels=ignore_channels)\n",
    "\n",
    "    # sklearn.metrics\n",
    "    pr1d = pr.view(-1).cpu().detach().numpy()\n",
    "    gt1d = gt.view(-1).cpu().detach().numpy()\n",
    "    score = jaccard_score(gt1d, pr1d, labels=labels, average=None)\n",
    "\n",
    "    return score\n",
    "    # conf_metric = ConfusionMatrix(num_classes, normalized)\n",
    "    # conf_metric.add(pr.view(-1), gt.view(-1))\n",
    "    # conf_matrix = conf_metric.value()\n",
    "    # if ignore_channels is not None:\n",
    "    #     conf_matrix[:, ignore_channels] = 0\n",
    "    #     conf_matrix[ignore_channels, :] = 0\n",
    "    # true_positive = np.diag(conf_matrix)\n",
    "    # false_positive = np.sum(conf_matrix, 0) - true_positive\n",
    "    # false_negative = np.sum(conf_matrix, 1) - true_positive\n",
    "    #\n",
    "    # # Just in case we get a division by 0, ignore/hide the error\n",
    "    # with np.errstate(divide='ignore', invalid='ignore'):\n",
    "    #     iou = true_positive / (true_positive + false_positive + false_negative)\n",
    "    #\n",
    "    # return iou, np.nanmean(iou)\n",
    "\n",
    "\n",
    "def f_score(pr, gt, labels, threshold=None, ignore_channels=None):\n",
    "    \"\"\"Calculate F-score between ground truth and prediction\n",
    "    Args:\n",
    "        pr (torch.Tensor): predicted tensor\n",
    "        gt (torch.Tensor):  ground truth tensor\n",
    "        threshold: threshold for outputs binarization\n",
    "    Returns:\n",
    "        numpy array float: F score\n",
    "    \"\"\"\n",
    "\n",
    "    pr = _threshold(pr, threshold=threshold)\n",
    "    pr, gt = _take_channels(pr, gt, ignore_channels=ignore_channels)\n",
    "\n",
    "    # sklearn.metrics\n",
    "    pr1d = pr.view(-1).cpu().detach().numpy()\n",
    "    gt1d = gt.view(-1).cpu().detach().numpy()\n",
    "    score = f1_score(gt1d, pr1d, labels=labels, average=None)\n",
    "\n",
    "    # tp = torch.sum(gt * pr)\n",
    "    # fp = torch.sum(pr) - tp\n",
    "    # fn = torch.sum(gt) - tp\n",
    "    #\n",
    "    # score = ((1 + beta ** 2) * tp + eps) \\\n",
    "    #         / ((1 + beta ** 2) * tp + beta ** 2 * fn + fp + eps)\n",
    "\n",
    "    return score\n",
    "\n",
    "\n",
    "def accuracy(pr, gt, labels, threshold=None, ignore_channels=None, ignore_index=-100):\n",
    "    \"\"\"Calculate accuracy score between ground truth and prediction\n",
    "    Args:\n",
    "        pr (torch.Tensor): predicted tensor\n",
    "        gt (torch.Tensor):  ground truth tensor\n",
    "        eps (float): epsilon to avoid zero division\n",
    "        threshold: threshold for outputs binarization\n",
    "    Returns:\n",
    "        float: precision score\n",
    "    \"\"\"\n",
    "    pr = _threshold(pr, threshold=threshold)\n",
    "    pr, gt = _take_channels(pr, gt, ignore_channels=ignore_channels)\n",
    "\n",
    "    # sklearn.metrics\n",
    "    pr1d = pr.view(-1).cpu().detach().numpy()\n",
    "    gt1d = gt.view(-1).cpu().detach().numpy()\n",
    "\n",
    "    matrix = confusion_matrix(gt1d, pr1d, labels=labels)\n",
    "    matrix = _remove_index(matrix, ignore_index)\n",
    "    diagonal = matrix.diagonal()\n",
    "    score_per_class = np.nan_to_num(diagonal / matrix.sum(axis=1))\n",
    "\n",
    "    score = diagonal.sum() / matrix.sum()\n",
    "\n",
    "    # corrects = (gt == pr).float()\n",
    "    # score = corrects.sum() / float(corrects.numel())\n",
    "\n",
    "    return score_per_class, score\n",
    "\n",
    "\n",
    "def precision(pr, gt, labels, threshold=None, ignore_channels=None):\n",
    "    \"\"\"Calculate precision score between ground truth and prediction\n",
    "    Args:\n",
    "        pr (torch.Tensor): predicted tensor\n",
    "        gt (torch.Tensor):  ground truth tensor\n",
    "        eps (float): epsilon to avoid zero division\n",
    "        threshold: threshold for outputs binarization\n",
    "    Returns:\n",
    "        numpy array: precision score\n",
    "    \"\"\"\n",
    "    pr = _threshold(pr, threshold=threshold)\n",
    "    pr, gt = _take_channels(pr, gt, ignore_channels=ignore_channels)\n",
    "\n",
    "    # sklearn.metrics\n",
    "    pr1d = pr.view(-1).cpu().detach().numpy()\n",
    "    gt1d = gt.view(-1).cpu().detach().numpy()\n",
    "    score = precision_score(gt1d, pr1d, labels=labels, average=None)\n",
    "\n",
    "    return score\n",
    "\n",
    "\n",
    "def recall(pr, gt, labels, threshold=None, ignore_channels=None):\n",
    "    \"\"\"Calculate Recall between ground truth and prediction\n",
    "    Args:\n",
    "        pr (torch.Tensor): A list of predicted elements\n",
    "        gt (torch.Tensor):  A list of elements that are to be predicted\n",
    "        eps (float): epsilon to avoid zero division\n",
    "        threshold: threshold for outputs binarization\n",
    "    Returns:\n",
    "        numpy array float: recall score\n",
    "    \"\"\"\n",
    "\n",
    "    pr = _threshold(pr, threshold=threshold)\n",
    "    pr, gt = _take_channels(pr, gt, ignore_channels=ignore_channels)\n",
    "\n",
    "    # sklearn.metrics\n",
    "    pr1d = pr.view(-1).cpu().detach().numpy()\n",
    "    gt1d = gt.view(-1).cpu().detach().numpy()\n",
    "    score = recall_score(gt1d, pr1d, labels=labels, average=None)  # per class measure\n",
    "\n",
    "    return score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f9710ef",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-26T16:49:43.972259Z",
     "iopub.status.busy": "2022-05-26T16:49:43.971901Z",
     "iopub.status.idle": "2022-05-26T16:49:43.976437Z",
     "shell.execute_reply": "2022-05-26T16:49:43.975615Z"
    },
    "papermill": {
     "duration": 0.054639,
     "end_time": "2022-05-26T16:49:43.981610",
     "exception": false,
     "start_time": "2022-05-26T16:49:43.926971",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "482c920b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-26T16:49:44.059594Z",
     "iopub.status.busy": "2022-05-26T16:49:44.059356Z",
     "iopub.status.idle": "2022-05-26T16:49:44.169340Z",
     "shell.execute_reply": "2022-05-26T16:49:44.168443Z"
    },
    "papermill": {
     "duration": 0.146486,
     "end_time": "2022-05-26T16:49:44.171364",
     "exception": false,
     "start_time": "2022-05-26T16:49:44.024878",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def compute_train_weights(train_loader):\n",
    "    beta = 0.9\n",
    "    samples_per_cls = torch.zeros(n_classes)\n",
    "    for batch_idx, data in enumerate(train_loader):\n",
    "        inputs, targets, _, _, _ = data\n",
    "        for cls in range(n_classes):\n",
    "            samples_per_cls[cls] += torch.sum(targets == cls)\n",
    "    max_occ = torch.max(samples_per_cls)\n",
    "    weights = torch.FloatTensor(max_occ / samples_per_cls)\n",
    "    # max_occ = torch.max(weights)\n",
    "    # weights = torch.FloatTensor(weights / max_occ)\n",
    "    if torch.cuda.is_available():\n",
    "        weights = weights.cuda()\n",
    "\n",
    "    return weights\n",
    "\n",
    "def train_epoch(dataloader, network, optimizer, loss, ep, loss_cls, cls_weights=None):\n",
    "    num_processed_samples = 0\n",
    "    num_train_samples = len(dataloader.dataset)\n",
    "    labels = list(range(numclasses))\n",
    "\n",
    "    conf_mat_metrics = ConfusionMatrix(labels, ignore_class=ignore_index)\n",
    "    num_cls = 18\n",
    "    batch_size = 18\n",
    "    labels = list(range(num_cls))\n",
    "    conf_mat = ConfusionMatrix(labels, ignore_class=0)\n",
    "\n",
    "    for e in range(5):\n",
    "        target = torch.randint(num_cls, (batch_size,))\n",
    "        output = torch.rand(batch_size, num_cls)\n",
    "        pred = torch.argmax(output, dim=1)\n",
    "        metrics_v, metrics_s = conf_mat(pred, target)\n",
    "        for k, v in metrics_v.items():\n",
    "            print(k, v)\n",
    "        for k, v in metrics_s.items():\n",
    "            print(k, v)\n",
    "        print()\n",
    "    loss_measure = AverageValueMeter()\n",
    "    am = AverageValueMeter()\n",
    "    am.add(numpy.array([1, 0.5, 0.2, 0.1, 0.8]))\n",
    "    mean, std = am.value()\n",
    "    print(mean, std, am.sum)\n",
    "\n",
    "    am.add(numpy.array([0.5, 0.5, 0.5, 0.5, 0.5]))\n",
    "    mean, std = am.value()\n",
    "    print(mean, std, am.sum)\n",
    "    var_learning_rate = learning_rate\n",
    "    for iteration, data in enumerate(dataloader):\n",
    "        # with torch.no_grad():\n",
    "        if optimizer == 'sgd':\n",
    "            var_learning_rate = adjust_learning_rate(ep, iteration,  len(dataloader),\n",
    "                                      n_epochs, learning_rate, lr_decay='cos')\n",
    "            for param_group in optimizer.param_groups:\n",
    "                param_group['lr'] = var_learning_rate\n",
    "        if cls_weights is not None:\n",
    "            loss_cls.weight = adjust_classes_weights(ep, iteration, len(dataloader),\n",
    "                                      n_epochs, cls_weights, descending=False)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        input, target, target_ndvi, _, _ = data\n",
    "        num_processed_samples += len(input)\n",
    "        if torch.cuda.is_available():\n",
    "            input = input.cuda()\n",
    "            target = target.cuda()\n",
    "            target_ndvi = target_ndvi.cuda()\n",
    "\n",
    "        if loss == 'batch':\n",
    "            output = network.forward(input)\n",
    "            samples_per_cls = torch.stack([(target == x_u).sum() for x_u in range(n_classes)])\n",
    "            l = loss(output, target, samples_per_cls)\n",
    "        elif loss == 'ndvi':\n",
    "            out_ndvi, output = network.forward(input)\n",
    "            samples_per_cls = None\n",
    "            if loss_weights:\n",
    "                samples_per_cls = torch.stack([(target == x_u).sum() for x_u in range(n_classes)])\n",
    "            l = loss(out_ndvi, output, target, target_ndvi, samples_per_cls, weight_mse=1.0)\n",
    "        else:\n",
    "            output = network.forward(input)\n",
    "            l = loss(output, target)\n",
    "\n",
    "        l.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        with torch.no_grad():\n",
    "            pred = torch.argmax(output, dim=1).squeeze(1)\n",
    "\n",
    "            metrics_v, metrics_scalar = conf_mat_metrics(pred, target)\n",
    "            str_metrics = ''.join(['%s| %f | ' % (key, value) for (key, value) in metrics_scalar.items()])\n",
    "            loss_measure.add(l.item())\n",
    "            str_metrics += 'loss| %f | ' % loss_measure.mean\n",
    "\n",
    "            train_info = 'Train on | {} | Epoch| {} | [{}/{} ({:.0f}%)] | lr| {:.5f} | {}    '.format(\n",
    "                dataloader.dataset.name, ep, num_processed_samples, num_train_samples,\n",
    "                100. * (iteration + 1) / len(dataloader), var_learning_rate, str_metrics)\n",
    "            sys.stdout.write('\\r' + train_info)\n",
    "\n",
    "    print()\n",
    "    with open(os.path.join(result_path, result_train), 'a+') as f:\n",
    "        f.write(train_info + '\\n')\n",
    "\n",
    "\n",
    "def test_epoch(dataloader, network, loss):\n",
    "    num_processed_samples = 0\n",
    "    num_test_samples = len(dataloader.dataset)\n",
    "    labels = list(range(numclasses))\n",
    "    conf_mat_metrics = ConfusionMatrix(labels, ignore_class=ignore_index)\n",
    "    num_cls = 18\n",
    "    batch_size = 18\n",
    "    labels = list(range(num_cls))\n",
    "    conf_mat = ConfusionMatrix(labels, ignore_class=0)\n",
    "\n",
    "    for e in range(5):\n",
    "        target = torch.randint(num_cls, (batch_size,))\n",
    "        output = torch.rand(batch_size, num_cls)\n",
    "        pred = torch.argmax(output, dim=1)\n",
    "        metrics_v, metrics_s = conf_mat(pred, target)\n",
    "        for k, v in metrics_v.items():\n",
    "            print(k, v)\n",
    "        for k, v in metrics_s.items():\n",
    "            print(k, v)\n",
    "        print()\n",
    "    loss_measure = AverageValueMeter()\n",
    "    am = AverageValueMeter()\n",
    "    am.add(numpy.array([1, 0.5, 0.2, 0.1, 0.8]))\n",
    "    mean, std = am.value()\n",
    "    print(mean, std, am.sum)\n",
    "\n",
    "    am.add(numpy.array([0.5, 0.5, 0.5, 0.5, 0.5]))\n",
    "    mean, std = am.value()\n",
    "    print(mean, std, am.sum)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for iteration, data in enumerate(dataloader):\n",
    "\n",
    "            input, target, target_ndvi, _, _ = data\n",
    "            num_processed_samples += len(input)\n",
    "            if torch.cuda.is_available():\n",
    "                input = input.cuda()\n",
    "                target = target.cuda()\n",
    "                target_ndvi = target_ndvi.cuda()\n",
    "\n",
    "            if loss == 'batch':\n",
    "                output = network.forward(input)\n",
    "                samples_per_cls = torch.stack([(target == x_u).sum() for x_u in range(n_classes)])\n",
    "                l = loss(output, target, samples_per_cls)\n",
    "            elif loss == 'ndvi':\n",
    "                out_ndvi, output = network.forward(input)\n",
    "                samples_per_cls = None\n",
    "                if loss_weights:\n",
    "                    samples_per_cls = torch.stack([(target == x_u).sum() for x_u in range(n_classes)])\n",
    "                l = loss(out_ndvi, output, target, target_ndvi, samples_per_cls)\n",
    "            else:\n",
    "                output = network.forward(input)\n",
    "                l = loss(output, target)\n",
    "\n",
    "            pred = torch.argmax(output, dim=1).squeeze(1)\n",
    "            metrics_v, metrics_scalar = conf_mat_metrics(pred, target)\n",
    "            str_metrics = ''.join(['%s| %f | ' % (key, value) for (key, value) in metrics_scalar.items()])\n",
    "            loss_measure.add(l.item())\n",
    "            str_metrics += 'loss| %f | ' % loss_measure.mean\n",
    "            test_info = 'Test on | {} | Epoch | {} | [{}/{} ({:.0f}%)] | {}  '.format(\n",
    "                dataloader.dataset.name, epoch, num_processed_samples,\n",
    "                num_test_samples, 100. * (iteration + 1) / len(dataloader),\n",
    "                str_metrics)\n",
    "            sys.stdout.write('\\r' + test_info)\n",
    "\n",
    "        is_best = metrics_scalar['OA'] > best_test_acc\n",
    "        best = '  **best result' if is_best else '         '\n",
    "        test_info += best\n",
    "\n",
    "        sys.stdout.write('\\r' + test_info + '\\n')\n",
    "        with open(os.path.join(result_path, result_validation), 'a+') as f:\n",
    "            f.write(test_info + '\\n')\n",
    "\n",
    "        if is_best:\n",
    "            cls_names = numpy.array(traindataset.classes)[conf_mat_metrics.get_labels()]\n",
    "            with open(os.path.join(result_path, \"per_class_metricsDeepLab.txt\"), 'a+') as f:\n",
    "                f.write('classes:\\n' + numpy.array2string(cls_names) + '\\n')\n",
    "                for k, v in metrics_v.items():\n",
    "                    f.write(k + '\\n')\n",
    "                    if len(v.shape) == 1:\n",
    "                        for ki, vi in zip(cls_names, v):\n",
    "                            f.write(\"%.2f\" % vi + '\\t' + ki + '\\n')\n",
    "                    elif len(v.shape) == 2:  # confusion matrix\n",
    "                        num_gt = numpy.sum(v, axis=1)\n",
    "                        f.write('\\n'.join(\n",
    "                            [''.join(['{:10}'.format(item) for item in row]) + '  ' + lab + '(%d)' % tot\n",
    "                             for row, lab, tot in zip(v, cls_names, num_gt)]))\n",
    "                        f.write('\\n')\n",
    "\n",
    "        return metrics_scalar['OA']  # test_acc\n",
    "\n",
    "\n",
    "def test_only(dataloader, network, loss, epoch, set_name):\n",
    "    num_processed_samples = 0\n",
    "    num_test_samples = len(dataloader.dataset)\n",
    "    labels = list(range(numclasses))\n",
    "    conf_mat_metrics = ConfusionMatrix(labels, ignore_class=ignore_index)\n",
    "    num_cls = 18\n",
    "    batch_size = 18\n",
    "    labels = list(range(num_cls))\n",
    "    conf_mat = ConfusionMatrix(labels, ignore_class=0)\n",
    "\n",
    "    for e in range(5):\n",
    "        target = torch.randint(num_cls, (batch_size,))\n",
    "        output = torch.rand(batch_size, num_cls)\n",
    "        pred = torch.argmax(output, dim=1)\n",
    "        metrics_v, metrics_s = conf_mat(pred, target)\n",
    "        for k, v in metrics_v.items():\n",
    "            print(k, v)\n",
    "        for k, v in metrics_s.items():\n",
    "            print(k, v)\n",
    "        print()\n",
    "    loss_measure = AverageValueMeter()\n",
    "    am = AverageValueMeter()\n",
    "    am.add(numpy.array([1, 0.5, 0.2, 0.1, 0.8]))\n",
    "    mean, std = am.value()\n",
    "    print(mean, std, am.sum)\n",
    "\n",
    "    am.add(numpy.array([0.5, 0.5, 0.5, 0.5, 0.5]))\n",
    "    mean, std = am.value()\n",
    "    print(mean, std, am.sum)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for iteration, data in enumerate(dataloader):\n",
    "\n",
    "            input, target, target_ndvi, dates, patch_id = data\n",
    "\n",
    "            num_processed_samples += len(input)\n",
    "            if torch.cuda.is_available():\n",
    "                input = input.cuda()\n",
    "                target = target.cuda()\n",
    "                target_ndvi = target_ndvi.cuda()\n",
    "\n",
    "            if loss == 'batch':\n",
    "                output = network.forward(input)\n",
    "                samples_per_cls = torch.stack([(target == x_u).sum() for x_u in range(n_classes)])\n",
    "                l = loss(output, target, samples_per_cls)\n",
    "            elif loss == 'ndvi':\n",
    "                out_ndvi, output = network.forward(input)\n",
    "                samples_per_cls = None\n",
    "                if loss_weights:\n",
    "                    samples_per_cls = torch.stack([(target == x_u).sum() for x_u in range(n_classes)])\n",
    "                l = loss(out_ndvi, output, target, target_ndvi, samples_per_cls)\n",
    "                if iteration % 100 == 0:\n",
    "                    write_signatures(model.module, target, output, target_ndvi, out_ndvi, dates, patch_id, set_name)\n",
    "            else:\n",
    "                output = network.forward(input)\n",
    "                l = loss(output, target)\n",
    "\n",
    "            pred = torch.argmax(output, dim=1).squeeze(1)\n",
    "            metrics_v, metrics_scalar = conf_mat_metrics(pred, target)\n",
    "            str_metrics = ''.join(['%s| %f | ' % (key, value) for (key, value) in metrics_scalar.items()])\n",
    "            loss_measure.add(l.item())\n",
    "            str_metrics += 'loss| %f | ' % loss_measure.mean\n",
    "            test_info = 'Test on | {} | Epoch | {} | [{}/{} ({:.0f}%)] | {}  '.format(\n",
    "                dataloader.dataset.name, epoch, num_processed_samples,\n",
    "                num_test_samples, 100. * (iteration + 1) / len(dataloader),\n",
    "                str_metrics)\n",
    "            if ((100. * (iteration + 1))%100 == 0):\n",
    "                sys.stdout.write('\\r' + test_info)\n",
    "            \n",
    "\n",
    "        cls_names = numpy.array(traindataset.classes)[conf_mat_metrics.get_labels()]\n",
    "        if ((100. * (iteration + 1))%100 == 0):\n",
    "            sys.stdout.write('\\r' + test_info)\n",
    "        with open(os.path.join(result_path, set_name + \"_per_class_metricsDeepLab.txt\"), 'w') as f:\n",
    "            f.write('classes:\\n' + numpy.array2string(cls_names) + '\\n')\n",
    "            sys.stdout.write('classes:\\n' + numpy.array2string(cls_names) + '\\n')\n",
    "            for k, v in metrics_v.items():\n",
    "                sys.stdout.write('\\n' + k + '\\n')\n",
    "                f.write('\\n' + k + '\\n')\n",
    "                if len(v.shape) == 1:\n",
    "                    for ki, vi in zip(cls_names, v):\n",
    "                        sys.stdout.write(\"%.2f\" % vi + '\\t' + ki + '\\n')\n",
    "                        f.write(\"%.2f\" % vi + '\\t' + ki + '\\n')\n",
    "                elif len(v.shape) == 2:  # confusion matrix\n",
    "                    num_gt = numpy.sum(v, axis=1)\n",
    "                    sys.stdout.write('\\n'.join(\n",
    "                        [''.join(['{:10}'.format(item) for item in row]) + '  ' + lab + '(%d)' % tot\n",
    "                         for row, lab, tot in zip(v, cls_names, num_gt)]))\n",
    "                    f.write('\\n'.join(\n",
    "                        [''.join(['{:10}'.format(item) for item in row]) + '  ' + lab + '(%d)' % tot\n",
    "                         for row, lab, tot in zip(v, cls_names, num_gt)]))\n",
    "                    sys.stdout.write('\\n')\n",
    "                    f.write('\\n')\n",
    "\n",
    "        if loss == 'ndvi':\n",
    "            print(\"\\nClass Activation Interval saved in:\", result_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3995354d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-26T16:49:44.241974Z",
     "iopub.status.busy": "2022-05-26T16:49:44.241712Z",
     "iopub.status.idle": "2022-05-26T16:49:44.391162Z",
     "shell.execute_reply": "2022-05-26T16:49:44.390489Z"
    },
    "papermill": {
     "duration": 0.188928,
     "end_time": "2022-05-26T16:49:44.393495",
     "exception": false,
     "start_time": "2022-05-26T16:49:44.204567",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ASPP(nn.Module):\n",
    "    def __init__(self, num_classes):\n",
    "        super(ASPP, self).__init__()\n",
    "\n",
    "        self.conv_1x1_1 = nn.Conv3d(512, 256, kernel_size=1)\n",
    "        self.bn_conv_1x1_1 = nn.BatchNorm3d(256)\n",
    "\n",
    "        self.conv_3x3_1 = nn.Conv3d(512, 256, kernel_size=3, stride=1, padding=6, dilation=6)\n",
    "        self.bn_conv_3x3_1 = nn.BatchNorm3d(256)\n",
    "\n",
    "        self.conv_3x3_2 = nn.Conv3d(512, 256, kernel_size=3, stride=1, padding=12, dilation=12)\n",
    "        self.bn_conv_3x3_2 = nn.BatchNorm3d(256)\n",
    "\n",
    "        self.conv_3x3_3 = nn.Conv3d(512, 256, kernel_size=3, stride=1, padding=18, dilation=18)\n",
    "        self.bn_conv_3x3_3 = nn.BatchNorm3d(256)\n",
    "\n",
    "        self.avg_pool = nn.AdaptiveAvgPool3d(1)\n",
    "\n",
    "        self.conv_1x1_2 = nn.Conv3d(512, 256, kernel_size=1)\n",
    "        self.bn_conv_1x1_2 = nn.BatchNorm3d(256)\n",
    "\n",
    "        self.conv_1x1_3 = nn.Conv3d(1280, 256, kernel_size=1)\n",
    "        self.bn_conv_1x1_3 = nn.BatchNorm3d(256)\n",
    "\n",
    "        self.conv_1x1_4 = nn.Conv3d(256, num_classes, kernel_size=1)\n",
    "\n",
    "    def forward(self, feature_map):\n",
    "        feature_map_h = feature_map.size()[2]\n",
    "        feature_map_w = feature_map.size()[3]\n",
    "        feature_map_c = feature_map.size()[4]\n",
    "\n",
    "        out_1x1 = F.relu(self.bn_conv_1x1_1(self.conv_1x1_1(feature_map)))\n",
    "        out_3x3_1 = F.relu(self.bn_conv_3x3_1(self.conv_3x3_1(feature_map)))\n",
    "        out_3x3_2 = F.relu(self.bn_conv_3x3_2(self.conv_3x3_2(feature_map)))\n",
    "        out_3x3_3 = F.relu(self.bn_conv_3x3_3(self.conv_3x3_3(feature_map)))\n",
    "\n",
    "        out_img = self.avg_pool(feature_map)\n",
    "        out_img = F.relu(self.bn_conv_1x1_2(self.conv_1x1_2(out_img)))\n",
    "        out_img = F.interpolate(out_img, size=(feature_map_h, feature_map_w, feature_map_c), mode='trilinear', align_corners=True)\n",
    "\n",
    "        out = torch.cat([out_1x1, out_3x3_1, out_3x3_2, out_3x3_3, out_img], 1)\n",
    "        out = F.relu(self.bn_conv_1x1_3(self.conv_1x1_3(out)))\n",
    "        out = self.conv_1x1_4(out)\n",
    "\n",
    "\n",
    "        return out\n",
    "\n",
    "class ASPP_Bottleneck(nn.Module):\n",
    "    def __init__(self, num_classes):\n",
    "        super(ASPP_Bottleneck, self).__init__()\n",
    "\n",
    "        self.conv_1x1_1 = nn.Conv3d(4*512, 256, kernel_size=1)\n",
    "        self.bn_conv_1x1_1 = nn.BatchNorm3d(256)\n",
    "\n",
    "        self.conv_3x3_1 = nn.Conv3d(4*512, 256, kernel_size=3, stride=1, padding=6, dilation=6)\n",
    "        self.bn_conv_3x3_1 = nn.BatchNorm3d(256)\n",
    "\n",
    "        self.conv_3x3_2 = nn.Conv3d(4*512, 256, kernel_size=3, stride=1, padding=12, dilation=12)\n",
    "        self.bn_conv_3x3_2 = nn.BatchNorm3d(256)\n",
    "\n",
    "        self.conv_3x3_3 = nn.Conv3d(4*512, 256, kernel_size=3, stride=1, padding=18, dilation=18)\n",
    "        self.bn_conv_3x3_3 = nn.BatchNorm3d(256)\n",
    "\n",
    "        self.avg_pool = nn.AdaptiveAvgPool3d(1)\n",
    "\n",
    "        self.conv_1x1_2 = nn.Conv3d(4*512, 256, kernel_size=1)\n",
    "        self.bn_conv_1x1_2 = nn.BatchNorm3d(256)\n",
    "\n",
    "        self.conv_1x1_3 = nn.Conv3d(1280, 256, kernel_size=1) # (1280 = 5*256)\n",
    "        self.bn_conv_1x1_3 = nn.BatchNorm3d(256)\n",
    "\n",
    "        self.conv_1x1_4 = nn.Conv3d(256, num_classes, kernel_size=1)\n",
    "\n",
    "    def forward(self, feature_map):\n",
    "        feature_map_h = feature_map.size()[2]\n",
    "        feature_map_w = feature_map.size()[3]\n",
    "        feature_map_c = feature_map.size()[4]\n",
    "\n",
    "        out_1x1 = F.relu(self.bn_conv_1x1_1(self.conv_1x1_1(feature_map)))\n",
    "        out_3x3_1 = F.relu(self.bn_conv_3x3_1(self.conv_3x3_1(feature_map)))\n",
    "        out_3x3_2 = F.relu(self.bn_conv_3x3_2(self.conv_3x3_2(feature_map)))\n",
    "        out_3x3_3 = F.relu(self.bn_conv_3x3_3(self.conv_3x3_3(feature_map)))\n",
    "\n",
    "        out_img = self.avg_pool(feature_map)\n",
    "        out_img = F.relu(self.bn_conv_1x1_2(self.conv_1x1_2(out_img)))\n",
    "        out_img = F.interpolate(out_img, size=(feature_map_h, feature_map_w, feature_map_c), mode='trilinear', align_corners=True)\n",
    "        \n",
    "        out = torch.cat([out_1x1, out_3x3_1, out_3x3_2, out_3x3_3, out_img], 1)\n",
    "        out = F.relu(self.bn_conv_1x1_3(self.conv_1x1_3(out)))\n",
    "        out = self.conv_1x1_4(out)\n",
    "\n",
    "        return out\n",
    "\n",
    "\n",
    "class DeepLabV3_3D(nn.Module):\n",
    "    def __init__(self, num_classes, input_channels, resnet, last_activation = None):\n",
    "        super(DeepLabV3_3D, self).__init__()\n",
    "        self.num_classes = num_classes\n",
    "        self.last_activation = last_activation\n",
    "\n",
    "        if resnet.lower() == 'resnet18_os16':\n",
    "            self.resnet = ResNet18_OS16(input_channels)\n",
    "        \n",
    "        elif resnet.lower() == 'resnet34_os16':\n",
    "            self.resnet = ResNet34_OS16(input_channels)\n",
    "        \n",
    "        elif resnet.lower() == 'resnet50_os16':\n",
    "            self.resnet = ResNet50_OS16(input_channels)\n",
    "        \n",
    "        elif resnet.lower() == 'resnet101_os16':\n",
    "            self.resnet = ResNet101_OS16(input_channels)\n",
    "        \n",
    "        elif resnet.lower() == 'resnet152_os16':\n",
    "            self.resnet = ResNet152_OS16(input_channels)\n",
    "        \n",
    "        elif resnet.lower() == 'resnet18_os8':\n",
    "            self.resnet = ResNet18_OS8(input_channels)\n",
    "        \n",
    "        elif resnet.lower() == 'resnet34_os8':\n",
    "            self.resnet = ResNet34_OS8(input_channels)\n",
    "\n",
    "        if resnet.lower() in ['resnet50_os16', 'resnet101_os16', 'resnet152_os16']:\n",
    "            self.aspp = ASPP_Bottleneck(num_classes=self.num_classes)\n",
    "        else:\n",
    "            self.aspp = ASPP(num_classes=self.num_classes)\n",
    "\n",
    "        num_classes = 18\n",
    "        self.final_conv = torch.nn.Conv3d(30, 1, kernel_size=1, stride=1, padding=0, bias=True)\n",
    "        \n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        h = x.size()[2]\n",
    "        w = x.size()[3]\n",
    "        c = x.size()[4]\n",
    "\n",
    "        feature_map = self.resnet(x)\n",
    "\n",
    "        output = self.aspp(feature_map)\n",
    "\n",
    "        output = F.interpolate(output, size=(h, w, c), mode='trilinear', align_corners=True)\n",
    "        \n",
    "        if self.last_activation.lower() == 'sigmoid':\n",
    "            output = nn.Sigmoid()(output)\n",
    "        \n",
    "        elif self.last_activation.lower() == 'softmax':\n",
    "            output = nn.Softmax()(output)\n",
    "\n",
    "        \n",
    "        output = torch.permute(output, (0,2,1,3,4))\n",
    "        print(output.shape)\n",
    "        output = self.final_conv(output)\n",
    "        output = torch.squeeze(output)\n",
    "        \n",
    "        \n",
    "        return output\n",
    "\n",
    "__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',\n",
    "           'resnet152']\n",
    "\n",
    "\n",
    "\n",
    "def conv3x3(in_planes, out_planes, stride=1):\n",
    "    \"\"\"3x3 convolution with padding\"\"\"\n",
    "    return nn.Conv3d(in_planes, out_planes, kernel_size=3, stride=stride,\n",
    "                     padding=1, bias=False)\n",
    "\n",
    "\n",
    "def conv1x1(in_planes, out_planes, stride=1):\n",
    "    \"\"\"1x1 convolution\"\"\"\n",
    "    return nn.Conv3d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)\n",
    "\n",
    "\n",
    "class BasicBlock(nn.Module):\n",
    "    expansion = 1\n",
    "\n",
    "    def __init__(self, inplanes, planes, stride=1, downsample=None):\n",
    "        super(BasicBlock, self).__init__()\n",
    "        self.conv1 = conv3x3(inplanes, planes, stride)\n",
    "        self.bn1 = nn.BatchNorm3d(planes)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.conv2 = conv3x3(planes, planes)\n",
    "        self.bn2 = nn.BatchNorm3d(planes)\n",
    "        self.downsample = downsample\n",
    "        self.stride = stride\n",
    "\n",
    "    def forward(self, x):\n",
    "        identity = x\n",
    "\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        out = self.relu(out)\n",
    "\n",
    "        out = self.conv2(out)\n",
    "        out = self.bn2(out)\n",
    "\n",
    "        if self.downsample is not None:\n",
    "            identity = self.downsample(x)\n",
    "\n",
    "        out += identity\n",
    "        out = self.relu(out)\n",
    "\n",
    "        return out\n",
    "\n",
    "\n",
    "class Bottleneck(nn.Module):\n",
    "    expansion = 4\n",
    "\n",
    "    def __init__(self, inplanes, planes, stride=1, downsample=None):\n",
    "        super(Bottleneck, self).__init__()\n",
    "        self.conv1 = conv1x1(inplanes, planes)\n",
    "        self.bn1 = nn.BatchNorm3d(planes)\n",
    "        self.conv2 = conv3x3(planes, planes, stride)\n",
    "        self.bn2 = nn.BatchNorm3d(planes)\n",
    "        self.conv3 = conv1x1(planes, planes * self.expansion)\n",
    "        self.bn3 = nn.BatchNorm3d(planes * self.expansion)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.downsample = downsample\n",
    "        self.stride = stride\n",
    "\n",
    "    def forward(self, x):\n",
    "        identity = x\n",
    "\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        out = self.relu(out)\n",
    "\n",
    "        out = self.conv2(out)\n",
    "        out = self.bn2(out)\n",
    "        out = self.relu(out)\n",
    "\n",
    "        out = self.conv3(out)\n",
    "        out = self.bn3(out)\n",
    "\n",
    "        if self.downsample is not None:\n",
    "            identity = self.downsample(x)\n",
    "\n",
    "        out += identity\n",
    "        out = self.relu(out)\n",
    "\n",
    "        return out\n",
    "\n",
    "\n",
    "class ResNet(nn.Module):\n",
    "\n",
    "    def __init__(self, input_channels, block, layers, num_classes=1, zero_init_residual=False):\n",
    "        super(ResNet, self).__init__()\n",
    "        self.inplanes = 64\n",
    "        self.conv1 = nn.Conv3d(input_channels, 64, kernel_size=7, stride=2, padding=3,\n",
    "                               bias=False)\n",
    "        self.bn1 = nn.BatchNorm3d(64)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.maxpool = nn.MaxPool3d(kernel_size=3, stride=2, padding=1)\n",
    "        self.layer1 = self._make_layer(block, 64, layers[0])\n",
    "        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n",
    "        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n",
    "        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)\n",
    "        self.avgpool = nn.AdaptiveAvgPool3d((1, 1))\n",
    "        self.fc = nn.Linear(512 * block.expansion, num_classes)\n",
    "\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv3d):\n",
    "                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n",
    "            elif isinstance(m, nn.BatchNorm3d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "\n",
    "        # Zero-initialize the last BN in each residual branch,\n",
    "        # so that the residual branch starts with zeros, and each residual block behaves like an identity.\n",
    "        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677\n",
    "        if zero_init_residual:\n",
    "            for m in self.modules():\n",
    "                if isinstance(m, Bottleneck):\n",
    "                    nn.init.constant_(m.bn3.weight, 0)\n",
    "                elif isinstance(m, BasicBlock):\n",
    "                    nn.init.constant_(m.bn2.weight, 0)\n",
    "\n",
    "    def _make_layer(self, block, planes, blocks, stride=1):\n",
    "        downsample = None\n",
    "        if stride != 1 or self.inplanes != planes * block.expansion:\n",
    "            downsample = nn.Sequential(\n",
    "                conv1x1(self.inplanes, planes * block.expansion, stride),\n",
    "                nn.BatchNorm3d(planes * block.expansion),\n",
    "            )\n",
    "\n",
    "        layers = []\n",
    "        layers.append(block(self.inplanes, planes, stride, downsample))\n",
    "        self.inplanes = planes * block.expansion\n",
    "        for _ in range(1, blocks):\n",
    "            layers.append(block(self.inplanes, planes))\n",
    "\n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.bn1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.maxpool(x)\n",
    "\n",
    "        x = self.layer1(x)\n",
    "        x = self.layer2(x)\n",
    "        x = self.layer3(x)\n",
    "        x = self.layer4(x)\n",
    "\n",
    "        x = self.avgpool(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.fc(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "def resnet18(input_channels, **kwargs):\n",
    "    model = ResNet(input_channels, BasicBlock, [2, 2, 2, 2], **kwargs)\n",
    "    return model\n",
    "\n",
    "\n",
    "def resnet34(input_channels, **kwargs):\n",
    "    model = ResNet(input_channels, BasicBlock, [3, 4, 6, 3], **kwargs)\n",
    "    return model\n",
    "\n",
    "\n",
    "def resnet50(input_channels, **kwargs):\n",
    "    model = ResNet(input_channels, Bottleneck, [3, 4, 6, 3], **kwargs)\n",
    "    return model\n",
    "\n",
    "\n",
    "def resnet101(input_channels, **kwargs):\n",
    "    model = ResNet(input_channels, Bottleneck, [3, 4, 23, 3], **kwargs)\n",
    "    return model\n",
    "\n",
    "\n",
    "def resnet152(input_channels, **kwargs):\n",
    "    model = ResNet(input_channels, Bottleneck, [3, 8, 36, 3], **kwargs)\n",
    "    return model\n",
    "\n",
    "\n",
    "# -------------------------------------- Resnet for Deeplab --------------------------------------\n",
    "def make_layer(block, in_channels, channels, num_blocks, stride=1, dilation=1):\n",
    "    strides = [stride] + [1]*(num_blocks - 1)\n",
    "\n",
    "    blocks = []\n",
    "    for stride in strides:\n",
    "        blocks.append(block(in_channels=in_channels, channels=channels, stride=stride, dilation=dilation))\n",
    "        in_channels = block.expansion*channels\n",
    "\n",
    "    layer = nn.Sequential(*blocks) # (*blocks: call with unpacked list entires as arguments)\n",
    "\n",
    "    return layer\n",
    "\n",
    "class BasicBlock(nn.Module):\n",
    "    expansion = 1\n",
    "\n",
    "    def __init__(self, in_channels, channels, stride=1, dilation=1):\n",
    "        super(BasicBlock, self).__init__()\n",
    "\n",
    "        out_channels = self.expansion*channels\n",
    "        \n",
    "        if type(dilation) != type(1):\n",
    "            dilation = 1\n",
    "            \n",
    "        self.conv1 = nn.Conv3d(in_channels, channels, kernel_size=3, stride=stride, padding=dilation, dilation=dilation, bias=False)\n",
    "        self.bn1 = nn.BatchNorm3d(channels)\n",
    "\n",
    "        self.conv2 = nn.Conv3d(channels, channels, kernel_size=3, stride=1, padding=dilation, dilation=dilation, bias=False)\n",
    "        self.bn2 = nn.BatchNorm3d(channels)\n",
    "\n",
    "        if (stride != 1) or (in_channels != out_channels):\n",
    "            conv = nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False)\n",
    "            bn = nn.BatchNorm3d(out_channels)\n",
    "            self.downsample = nn.Sequential(conv, bn)\n",
    "        else:\n",
    "            self.downsample = nn.Sequential()\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = F.relu(self.bn1(self.conv1(x)))\n",
    "        out = self.bn2(self.conv2(out))\n",
    "\n",
    "        out = out + self.downsample(x)\n",
    "\n",
    "        out = F.relu(out)\n",
    "\n",
    "        return out\n",
    "\n",
    "class Bottleneck(nn.Module):\n",
    "    expansion = 4\n",
    "\n",
    "    def __init__(self, in_channels, channels, stride=1, dilation=1):\n",
    "        super(Bottleneck, self).__init__()\n",
    "\n",
    "        out_channels = self.expansion*channels\n",
    "\n",
    "        self.conv1 = nn.Conv3d(in_channels, channels, kernel_size=1, bias=False)\n",
    "        self.bn1 = nn.BatchNorm3d(channels)\n",
    "\n",
    "        self.conv2 = nn.Conv3d(channels, channels, kernel_size=3, stride=stride, padding=dilation, dilation=dilation, bias=False)\n",
    "        self.bn2 = nn.BatchNorm3d(channels)\n",
    "\n",
    "        self.conv3 = nn.Conv3d(channels, out_channels, kernel_size=1, bias=False)\n",
    "        self.bn3 = nn.BatchNorm3d(out_channels)\n",
    "\n",
    "        if (stride != 1) or (in_channels != out_channels):\n",
    "            conv = nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False)\n",
    "            bn = nn.BatchNorm3d(out_channels)\n",
    "            self.downsample = nn.Sequential(conv, bn)\n",
    "        else:\n",
    "            self.downsample = nn.Sequential()\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = F.relu(self.bn1(self.conv1(x)))\n",
    "        out = F.relu(self.bn2(self.conv2(out)))\n",
    "        out = self.bn3(self.conv3(out))\n",
    "\n",
    "        out = out + self.downsample(x)\n",
    "\n",
    "        out = F.relu(out)\n",
    "\n",
    "        return out\n",
    "\n",
    "class ResNet_Bottleneck_OS16(nn.Module):\n",
    "    def __init__(self, num_layers, input_channels):\n",
    "        super(ResNet_Bottleneck_OS16, self).__init__()\n",
    "\n",
    "        if num_layers == 50:\n",
    "            resnet = resnet50(input_channels)\n",
    "            self.resnet = nn.Sequential(*list(resnet.children())[:-3])\n",
    "        elif num_layers == 101:\n",
    "            resnet = resnet101(input_channels)\n",
    "            self.resnet = nn.Sequential(*list(resnet.children())[:-3])\n",
    "        elif num_layers == 152:\n",
    "            resnet = resnet152(input_channels)\n",
    "            self.resnet = nn.Sequential(*list(resnet.children())[:-3])\n",
    "        else:\n",
    "            raise Exception(\"num_layers must be in {50, 101, 152}!\")\n",
    "\n",
    "        self.layer5 = make_layer(Bottleneck, in_channels=4*256, channels=512, num_blocks=3, stride=1, dilation=2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        c4 = self.resnet(x)\n",
    "\n",
    "        output = self.layer5(c4)\n",
    "\n",
    "        return output\n",
    "\n",
    "class ResNet_BasicBlock_OS16(nn.Module):\n",
    "    def __init__(self, num_layers, input_channels):\n",
    "        super(ResNet_BasicBlock_OS16, self).__init__()\n",
    "\n",
    "        if num_layers == 18:\n",
    "            resnet = resnet18(input_channels)\n",
    "            self.resnet = nn.Sequential(*list(resnet.children())[:-3])\n",
    "\n",
    "            num_blocks = 2\n",
    "\n",
    "        elif num_layers == 34:\n",
    "            resnet = resnet34(input_channels)\n",
    "            self.resnet = nn.Sequential(*list(resnet.children())[:-3])\n",
    "\n",
    "            num_blocks = 3\n",
    "        else:\n",
    "            raise Exception(\"num_layers must be in {18, 34}!\")\n",
    "    \n",
    "        self.layer5 = make_layer(BasicBlock, in_channels=256, channels=512, num_blocks=num_blocks, stride=1, dilation=2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        c4 = self.resnet(x)\n",
    "\n",
    "        output = self.layer5(c4)\n",
    "\n",
    "        return output\n",
    "\n",
    "class ResNet_BasicBlock_OS8(nn.Module):\n",
    "    def __init__(self, num_layers, input_channels):\n",
    "        super(ResNet_BasicBlock_OS8, self).__init__()\n",
    "\n",
    "        if num_layers == 18:\n",
    "            resnet = resnet18(input_channels)\n",
    "            \n",
    "            self.resnet = nn.Sequential(*list(resnet.children())[:-4])\n",
    "\n",
    "            num_blocks_layer_4 = 2\n",
    "            num_blocks_layer_5 = 2\n",
    "\n",
    "        elif num_layers == 34:\n",
    "            resnet = resnet34(input_channels)\n",
    "            \n",
    "            self.resnet = nn.Sequential(*list(resnet.children())[:-4])\n",
    "\n",
    "            num_blocks_layer_4 = 6\n",
    "            num_blocks_layer_5 = 3\n",
    "        else:\n",
    "            raise Exception(\"num_layers must be in {18, 34}!\")\n",
    "\n",
    "        self.layer4 = make_layer(BasicBlock, in_channels=128, channels=256, num_blocks=num_blocks_layer_4, stride=1, dilation=2)\n",
    "\n",
    "        self.layer5 = make_layer(BasicBlock, in_channels=256, channels=512, num_blocks=num_blocks_layer_5, stride=1, dilation=4)\n",
    "\n",
    "        \n",
    "\n",
    "    def forward(self, x):\n",
    "        c3 = self.resnet(x)\n",
    "\n",
    "        output = self.layer4(c3)\n",
    "        output = self.layer5(output)\n",
    "        \n",
    "\n",
    "        return output\n",
    "\n",
    "def ResNet18_OS16(input_channels):\n",
    "    return ResNet_BasicBlock_OS16(num_layers=18, input_channels=input_channels)\n",
    "\n",
    "def ResNet50_OS16(input_channels):\n",
    "    return ResNet_Bottleneck_OS16(num_layers=50, input_channels=input_channels)\n",
    "\n",
    "def ResNet101_OS16(input_channels):\n",
    "    return ResNet_Bottleneck_OS16(num_layers=101, input_channels=input_channels)\n",
    "\n",
    "def ResNet152_OS16(input_channels):\n",
    "    return ResNet_Bottleneck_OS16(num_layers=152, input_channels=input_channels)\n",
    "\n",
    "def ResNet34_OS16(input_channels):\n",
    "    return ResNet_BasicBlock_OS16(num_layers=34, input_channels=input_channels)\n",
    "\n",
    "def ResNet18_OS8(input_channels):\n",
    "    return ResNet_BasicBlock_OS8(num_layers=18, input_channels=input_channels)\n",
    "\n",
    "def ResNet34_OS8(input_channels):\n",
    "    return ResNet_BasicBlock_OS8(num_layers=34, input_channels=input_channels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38912d11",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-26T16:49:44.449658Z",
     "iopub.status.busy": "2022-05-26T16:49:44.449413Z",
     "iopub.status.idle": "2022-05-26T16:49:49.476185Z",
     "shell.execute_reply": "2022-05-26T16:49:49.475384Z"
    },
    "papermill": {
     "duration": 5.052784,
     "end_time": "2022-05-26T16:49:49.478799",
     "exception": false,
     "start_time": "2022-05-26T16:49:44.426015",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "input_channels = 13          \n",
    "resnet = 'resnet34_os8'    \n",
    "last_activation = 'softmax' \n",
    "\n",
    "model = DeepLabV3_3D(num_classes = 18, input_channels = input_channels, resnet = resnet, last_activation = last_activation).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdf78773",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-26T16:49:49.523872Z",
     "iopub.status.busy": "2022-05-26T16:49:49.523615Z",
     "iopub.status.idle": "2022-05-27T01:39:02.126329Z",
     "shell.execute_reply": "2022-05-27T01:39:02.125436Z"
    },
    "papermill": {
     "duration": 31752.62749,
     "end_time": "2022-05-27T01:39:02.128329",
     "exception": false,
     "start_time": "2022-05-26T16:49:49.500839",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "labels = list(range(numclasses))\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    model = torch.nn.DataParallel(model).cuda()\n",
    "if ottimizzatore == 'adam':\n",
    "    optimizer = torch.optim.Adam(model.parameters(),\n",
    "                                lr=learning_rate,\n",
    "                                weight_decay=weight_decay)\n",
    "elif ottimizzatore == 'sgd':\n",
    "    optimizer = torch.optim.SGD(model.parameters(),\n",
    "                                lr=learning_rate,\n",
    "                                momentum=momentum,\n",
    "                                weight_decay=weight_decay)\n",
    "\n",
    "# Define Criterion\n",
    "weights = None\n",
    "if not (loss == 'batch' or loss == 'ndvi') and loss_weights:\n",
    "    print('Computing weights per classes...')\n",
    "    weights = compute_train_weights(traindataloader)\n",
    "    weights = torch.sqrt(weights)\n",
    "    print(\"weights per classes:\", weights)\n",
    "\n",
    "loss_cls = SegmentationLosses(weight=weights, cuda=True, ignore_index=ignore_index)\n",
    "loss = loss_cls.build_loss(mode='ce')\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "start_epoch = 0\n",
    "if resume_path:  # Empty strings are considered false\n",
    "    print('trying to resume previous saved model...')\n",
    "    state = resume(resume_path, model=model, optimizer=optimizer)\n",
    "\n",
    "    if \"epoch\" in state.keys():\n",
    "        start_epoch = state[\"epoch\"]\n",
    "        best_test_acc = state['best_test_acc']\n",
    "\n",
    "\n",
    "for epoch in range(start_epoch, n_epochs):\n",
    "    train_epoch(traindataloader, model, optimizer, loss, epoch, loss_cls, cls_weights=None)\n",
    "    val_acc = test_epoch(validationdataloader, model, loss)\n",
    "\n",
    "    is_best = val_acc > best_test_acc\n",
    "    if is_best:\n",
    "        epochs_best_acc = epoch\n",
    "        best_test_acc = val_acc\n",
    "        if result_path:\n",
    "            checkpoint_name = os.path.join(result_path, \"best_model.pth\")\n",
    "            save(checkpoint_name, model, optimizer,\n",
    "                epoch=epoch, best_test_acc=best_test_acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a74081db",
   "metadata": {
    "papermill": {
     "duration": 13.739654,
     "end_time": "2022-05-27T01:39:29.113143",
     "exception": false,
     "start_time": "2022-05-27T01:39:15.373489",
     "status": "completed"
    },
    "tags": []
   },
   "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.10.7"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 32383.074567,
   "end_time": "2022-05-27T01:39:45.508980",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2022-05-26T16:40:02.434413",
   "version": "2.3.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}