SAI_DL_Ori - Git.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7c9e9e57-0942-499d-9469-82fae400ab10",
   "metadata": {
    "id": "7c9e9e57-0942-499d-9469-82fae400ab10"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "import cv2 as cv\n",
    "import pathlib, glob, os\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "\n",
    "import torch\n",
    "torch.cuda.empty_cache()\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torchsummary import summary\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torchvision import datasets, transforms\n",
    "from torch.autograd import Variable\n",
    "\n",
    "import albumentations as album"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ed5b0894-9143-4a90-b0b6-0151400de159",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "values = np.array([11, 30, 100, 41, 19, 0, 10, 55, 78, 90, -10, -11, 6, 105])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5fa1f08a-30fb-43b9-b7b6-fea8141b45df",
   "metadata": {
    "id": "5fa1f08a-30fb-43b9-b7b6-fea8141b45df"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "import re\n",
    "\n",
    "def read_pfm(fpath, expected_identifier=\"Pf\"):\n",
    "    # PFM format definition: http://netpbm.sourceforge.net/doc/pfm.html\n",
    "    \n",
    "    def _get_next_line(f):\n",
    "        next_line = f.readline().decode('utf-8').rstrip()\n",
    "        # ignore comments\n",
    "        while next_line.startswith('#'):\n",
    "            next_line = f.readline().rstrip()\n",
    "        return next_line\n",
    "    \n",
    "    with open(fpath, 'rb') as f:\n",
    "        #  header\n",
    "        identifier = _get_next_line(f)\n",
    "        if identifier != expected_identifier:\n",
    "            raise Exception('Unknown identifier. Expected: \"%s\", got: \"%s\".' % (expected_identifier, identifier))\n",
    "\n",
    "        try:\n",
    "            line_dimensions = _get_next_line(f)\n",
    "            dimensions = line_dimensions.split(' ')\n",
    "            width = int(dimensions[0].strip())\n",
    "            height = int(dimensions[1].strip())\n",
    "        except:\n",
    "            raise Exception('Could not parse dimensions: \"%s\". '\n",
    "                            'Expected \"width height\", e.g. \"512 512\".' % line_dimensions)\n",
    "\n",
    "        try:\n",
    "            line_scale = _get_next_line(f)\n",
    "            scale = float(line_scale)\n",
    "            assert scale != 0\n",
    "            if scale < 0:\n",
    "                endianness = \"<\"\n",
    "            else:\n",
    "                endianness = \">\"\n",
    "        except:\n",
    "            raise Exception('Could not parse max value / endianess information: \"%s\". '\n",
    "                            'Should be a non-zero number.' % line_scale)\n",
    "\n",
    "        try:\n",
    "            data = np.fromfile(f, \"%sf\" % endianness)\n",
    "            data = np.reshape(data, (height, width))\n",
    "            data = np.flipud(data)\n",
    "            with np.errstate(invalid=\"ignore\"):\n",
    "                data *= abs(scale)\n",
    "        except:\n",
    "            raise Exception('Invalid binary values. Could not create %dx%d array from input.' % (height, width))\n",
    "\n",
    "        return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "95ccb198-5a52-48e2-96e6-8cb7159f2cf1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' 1rst time load training data\\n#Datasets : as dataset is the same for each epoch we save the data after the 1rst reading and reuse for next epoches to save training time\\ndef ImageDataset(path):\\n    imgs_path = path + \"/*\"\\n    paths = []\\n    for classes in glob.glob(imgs_path):\\n        class_name = classes.split(\"\\\\\")[-1]#/\\n        paths.append(class_name)\\n    print(len(paths), paths)\\n\\n    # Transforms for low resolution images and high resolution images\\n    X = [0, 64, 128, 192, 256, 320, 384, 448,\\n         0, 64, 128, 192, 256, 320, 384, 448,\\n         0, 64, 128, 192, 256, 320, 384, 448,\\n         0, 64, 128, 192, 256, 320, 384, 448,\\n         0, 64, 128, 192, 256, 320, 384, 448,\\n         0, 64, 128, 192, 256, 320, 384, 448,\\n         0, 64, 128, 192, 256, 320, 384, 448,\\n         0, 64, 128, 192, 256, 320, 384, 448]\\n    Y = [0, 0, 0, 0, 0, 0, 0, 0,     \\n         64, 64, 64, 64, 64, 64, 64, 64,\\n         128, 128, 128, 128, 128, 128, 128, 128,\\n         192, 192, 192, 192, 192, 192, 192, 192,\\n         256, 256, 256, 256, 256, 256, 256, 256,\\n         320, 320, 320, 320, 320, 320, 320, 320, \\n         384, 384, 384, 384, 384, 384, 384, 384,\\n         448, 448, 448, 448, 448, 448, 448, 448]\\n\\n    for i in range (len(paths)):\\n        Img = np.zeros((81, 512, 512, 3))\\n        Imgs = np.zeros((81, 64, 64, 3))    \\n        \\n        GT_disparity = read_pfm(str(\\'%s/%s/%s.pfm\\') %(path, paths[i], \"gt_disp_lowres\"))\\n        for nb in range (81):\\n            Img[nb] = cv.imread(str(\\'%s/%s/%s%03d.png\\') %(path, paths[i], \"input_Cam\", nb))\\n            \\n        for idx in range (len(X)):            \\n            x_min=X[idx]\\n            y_min=Y[idx]\\n            \\n            aug = 0#np.random.randint(0,3)         \\n            for j in range (81): \\n                if aug == 1:\\n                    Imgs[j] = cv.flip(np.squeeze(Img[j, x_min:x_min+64, y_min:y_min+64]), 0)\\n                elif aug == 2:\\n                    Imgs[j] = cv.flip(np.squeeze(Img[j, x_min:x_min+64, y_min:y_min+64]), 1)   \\n                else:\\n                    Imgs[j] = Img[j, x_min:x_min+64, y_min:y_min+64]\\n                    \\n                if j == 0:\\n                    cropped_Imgs = Imgs[j]\\n                else:\\n                    cropped_Imgs = np.concatenate((cropped_Imgs, Imgs[j]), axis=-1)\\n\\n            cropped_GT_disparity = GT_disparity[x_min:x_min+64, y_min:y_min+64]\\n            if aug == 1:\\n                cropped_GT_disparity = - cv.flip(cropped_GT_disparity, 0)\\n            elif aug == 2:\\n                cropped_GT_disparity = - cv.flip(cropped_GT_disparity, 1)\\n            \\n            img_lr = (np.expand_dims(np.transpose(cropped_Imgs, (2, 0, 1)),0)/255.).astype(\\'float32\\')\\n            gt_lr  = (np.expand_dims(np.expand_dims(cropped_GT_disparity,0),0)).astype(\\'float32\\')\\n            \\n            if np.bitwise_and(i==0, idx ==0):\\n                Img_lr = img_lr\\n                Gt_lr  = gt_lr\\n            else:\\n                Img_lr =  np.concatenate((Img_lr, img_lr), axis=0)\\n                Gt_lr  =  np.concatenate((Gt_lr, gt_lr), axis=0)\\n        \\n    return Img_lr, Gt_lr\\n\\nbatch_size = 32\\n\\nx_train, y_train = ImageDataset(\"hci_dataset/additional\")\\nnp.save(\\'x_train_64_0.npy\\', x_train)\\nnp.save(\\'y_train_64_0.npy\\', y_train)\\n\\ndata_train = torch.utils.data.TensorDataset(torch.from_numpy(x_train), torch.from_numpy(y_train))\\ntrainloader = DataLoader(data_train, batch_size = batch_size, shuffle=True, drop_last=True)\\nprint(len(trainloader))\\n'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\" 1rst time load training data\n",
    "#Datasets : as dataset is the same for each epoch we save the data after the 1rst reading and reuse for next epoches to save training time\n",
    "def ImageDataset(path):\n",
    "    imgs_path = path + \"/*\"\n",
    "    paths = []\n",
    "    for classes in glob.glob(imgs_path):\n",
    "        class_name = classes.split(\"\\\\\")[-1]#/\n",
    "        paths.append(class_name)\n",
    "    print(len(paths), paths)\n",
    "\n",
    "    # Transforms for low resolution images and high resolution images\n",
    "    X = [0, 64, 128, 192, 256, 320, 384, 448,\n",
    "         0, 64, 128, 192, 256, 320, 384, 448,\n",
    "         0, 64, 128, 192, 256, 320, 384, 448,\n",
    "         0, 64, 128, 192, 256, 320, 384, 448,\n",
    "         0, 64, 128, 192, 256, 320, 384, 448,\n",
    "         0, 64, 128, 192, 256, 320, 384, 448,\n",
    "         0, 64, 128, 192, 256, 320, 384, 448,\n",
    "         0, 64, 128, 192, 256, 320, 384, 448]\n",
    "    Y = [0, 0, 0, 0, 0, 0, 0, 0,     \n",
    "         64, 64, 64, 64, 64, 64, 64, 64,\n",
    "         128, 128, 128, 128, 128, 128, 128, 128,\n",
    "         192, 192, 192, 192, 192, 192, 192, 192,\n",
    "         256, 256, 256, 256, 256, 256, 256, 256,\n",
    "         320, 320, 320, 320, 320, 320, 320, 320, \n",
    "         384, 384, 384, 384, 384, 384, 384, 384,\n",
    "         448, 448, 448, 448, 448, 448, 448, 448]\n",
    "\n",
    "    for i in range (len(paths)):\n",
    "        Img = np.zeros((81, 512, 512, 3))\n",
    "        Imgs = np.zeros((81, 64, 64, 3))    \n",
    "        \n",
    "        GT_disparity = read_pfm(str('%s/%s/%s.pfm') %(path, paths[i], \"gt_disp_lowres\"))\n",
    "        for nb in range (81):\n",
    "            Img[nb] = cv.imread(str('%s/%s/%s%03d.png') %(path, paths[i], \"input_Cam\", nb))\n",
    "            \n",
    "        for idx in range (len(X)):            \n",
    "            x_min=X[idx]\n",
    "            y_min=Y[idx]\n",
    "            \n",
    "            aug = 0#np.random.randint(0,3)         \n",
    "            for j in range (81): \n",
    "                if aug == 1:\n",
    "                    Imgs[j] = cv.flip(np.squeeze(Img[j, x_min:x_min+64, y_min:y_min+64]), 0)\n",
    "                elif aug == 2:\n",
    "                    Imgs[j] = cv.flip(np.squeeze(Img[j, x_min:x_min+64, y_min:y_min+64]), 1)   \n",
    "                else:\n",
    "                    Imgs[j] = Img[j, x_min:x_min+64, y_min:y_min+64]\n",
    "                    \n",
    "                if j == 0:\n",
    "                    cropped_Imgs = Imgs[j]\n",
    "                else:\n",
    "                    cropped_Imgs = np.concatenate((cropped_Imgs, Imgs[j]), axis=-1)\n",
    "\n",
    "            cropped_GT_disparity = GT_disparity[x_min:x_min+64, y_min:y_min+64]\n",
    "            if aug == 1:\n",
    "                cropped_GT_disparity = - cv.flip(cropped_GT_disparity, 0)\n",
    "            elif aug == 2:\n",
    "                cropped_GT_disparity = - cv.flip(cropped_GT_disparity, 1)\n",
    "            \n",
    "            img_lr = (np.expand_dims(np.transpose(cropped_Imgs, (2, 0, 1)),0)/255.).astype('float32')\n",
    "            gt_lr  = (np.expand_dims(np.expand_dims(cropped_GT_disparity,0),0)).astype('float32')\n",
    "            \n",
    "            if np.bitwise_and(i==0, idx ==0):\n",
    "                Img_lr = img_lr\n",
    "                Gt_lr  = gt_lr\n",
    "            else:\n",
    "                Img_lr =  np.concatenate((Img_lr, img_lr), axis=0)\n",
    "                Gt_lr  =  np.concatenate((Gt_lr, gt_lr), axis=0)\n",
    "        \n",
    "    return Img_lr, Gt_lr\n",
    "\n",
    "batch_size = 32\n",
    "\n",
    "x_train, y_train = ImageDataset(\"hci_dataset/additional\")\n",
    "np.save('x_train_64_0.npy', x_train)\n",
    "np.save('y_train_64_0.npy', y_train)\n",
    "\n",
    "data_train = torch.utils.data.TensorDataset(torch.from_numpy(x_train), torch.from_numpy(y_train))\n",
    "trainloader = DataLoader(data_train, batch_size = batch_size, shuffle=True, drop_last=True)\n",
    "print(len(trainloader))\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0124c670-79d2-42cd-86fd-9914f07eb8ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "32\n"
     ]
    }
   ],
   "source": [
    "#Datasets\n",
    "batch_size = 16*2\n",
    "data_train = torch.utils.data.TensorDataset(torch.from_numpy(np.load('x_train_64_0.npy')), torch.from_numpy(np.load('y_train_64_0.npy')))\n",
    "trainloader = DataLoader(data_train, batch_size = batch_size, shuffle=True, drop_last=True)\n",
    "print(len(trainloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa525d6f-139e-4af6-989a-78e0430b0ca6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4 ['boxes', 'cotton', 'dino', 'sideboard']\n"
     ]
    }
   ],
   "source": [
    "#Datasets\n",
    "def ImageDataset_T(path):\n",
    "    imgs_path = path + \"/*\"\n",
    "    paths = []\n",
    "    for classes in glob.glob(imgs_path):\n",
    "        class_name = classes.split(\"\\\\\")[-1]#/\n",
    "        paths.append(class_name)\n",
    "    print(len(paths), paths)\n",
    "\n",
    "    # Transforms for low resolution images and high resolution images\n",
    "    X = [0, 64, 128, 192, 256, 320, 384, 448,\n",
    "         0, 64, 128, 192, 256, 320, 384, 448,\n",
    "         0, 64, 128, 192, 256, 320, 384, 448,\n",
    "         0, 64, 128, 192, 256, 320, 384, 448,\n",
    "         0, 64, 128, 192, 256, 320, 384, 448,\n",
    "         0, 64, 128, 192, 256, 320, 384, 448,\n",
    "         0, 64, 128, 192, 256, 320, 384, 448,\n",
    "         0, 64, 128, 192, 256, 320, 384, 448]\n",
    "    Y = [0, 0, 0, 0, 0, 0, 0, 0,     \n",
    "         64, 64, 64, 64, 64, 64, 64, 64,\n",
    "         128, 128, 128, 128, 128, 128, 128, 128,\n",
    "         192, 192, 192, 192, 192, 192, 192, 192,\n",
    "         256, 256, 256, 256, 256, 256, 256, 256,\n",
    "         320, 320, 320, 320, 320, 320, 320, 320, \n",
    "         384, 384, 384, 384, 384, 384, 384, 384,\n",
    "         448, 448, 448, 448, 448, 448, 448, 448]\n",
    "\n",
    "    for i in range (len(paths)):\n",
    "        Img = np.zeros((81, 512, 512, 3))\n",
    "        Imgs = np.zeros((81, 64, 64, 3))    \n",
    "        \n",
    "        GT_disparity = read_pfm(str('%s/%s/%s.pfm') %(path, paths[i], \"gt_disp_lowres\"))\n",
    "        for nb in range (81):\n",
    "            Img[nb] = cv.imread(str('%s/%s/%s%03d.png') %(path, paths[i], \"input_Cam\", nb))\n",
    "            \n",
    "        for idx in range (len(X)):            \n",
    "            x_min=X[idx]\n",
    "            y_min=Y[idx]       \n",
    "            for j in range (81): \n",
    "                Imgs[j] = Img[j, x_min:x_min+64, y_min:y_min+64]\n",
    "                if j == 0:\n",
    "                    cropped_Imgs = Imgs[j]\n",
    "                else:\n",
    "                    cropped_Imgs = np.concatenate((cropped_Imgs, Imgs[j]), axis=-1)\n",
    "\n",
    "            cropped_GT_disparity = GT_disparity[x_min:x_min+64, y_min:y_min+64]\n",
    "            \n",
    "            img_lr = (np.expand_dims(np.transpose(cropped_Imgs, (2, 0, 1)),0)/255.).astype('float32')\n",
    "            gt_lr  = (np.expand_dims(np.expand_dims(cropped_GT_disparity,0),0)).astype('float32')\n",
    "            \n",
    "            if np.bitwise_and(i==0, idx ==0):\n",
    "                Img_lr = img_lr\n",
    "                Gt_lr  = gt_lr\n",
    "            else:\n",
    "                Img_lr =  np.concatenate((Img_lr, img_lr), axis=0)\n",
    "                Gt_lr  =  np.concatenate((Gt_lr, gt_lr), axis=0)\n",
    "        \n",
    "    return Img_lr, Gt_lr\n",
    "\n",
    "batch_size = 4*4\n",
    "\n",
    "x_test, y_test = ImageDataset_T(\"hci_dataset/training\")\n",
    "data_test = torch.utils.data.TensorDataset(torch.from_numpy(x_test), torch.from_numpy(y_test))\n",
    "testloader = DataLoader(data_test, batch_size = batch_size, shuffle=False, drop_last=True)\n",
    "print(len(testloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc8c1e01-ec99-4b57-9167-0749589f24c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "#Datasets : traditional methods - load data each epoch\n",
    "class ImageDataset_Iter(Dataset):\n",
    "    def __init__(self, path):\n",
    "        self.O_path = path\n",
    "        self.imgs_path = path + \"/*\"\n",
    "        self.paths = []\n",
    "        for classes in glob.glob(self.imgs_path):\n",
    "            class_name = classes.split(\"\\\\\")[-1]\n",
    "            self.paths.append(class_name)\n",
    "        print(len(self.paths), self.paths)\n",
    "            \n",
    "        # Transforms for low resolution images and high resolution images\n",
    "        self.X = [0, 64, 128, 192, 256, 320, 384, 448,\n",
    "                 0, 64, 128, 192, 256, 320, 384, 448,\n",
    "                 0, 64, 128, 192, 256, 320, 384, 448,\n",
    "                 0, 64, 128, 192, 256, 320, 384, 448,\n",
    "                 0, 64, 128, 192, 256, 320, 384, 448,\n",
    "                 0, 64, 128, 192, 256, 320, 384, 448,\n",
    "                 0, 64, 128, 192, 256, 320, 384, 448,\n",
    "                 0, 64, 128, 192, 256, 320, 384, 448]\n",
    "        self.Y = [0, 0, 0, 0, 0, 0, 0, 0,     \n",
    "                 64, 64, 64, 64, 64, 64, 64, 64,\n",
    "                 128, 128, 128, 128, 128, 128, 128, 128,\n",
    "                 192, 192, 192, 192, 192, 192, 192, 192,\n",
    "                 256, 256, 256, 256, 256, 256, 256, 256,\n",
    "                 320, 320, 320, 320, 320, 320, 320, 320, \n",
    "                 384, 384, 384, 384, 384, 384, 384, 384,\n",
    "                 448, 448, 448, 448, 448, 448, 448, 448]\n",
    "        self.lr_transform = transforms.Compose(\n",
    "            [\n",
    "                transforms.ToTensor(),\n",
    "            ]\n",
    "        )\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.paths*64)\n",
    "    \n",
    "    def __getitem__(self, idxx):#\n",
    "\n",
    "        idx = idxx//64\n",
    "        x_min=self.X[idxx%64] \n",
    "        y_min=self.Y[idxx%64] \n",
    "\n",
    "        Imgs = cv.imread(str('%s/%s/%s%03d.png') %(self.O_path, self.paths[idx], \"input_Cam\", 0))\n",
    "        cropped_Imgs = Imgs[x_min:x_min+64, y_min:y_min+64]\n",
    "        for i in range(1,81):\n",
    "            img = cv.imread(str('%s/%s/%s%03d.png') %(self.O_path, self.paths[idx], \"input_Cam\", i))\n",
    "            cropped_img = img[x_min:x_min+64, y_min:y_min+64]\n",
    "            cropped_Imgs =  np.concatenate((cropped_Imgs, cropped_img), axis=-1)\n",
    "\n",
    "        GT_disparity = read_pfm(str('%s/%s/%s.pfm') %(self.O_path, self.paths[idx], \"gt_disp_lowres\"))\n",
    "        cropped_GT_disparity = GT_disparity[x_min:x_min+64, y_min:y_min+64]\n",
    "        \n",
    "        img_lr = self.lr_transform(cropped_Imgs.copy())\n",
    "        gt_lr  = self.lr_transform(cropped_GT_disparity.copy())\n",
    "        \n",
    "        return {\"lr\": img_lr, \"GT\": gt_lr}\n",
    "\n",
    "\n",
    "batch_size = 4\n",
    "data_train = ImageDataset_Iter(\"hci_dataset/additional\")\n",
    "data_test  = ImageDataset_Iter(\"hci_dataset/training\")\n",
    "\n",
    "trainloader = DataLoader(data_train, batch_size = batch_size, shuffle=True, drop_last=True)\n",
    "testloader  = DataLoader(data_test, batch_size = batch_size, shuffle=True, drop_last=True)\n",
    "print(len(trainloader), len(testloader))\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b256363-6635-428e-b9ff-cf9faebdc709",
   "metadata": {
    "id": "4b256363-6635-428e-b9ff-cf9faebdc709"
   },
   "outputs": [],
   "source": [
    "class _ASPP(nn.Module):\n",
    "    \"\"\"\n",
    "    Atrous spatial pyramid pooling with image-level feature\n",
    "    \"\"\"\n",
    "    def __init__(self, in_ch, out_ch, rates):\n",
    "        super().__init__()\n",
    "        self.stages = nn.Module()\n",
    "        for i, rate in enumerate(rates):\n",
    "            self.stages.add_module(\n",
    "                \"C{}\".format(i + 1), nn.Conv2d(in_ch, out_ch, 3, 1, padding=rate, dilation=rate, bias=False))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return torch.cat([stage(x) for stage in self.stages.children()], dim=1)\n",
    "        \n",
    "#https://towardsdatascience.com/a-basic-introduction-to-separable-convolutions-b99ec3102728\n",
    "#Xception\n",
    "class _depthwise_separable_conv(nn.Module):#nn.Module\n",
    "    def __init__(self, in_ch, out_ch): #, kernels_per_layer\n",
    "        super().__init__() \n",
    "        self.depthwise = nn.Conv2d(in_ch, in_ch, kernel_size=3, padding='same', groups=in_ch) #in_ch * kernels_per_layer\n",
    "        self.pointwise = nn.Conv2d(in_ch, out_ch, kernel_size=1, padding='same') #in_ch * kernels_per_layer\n",
    "\n",
    "    def forward(self, x): \n",
    "        out = self.depthwise(x) \n",
    "        out = self.pointwise(out) \n",
    "        return out\n",
    "    \n",
    "class _DWSConvBnReLU(nn.Sequential):\n",
    "    \"\"\"\n",
    "    Cascade of depthwise separable 2D convolution, batch norm, and ReLU.\n",
    "    \"\"\"\n",
    "    def __init__(self, in_ch, out_ch, kernel_size=3, relu=True):\n",
    "        super().__init__()\n",
    "        self.add_module(\"DWS_Conv\", _depthwise_separable_conv(in_ch, out_ch))\n",
    "        self.add_module(\"BN\", nn.BatchNorm2d(out_ch))\n",
    "\n",
    "        if relu:\n",
    "            self.add_module(\"RELU\", nn.ReLU())#inplace=True\n",
    "            \n",
    "class _Bloc_1(nn.Sequential):\n",
    "    def __init__(self, kernel_size=3):\n",
    "        super().__init__()\n",
    "        self.add_module(\"DWS_Conv_1\", _DWSConvBnReLU(256, 128))\n",
    "        self.add_module(\"DWS_Conv_2\", _DWSConvBnReLU(128, 128))\n",
    "        self.add_module(\"DWS_Conv_3\", _DWSConvBnReLU(128, 128, relu=False))\n",
    "        \n",
    "class _All_Bloc_1(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.B_1 = _DWSConvBnReLU(256, 128, relu=False)\n",
    "        self.B_2 = _Bloc_1()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_ini = x\n",
    "        x1 = self.B_1(x_ini)\n",
    "        x2 = self.B_2(x_ini)\n",
    "        x_f = x1 + x2\n",
    "\n",
    "        return x_f    \n",
    "    \n",
    "class _Bloc(nn.Sequential):\n",
    "    def __init__(self, kernel_size=3):\n",
    "        super().__init__()\n",
    "        self.add_module(\"DWS_Conv_1\", _DWSConvBnReLU(128, 128))\n",
    "        self.add_module(\"DWS_Conv_2\", _DWSConvBnReLU(128, 128))\n",
    "        self.add_module(\"DWS_Conv_3\", _DWSConvBnReLU(128, 128, relu=False))\n",
    "        \n",
    "class _All_Bloc(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.B_1 = _DWSConvBnReLU(128, 128, relu=False)\n",
    "        self.B_2 = _Bloc()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_ini = x\n",
    "        x1 = self.B_1(x_ini)\n",
    "        x2 = self.B_2(x_ini)\n",
    "        x_f = x1 + x2\n",
    "\n",
    "        return x_f     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b175255b-cdb6-4840-b472-e739c7308dfb",
   "metadata": {
    "id": "b175255b-cdb6-4840-b472-e739c7308dfb"
   },
   "outputs": [],
   "source": [
    "class SAI_Model(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.in_ch = 81*3\n",
    "        self.rates = [1, 2, 4, 8, 16, 32]#[6, 12, 18] #\n",
    "\n",
    "        self.ASPP_1 = _ASPP(self.in_ch, self.in_ch, self.rates)\n",
    "        self.ASPP_2 = _DWSConvBnReLU(self.in_ch * len(self.rates), 256)#len(self.rates)\n",
    "        \n",
    "        self.bloc1 = _All_Bloc_1()\n",
    "        self.bloc  = _All_Bloc()\n",
    "        self.final_1 = _DWSConvBnReLU(128,128)\n",
    "        self.final_2 = nn.Conv2d(128,1, kernel_size=1, padding='same')\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.ASPP_1(x)\n",
    "        x = self.ASPP_2(x)\n",
    "        x = self.bloc1(x)\n",
    "        for i in range(3):#7\n",
    "            x = self.bloc(x)\n",
    "        x = self.final_1(x)\n",
    "        x = self.final_2(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "#Instantiate the model\n",
    "model = SAI_Model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30f864ae-e81a-4226-b24b-f732049f563b",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "30f864ae-e81a-4226-b24b-f732049f563b",
    "outputId": "87f2b765-2df0-4b6d-8615-bd94af466cc6"
   },
   "outputs": [],
   "source": [
    "def get_device():\n",
    "    if torch.cuda.is_available():\n",
    "        device = 'cuda:0'\n",
    "    else:\n",
    "        device = 'cpu'\n",
    "    return device\n",
    "\n",
    "device = get_device()\n",
    "print(device)\n",
    "model.to(device)\n",
    "summary(model,(81*3,64,64))\n",
    "\n",
    "Tensor = torch.cuda.FloatTensor if device == 'cuda:0' else torch.FloatTensor\n",
    "\n",
    "# Load pretrained models \n",
    "load_pretrained_models = True\n",
    "model_path = \"SAI_model_64_3L_best.pth\"\n",
    "if load_pretrained_models:\n",
    "    model.load_state_dict(torch.load(model_path,map_location=device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ZSqtrD6AAerG",
   "metadata": {
    "id": "ZSqtrD6AAerG"
   },
   "outputs": [],
   "source": [
    "class Sobel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Sobel, self).__init__()\n",
    "        self.edge_conv = nn.Conv2d(1, 2, kernel_size=3, stride=1, padding=1, bias=False)\n",
    "        edge_kx = np.array([[1, 0, -1], [2, 0, -2], [1, 0, -1]])\n",
    "        edge_ky = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])\n",
    "        edge_k = np.stack((edge_kx, edge_ky))\n",
    "\n",
    "        edge_k = torch.from_numpy(edge_k).float().view(2, 1, 3, 3)\n",
    "        self.edge_conv.weight = nn.Parameter(edge_k)\n",
    "        \n",
    "        for param in self.parameters():\n",
    "            param.requires_grad = False\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.edge_conv(x) \n",
    "        out = out.contiguous().view(-1, 2, x.size(2), x.size(3))\n",
    "  \n",
    "        return out\n",
    "\n",
    "get_gradient = Sobel().cuda()\n",
    "cos = nn.CosineSimilarity(dim=1, eps=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c5dc085-7d86-4b47-a65b-6fee2e8b24ab",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6c5dc085-7d86-4b47-a65b-6fee2e8b24ab",
    "outputId": "a8d66bb7-7784-489d-afe5-161213c2eb1c"
   },
   "outputs": [],
   "source": [
    "# Set nb of epoch for training (=1 for testing only)\n",
    "epoch = 1\n",
    "# Optimizer\n",
    "lr = 0.001\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=lr)#weight_decay=0.001\n",
    "\n",
    "val_min = 1\n",
    "\n",
    "start_time = time.time()\n",
    "\n",
    "train_losses=[]\n",
    "for epoch in range(1,epoch):  # loop over the dataset multiple times\n",
    "    model.train()\n",
    "    train_loss = 0.0\n",
    "    for i, data in enumerate(trainloader):#(trainloader,0)\n",
    "        # get the inputs; data is a list of [inputs, labels]\n",
    "        inputs, labels = data#[\"lr\"], data[\"GT\"]\n",
    "        inputs.cuda()\n",
    "        labels.cuda()\n",
    "\n",
    "        inputs = torch.autograd.Variable(inputs.type(Tensor))\n",
    "        labels = torch.autograd.Variable(labels.type(Tensor))\n",
    "        \n",
    "        ones = torch.ones(labels.size(0), 1, labels.size(2), labels.size(3)).float().cuda()\n",
    "        ones = torch.autograd.Variable(ones)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        output = model(inputs)\n",
    "        \n",
    "        depth_grad = get_gradient(labels)\n",
    "        output_grad = get_gradient(output)\n",
    "        depth_grad_dx = depth_grad[:, 0, :, :].contiguous().view_as(labels)\n",
    "        depth_grad_dy = depth_grad[:, 1, :, :].contiguous().view_as(labels)\n",
    "        output_grad_dx = output_grad[:, 0, :, :].contiguous().view_as(labels)\n",
    "        output_grad_dy = output_grad[:, 1, :, :].contiguous().view_as(labels)\n",
    "\n",
    "        depth_normal = torch.cat((-depth_grad_dx, -depth_grad_dy, ones), 1)\n",
    "        output_normal = torch.cat((-output_grad_dx, -output_grad_dy, ones), 1)\n",
    "\n",
    "        loss_1 = torch.abs(output - labels).mean() #criterion(output, labels)\n",
    "        loss_dx = torch.abs(output_grad_dx - depth_grad_dx).mean()\n",
    "        loss_dy = torch.abs(output_grad_dy - depth_grad_dy).mean()\n",
    "        loss_normal = torch.abs(1 - cos(output_normal, depth_normal)).mean()\n",
    "\n",
    "        loss = 1/3*loss_1 + 1/3*(loss_dx + loss_dy) + 1/3*loss_normal\n",
    "        \n",
    "        #loss = criterion(output, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        train_loss += loss.item()\n",
    "\n",
    "    train_loss = train_loss/len(trainloader)\n",
    "    # append training loss\n",
    "    train_losses.append(train_loss)\n",
    "    \n",
    "    #validation \n",
    "    model.eval()\n",
    "    val_losses=0\n",
    "    with torch.no_grad():\n",
    "        #Batch of test images\n",
    "        for j, imgs in enumerate(testloader):\n",
    "            images, GT = imgs#[\"lr\"], imgs[\"GT\"]\n",
    "            #Sample outputs\n",
    "            output = model(images.type(Tensor))\n",
    "\n",
    "            val_loss = torch.abs(output - GT.type(Tensor)).mean() \n",
    "            val_losses += val_loss.item()#*inputs.size(0)\n",
    "\n",
    "        val_losses = val_losses/len(testloader)\n",
    "        \n",
    "    print('Epoch: {} \\tTraining Loss: {:.5f}, \\tValidation Loss: {:.5f}'.format(epoch, train_loss, val_losses))\n",
    "    if val_losses < val_min:\n",
    "        torch.save(model.state_dict(), \"SAI_model_64_3L_best.pth\")\n",
    "        val_min = val_losses\n",
    "    \n",
    "    # Save model checkpoints\n",
    "    if (epoch)%5 == 0:\n",
    "        torch.save(model.state_dict(), \"SAI_model_64_3L.pth\")\n",
    "        #save model\n",
    "        torch.save(model.state_dict(), f\"SAI_model_64_3L_{epoch}.pth\")\n",
    "        print(f\"save model: {epoch}\")\n",
    "        \n",
    "print('Finished Training')\n",
    "print(\"--- %s seconds ---\" % (time.time() - start_time))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "vy-QgLqtA-s9",
   "metadata": {
    "id": "vy-QgLqtA-s9"
   },
   "outputs": [],
   "source": [
    "def Test(data, X, Y): \n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        Disparity_lr = np.zeros((512, 512))\n",
    "        GT_disp      = np.zeros((512, 512))\n",
    "        \n",
    "        for i in range(0,9):\n",
    "            for j in range(0,9):\n",
    "                if np.logical_and(i==0, j==0):\n",
    "                    Imgs = cv.imread(str('%s/%s%03d.png') %(data, \"input_Cam\", 0))\n",
    "                    #cropped_Imgs = Imgs[x_min:x_min+64, y_min:y_min+64]\n",
    "                else:\n",
    "                    img = cv.imread(str('%s/%s%03d.png') %(data, \"input_Cam\", i*9 + j))\n",
    "                    #cropped_img = img[x_min:x_min+64, y_min:y_min+64]\n",
    "                    Imgs = np.concatenate((Imgs, img), axis=-1)\n",
    "\n",
    "        for ii in range(8):\n",
    "            img_lr = torch.from_numpy(np.zeros((8, 81*3, 64, 64)).astype('float32'))\n",
    "            gt_lr  = np.zeros((8, 64, 64)).astype('float32')\n",
    "            for jj in range (8):\n",
    "                x_min=X[ii + jj*8] \n",
    "                y_min=Y[ii + jj*8]\n",
    "\n",
    "                cropped_Imgs = Imgs[x_min:x_min+64, y_min:y_min+64]\n",
    "\n",
    "                GT_disparity = read_pfm(str('%s/%s.pfm') %(data, \"gt_disp_lowres\"))\n",
    "                cropped_GT_disparity = GT_disparity[x_min:x_min+64, y_min:y_min+64]\n",
    "\n",
    "                cropped_Imgs_tensor = (torch.from_numpy(cropped_Imgs)).permute(2, 0, 1)\n",
    "\n",
    "                cropped_Imgs_tensor = (cropped_Imgs_tensor / 255.0)\n",
    "\n",
    "                img_lr[jj] = cropped_Imgs_tensor\n",
    "                gt_lr[jj]  = cropped_GT_disparity#_tensor\n",
    "\n",
    "            #Sample outputs\n",
    "            output = model(img_lr.type(Tensor))\n",
    "\n",
    "            output = output.cpu()\n",
    "            output = output.view(8, 1, 64, 64)\n",
    "            output = output.detach().numpy()\n",
    "            for k in range(8):\n",
    "                Disparity_lr[ii*64:(ii+1)*64, k*64:(k+1)*64] = np.squeeze(output[k,...])\n",
    "                GT_disp     [ii*64:(ii+1)*64, k*64:(k+1)*64] = gt_lr[k,...] #np.squeeze(gt_lr[k,...])\n",
    "\n",
    "        fig, axes = plt.subplots(nrows=1, ncols=2, sharex=True, sharey=True, figsize=(20,10))\n",
    "        ax_1 = fig.add_subplot(1, 2, 1, xticks=[], yticks=[])\n",
    "        plt.imshow(GT_disp, cmap ='gray')\n",
    "        ax_2 = fig.add_subplot(1, 2, 2, xticks=[], yticks=[])\n",
    "        plt.imshow(Disparity_lr, cmap ='gray')\n",
    "        plt.show()\n",
    "\n",
    "        fig, axes = plt.subplots(nrows=1, ncols=2, sharex=True, sharey=True, figsize=(20,10))\n",
    "        ax_1 = fig.add_subplot(1, 2, 1, xticks=[], yticks=[])\n",
    "        plt.imshow(GT_disp-Disparity_lr, cmap ='gray')\n",
    "        ax_2 = fig.add_subplot(1, 2, 2, xticks=[], yticks=[])\n",
    "        plt.imshow(np.abs(GT_disp-Disparity_lr)<0.07, cmap ='gray')\n",
    "        plt.show()\n",
    "\n",
    "        print(np.mean(np.abs(GT_disp-Disparity_lr)<0.07)*100)\n",
    "        print(np.mean(np.abs(GT_disp-Disparity_lr)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "svDtV8jC_UNW",
   "metadata": {
    "id": "svDtV8jC_UNW"
   },
   "outputs": [],
   "source": [
    "X = [0, 64, 128, 192, 256, 320, 384, 448,\n",
    "     0, 64, 128, 192, 256, 320, 384, 448,\n",
    "     0, 64, 128, 192, 256, 320, 384, 448,\n",
    "     0, 64, 128, 192, 256, 320, 384, 448,\n",
    "     0, 64, 128, 192, 256, 320, 384, 448,\n",
    "     0, 64, 128, 192, 256, 320, 384, 448,\n",
    "     0, 64, 128, 192, 256, 320, 384, 448,\n",
    "     0, 64, 128, 192, 256, 320, 384, 448]\n",
    "\n",
    "Y = [0, 0, 0, 0, 0, 0, 0, 0,     \n",
    "     64, 64, 64, 64, 64, 64, 64, 64,\n",
    "     128, 128, 128, 128, 128, 128, 128, 128,\n",
    "     192, 192, 192, 192, 192, 192, 192, 192,\n",
    "     256, 256, 256, 256, 256, 256, 256, 256,\n",
    "     320, 320, 320, 320, 320, 320, 320, 320, \n",
    "     384, 384, 384, 384, 384, 384, 384, 384,\n",
    "     448, 448, 448, 448, 448, 448, 448, 448]\n",
    "\n",
    "data = \"hci_dataset/training/cotton\"\n",
    "Test(data, X, Y)\n",
    "\n",
    "data = \"hci_dataset/training/dino\"\n",
    "Test(data, X, Y)\n",
    "\n",
    "data = \"hci_dataset/training/boxes\"\n",
    "Test(data, X, Y)\n",
    "\n",
    "data = \"hci_dataset/training/sideboard\"\n",
    "Test(data, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "yK0rNEZa_Whu",
   "metadata": {
    "id": "yK0rNEZa_Whu"
   },
   "outputs": [],
   "source": [
    "data = \"hci_dataset/stratified/backgammon\"\n",
    "Test(data, X, Y)\n",
    "\n",
    "data = \"hci_dataset/stratified/dots\"\n",
    "Test(data, X, Y)\n",
    "\n",
    "data = \"hci_dataset/stratified/pyramids\"\n",
    "Test(data, X, Y)\n",
    "\n",
    "data = \"hci_dataset/stratified/stripes\"\n",
    "Test(data, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2596a5a4-d09a-419c-b412-3815b66e4798",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "SAI_DL.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}