SIIE.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6VIDOcD4tWNj"
      },
      "source": [
        "**Sensor Independent Illumination Estimation for DNN Models**\n",
        "\n",
        "\n",
        "This is a PyTorch implementation of the paper: \"Sensor Independent Illumination Estimation for DNN Models.\" It is important to note that this implementation deviates from the original paper's code. For instance, it utilizes a lighter model architecture and employs different training parameters compared to those used in the original paper. Additionally, the code does not include the histogram learnable fall-off and scale factor introduced in the paper.\n",
        "\n",
        "It is important to note that this code has not been thoroughly tested, and there is no guarantee that it will produce similar results to the paper. Furthermore, the code is licensed under the same license as the main GitHub page, which can be found at the following link: https://github.com/mahmoudnafifi/SIIE\n",
        "\n",
        "\n",
        "Before you begin, please ensure to follow these instructions:\n",
        "\n",
        "\n",
        "1. Create a training folder and upload the training images in PNG-16 format. Each image should have a corresponding .JSON file with the postfix '_metadata.json'. Inside the JSON file, include a key called \"illuminant_color_raw\" that specifies the ground-truth illuminant color for that particular image. It is **essential** to capture the training images using different sensors or cameras, as mentioned in the paper. This variation in sensor sources will assist the neural network in learning implicit sensor mappings.\n",
        "\n",
        "\n",
        "2. Similarly, create a validation folder and upload the validation images using the same format and structure as described in step #1.\n",
        "\n",
        "3. Create a testing folder and upload the testing images in the same format and structure as mentioned in step #1. It is optional to provide the ground-truth information for the testing images.\n",
        "\n",
        "In case you encounter memory issues, consider implementing the following to mitigate the problem: 1) reduce the batch size, 2) decrease the input image size, and/or 3) trim model weights.\n",
        "\n",
        "\n",
        "If you find this code useful, please cite the following paper:\n",
        "\n",
        "Mahmoud Afifi and Michael S. Brown. \"Sensor Independent Illumination Estimation\n",
        "for DNN Models\". In BMVC, 2019.\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "```\n",
        "@inproceedings{afifi2019SIIE,\n",
        " title={Sensor-Independent Illumination Estimation for DNN Models},\n",
        " author={Afifi, Mahmoud and Brown, Michael S},\n",
        " booktitle={British Machine Vision Conference (BMVC)},\n",
        " pages={},\n",
        " year={2019}\n",
        "}\n",
        "```\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZJ1ZGv4osQ9r"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision.transforms as transforms\n",
        "from torch.utils.data import Dataset\n",
        "from torch.utils.data import DataLoader\n",
        "import numpy as np\n",
        "import cv2\n",
        "import os\n",
        "from os.path import join\n",
        "from os import listdir\n",
        "import numpy as np\n",
        "from glob import glob\n",
        "import logging\n",
        "import json\n",
        "from torch import optim\n",
        "from tqdm import tqdm\n",
        "import math\n",
        "import sys\n",
        "import pickle\n",
        "\n",
        "EPS = 1e-9\n",
        "PI = 22.0 / 7.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CPEKgnZnwfEs"
      },
      "source": [
        "Input/output information"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WBkRBGo-whpY"
      },
      "outputs": [],
      "source": [
        "trdir = 'drive/training'  # training dir\n",
        "valdir = 'drive/validation'  # validation dir\n",
        "testing_dir = 'drive/testing'  # testing dir\n",
        "model_dir = 'drive/models/'  # directory to save final model\n",
        "dir_checkpoint = 'drive/checkpoints/'  # directory to save checkpoints\n",
        "result_dir = 'drive/result/'  # directory to save testing results\n",
        "initial_est = False\n",
        "epochs = 120  # number of epochs\n",
        "batchsize = 16 # batch size\n",
        "lr = 0.00001 # learning rate\n",
        "l2reg = 0.0001 # L2 regularization factor\n",
        "max_img_size = 120\n",
        "validation_freq = 2  # validation frequency\n",
        "chkpointperiod = 2  # checkpoint frequency"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YkoffVlRskVS"
      },
      "source": [
        "Histogram block"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iHzMrp_jsjlo"
      },
      "outputs": [],
      "source": [
        "\"\"\"\n",
        " If you find this histogram block useful, please cite our papers:\n",
        "\n",
        " Mahmoud Afifi and Michael S. Brown. \"Sensor Independent Illumination Estimation\n",
        " for DNN Models\". In BMVC, 2019.\n",
        "\n",
        " Mahmoud Afifi, Marcus A. Brubaker, and Michael S. Brown. \"HistoGAN:\n",
        " Controlling Colors of GAN-Generated and Real Images via Color Histograms.\"\n",
        " In CVPR, 2021.\n",
        "\n",
        " @inproceedings{afifi2019SIIE,\n",
        "  title={Sensor-Independent Illumination Estimation for DNN Models},\n",
        "  author={Afifi, Mahmoud and Brown, Michael S},\n",
        "  booktitle={British Machine Vision Conference (BMVC)},\n",
        "  pages={},\n",
        "  year={2019}\n",
        " }\n",
        "\n",
        " @inproceedings{afifi2021histogan,\n",
        "  title={Histo{GAN}: Controlling Colors of {GAN}-Generated and Real Images via\n",
        "  Color Histograms},\n",
        "  author={Afifi, Mahmoud and Brubaker, Marcus A. and Brown, Michael S.},\n",
        "  booktitle={CVPR},\n",
        "  year={2021}\n",
        "}\n",
        "\"\"\"\n",
        "\n",
        "class RGBuvHistBlock(nn.Module):\n",
        "  def __init__(self, h=64, insz=max_img_size, resizing='interpolation',\n",
        "               method='inverse-quadratic', sigma=0.02, intensity_scale=True,\n",
        "               hist_boundary=None, green_only=False, device='cuda'):\n",
        "    \"\"\" Computes the RGB-uv histogram feature of a given image.\n",
        "    Args:\n",
        "      h: histogram dimension size (scalar). The default value is 64.\n",
        "      insz: maximum size of the input image; if it is larger than this size, the\n",
        "        image will be resized (scalar).\n",
        "      resizing: resizing method if applicable. Options are: 'interpolation' or\n",
        "        'sampling'. Default is 'interpolation'.\n",
        "      method: the method used to count the number of pixels for each bin in the\n",
        "        histogram feature. Options are: 'thresholding', 'RBF' (radial basis\n",
        "        function), or 'inverse-quadratic'. Default value is 'inverse-quadratic'.\n",
        "      sigma: if the method value is 'RBF' or 'inverse-quadratic', then this is\n",
        "        the learnable sigma parameter of the kernel function.\n",
        "      intensity_scale: boolean variable to use the intensity scale (I_y in\n",
        "        Equation 2). Default value is True.\n",
        "      hist_boundary: a list of histogram boundary values.\n",
        "      green_only: boolean variable to use only the log(g/r), log(g/b) channels.\n",
        "        Default is False.\n",
        "\n",
        "    Methods:\n",
        "      forward: accepts input image and returns its histogram feature. Note that\n",
        "        unless the method is 'thresholding', this is a differentiable function\n",
        "        and can be easily integrated with the loss function. As mentioned in the\n",
        "         paper, the 'inverse-quadratic' was found more stable than 'RBF' in our\n",
        "         training.\n",
        "    \"\"\"\n",
        "    super(RGBuvHistBlock, self).__init__()\n",
        "    self.h = h\n",
        "    self.insz = insz\n",
        "    self.device = device\n",
        "    self.resizing = resizing\n",
        "    self.method = method\n",
        "    self.intensity_scale = intensity_scale\n",
        "    self.green_only = green_only\n",
        "    if hist_boundary is None:\n",
        "      hist_boundary = [-2.6, 2.6]\n",
        "    hist_boundary.sort()\n",
        "    self.hist_boundary = hist_boundary\n",
        "    if self.method == 'thresholding':\n",
        "      self.eps = (abs(hist_boundary[0]) + abs(hist_boundary [1])) / h\n",
        "    else:\n",
        "      self.sigma = sigma\n",
        "\n",
        "  def forward(self, x):\n",
        "    x = torch.clamp(x, 0, 1)\n",
        "    if x.shape[2] > self.insz or x.shape[3] > self.insz:\n",
        "      if self.resizing == 'interpolation':\n",
        "        x_sampled = F.interpolate(x, size=(self.insz, self.insz),\n",
        "                                  mode='bilinear', align_corners=False)\n",
        "      elif self.resizing == 'sampling':\n",
        "        inds_1 = torch.LongTensor(\n",
        "          np.linspace(0, x.shape[2], self.h, endpoint=False)).to(\n",
        "          device=self.device)\n",
        "        inds_2 = torch.LongTensor(\n",
        "          np.linspace(0, x.shape[3], self.h, endpoint=False)).to(\n",
        "          device=self.device)\n",
        "        x_sampled = x.index_select(2, inds_1)\n",
        "        x_sampled = x_sampled.index_select(3, inds_2)\n",
        "      else:\n",
        "        raise Exception(\n",
        "          f'Wrong resizing method. It should be: interpolation or sampling. '\n",
        "          f'But the given value is {self.resizing}.')\n",
        "    else:\n",
        "      x_sampled = x\n",
        "\n",
        "    L = x_sampled.shape[0]  # size of mini-batch\n",
        "    if x_sampled.shape[1] > 3:\n",
        "      x_sampled = x_sampled[:, :3, :, :]\n",
        "    X = torch.unbind(x_sampled, dim=0)\n",
        "    hists = torch.zeros((x_sampled.shape[0], 1 + int(not self.green_only) * 2,\n",
        "                         self.h, self.h)).to(device=self.device)\n",
        "    for l in range(L):\n",
        "      I = torch.t(torch.reshape(X[l], (3, -1)))\n",
        "      II = torch.pow(I, 2)\n",
        "      if self.intensity_scale:\n",
        "        Iy = torch.unsqueeze(torch.sqrt(II[:, 0] + II[:, 1] + II[:, 2] + EPS),\n",
        "                             dim=1)\n",
        "      else:\n",
        "        Iy = 1\n",
        "      if not self.green_only:\n",
        "        Iu0 = torch.unsqueeze(torch.log(I[:, 0] + EPS) - torch.log(I[:, 1] +\n",
        "                                                                   EPS), dim=1)\n",
        "        Iv0 = torch.unsqueeze(torch.log(I[:, 0] + EPS) - torch.log(I[:, 2] +\n",
        "                                                                   EPS), dim=1)\n",
        "        diff_u0 = abs(\n",
        "          Iu0 - torch.unsqueeze(torch.tensor(np.linspace(\n",
        "              self.hist_boundary[0], self.hist_boundary[1], num=self.h)),\n",
        "              dim=0).to(self.device))\n",
        "        diff_v0 = abs(\n",
        "          Iv0 - torch.unsqueeze(torch.tensor(np.linspace(\n",
        "              self.hist_boundary[0], self.hist_boundary[1], num=self.h)),\n",
        "              dim=0).to(self.device))\n",
        "        if self.method == 'thresholding':\n",
        "          diff_u0 = torch.reshape(diff_u0, (-1, self.h)) <= self.eps / 2\n",
        "          diff_v0 = torch.reshape(diff_v0, (-1, self.h)) <= self.eps / 2\n",
        "        elif self.method == 'RBF':\n",
        "          diff_u0 = torch.pow(torch.reshape(diff_u0, (-1, self.h)),\n",
        "                              2) / self.sigma ** 2\n",
        "          diff_v0 = torch.pow(torch.reshape(diff_v0, (-1, self.h)),\n",
        "                              2) / self.sigma ** 2\n",
        "          diff_u0 = torch.exp(-diff_u0)  # Radial basis function\n",
        "          diff_v0 = torch.exp(-diff_v0)\n",
        "        elif self.method == 'inverse-quadratic':\n",
        "          diff_u0 = torch.pow(torch.reshape(diff_u0, (-1, self.h)),\n",
        "                              2) / self.sigma ** 2\n",
        "          diff_v0 = torch.pow(torch.reshape(diff_v0, (-1, self.h)),\n",
        "                              2) / self.sigma ** 2\n",
        "          diff_u0 = 1 / (1 + diff_u0)  # Inverse quadratic\n",
        "          diff_v0 = 1 / (1 + diff_v0)\n",
        "        else:\n",
        "          raise Exception(\n",
        "            f'Wrong kernel method. It should be either thresholding, RBF,'\n",
        "            f' inverse-quadratic. But the given value is {self.method}.')\n",
        "        diff_u0 = diff_u0.type(torch.float32)\n",
        "        diff_v0 = diff_v0.type(torch.float32)\n",
        "        a = torch.t(Iy * diff_u0)\n",
        "        hists[l, 0, :, :] = torch.mm(a, diff_v0)\n",
        "\n",
        "      Iu1 = torch.unsqueeze(torch.log(I[:, 1] + EPS) - torch.log(I[:, 0] + EPS),\n",
        "                            dim=1)\n",
        "      Iv1 = torch.unsqueeze(torch.log(I[:, 1] + EPS) - torch.log(I[:, 2] + EPS),\n",
        "                            dim=1)\n",
        "      diff_u1 = abs(\n",
        "        Iu1 - torch.unsqueeze(torch.tensor(np.linspace(\n",
        "            self.hist_boundary[0], self.hist_boundary[1], num=self.h)),\n",
        "            dim=0).to(self.device))\n",
        "      diff_v1 = abs(\n",
        "        Iv1 - torch.unsqueeze(torch.tensor(np.linspace(\n",
        "            self.hist_boundary[0], self.hist_boundary[1], num=self.h)),\n",
        "            dim=0).to(self.device))\n",
        "\n",
        "      if self.method == 'thresholding':\n",
        "        diff_u1 = torch.reshape(diff_u1, (-1, self.h)) <= self.eps / 2\n",
        "        diff_v1 = torch.reshape(diff_v1, (-1, self.h)) <= self.eps / 2\n",
        "      elif self.method == 'RBF':\n",
        "        diff_u1 = torch.pow(torch.reshape(diff_u1, (-1, self.h)),\n",
        "                            2) / self.sigma ** 2\n",
        "        diff_v1 = torch.pow(torch.reshape(diff_v1, (-1, self.h)),\n",
        "                            2) / self.sigma ** 2\n",
        "        diff_u1 = torch.exp(-diff_u1)  # Gaussian\n",
        "        diff_v1 = torch.exp(-diff_v1)\n",
        "      elif self.method == 'inverse-quadratic':\n",
        "        diff_u1 = torch.pow(torch.reshape(diff_u1, (-1, self.h)),\n",
        "                            2) / self.sigma ** 2\n",
        "        diff_v1 = torch.pow(torch.reshape(diff_v1, (-1, self.h)),\n",
        "                            2) / self.sigma ** 2\n",
        "        diff_u1 = 1 / (1 + diff_u1)  # Inverse quadratic\n",
        "        diff_v1 = 1 / (1 + diff_v1)\n",
        "\n",
        "      diff_u1 = diff_u1.type(torch.float32)\n",
        "      diff_v1 = diff_v1.type(torch.float32)\n",
        "      a = torch.t(Iy * diff_u1)\n",
        "      if not self.green_only:\n",
        "        hists[l, 1, :, :] = torch.mm(a, diff_v1)\n",
        "      else:\n",
        "        hists[l, 0, :, :] = torch.mm(a, diff_v1)\n",
        "\n",
        "      if not self.green_only:\n",
        "        Iu2 = torch.unsqueeze(torch.log(I[:, 2] + EPS) - torch.log(I[:, 0] +\n",
        "                                                                   EPS), dim=1)\n",
        "        Iv2 = torch.unsqueeze(torch.log(I[:, 2] + EPS) - torch.log(I[:, 1] +\n",
        "                                                                   EPS), dim=1)\n",
        "        diff_u2 = abs(\n",
        "          Iu2 - torch.unsqueeze(torch.tensor(np.linspace(\n",
        "              self.hist_boundary[0], self.hist_boundary[1], num=self.h)),\n",
        "              dim=0).to(self.device))\n",
        "        diff_v2 = abs(\n",
        "          Iv2 - torch.unsqueeze(torch.tensor(np.linspace(\n",
        "              self.hist_boundary[0], self.hist_boundary[1], num=self.h)),\n",
        "              dim=0).to(self.device))\n",
        "        if self.method == 'thresholding':\n",
        "          diff_u2 = torch.reshape(diff_u2, (-1, self.h)) <= self.eps / 2\n",
        "          diff_v2 = torch.reshape(diff_v2, (-1, self.h)) <= self.eps / 2\n",
        "        elif self.method == 'RBF':\n",
        "          diff_u2 = torch.pow(torch.reshape(diff_u2, (-1, self.h)),\n",
        "                              2) / self.sigma ** 2\n",
        "          diff_v2 = torch.pow(torch.reshape(diff_v2, (-1, self.h)),\n",
        "                              2) / self.sigma ** 2\n",
        "          diff_u2 = torch.exp(-diff_u2)  # Gaussian\n",
        "          diff_v2 = torch.exp(-diff_v2)\n",
        "        elif self.method == 'inverse-quadratic':\n",
        "          diff_u2 = torch.pow(torch.reshape(diff_u2, (-1, self.h)),\n",
        "                              2) / self.sigma ** 2\n",
        "          diff_v2 = torch.pow(torch.reshape(diff_v2, (-1, self.h)),\n",
        "                              2) / self.sigma ** 2\n",
        "          diff_u2 = 1 / (1 + diff_u2)  # Inverse quadratic\n",
        "          diff_v2 = 1 / (1 + diff_v2)\n",
        "        diff_u2 = diff_u2.type(torch.float32)\n",
        "        diff_v2 = diff_v2.type(torch.float32)\n",
        "        a = torch.t(Iy * diff_u2)\n",
        "        hists[l, 2, :, :] = torch.mm(a, diff_v2)\n",
        "\n",
        "    # normalization\n",
        "    hists_normalized = hists / (\n",
        "        ((hists.sum(dim=1)).sum(dim=1)).sum(dim=1).view(-1, 1, 1, 1) + EPS)\n",
        "\n",
        "    return hists_normalized"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U0HbjOvrtNiy"
      },
      "source": [
        "Data loader"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0LWy4GijtOYT"
      },
      "outputs": [],
      "source": [
        "from types import prepare_class\n",
        "class BasicDataset(Dataset):\n",
        "    def __init__(self, imgs_dir):\n",
        "        self.imgs_dir = imgs_dir\n",
        "        self.imgfiles = [join(imgs_dir, file) for file in listdir(imgs_dir)\n",
        "                        if file.endswith('.png') and not file.startswith('.')]\n",
        "        logging.info(f'Creating dataset with {len(self.imgfiles)} examples')\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.imgfiles)\n",
        "\n",
        "    @classmethod\n",
        "    def preprocess(cls, img):\n",
        "        img = cv2.resize(img, (max_img_size, max_img_size))\n",
        "        # HWC to CHW\n",
        "        img_chw = img.transpose((2, 0, 1))\n",
        "        return img_chw\n",
        "\n",
        "    def __getitem__(self, i):\n",
        "        gt_ext = ('_metadata.json')\n",
        "        img_file = self.imgfiles[i]\n",
        "        gt_file = img_file.replace('.png', gt_ext)\n",
        "        img = cv2.imread(img_file, -1) / ((2 ** 16) - 1)\n",
        "        img = img[..., ::-1]\n",
        "        img = self.preprocess(img)\n",
        "        with open(gt_file, 'r') as f:\n",
        "          data = json.load(f)\n",
        "          gt_ill = np.array(data['illuminant_color_raw'])\n",
        "        return {'image': torch.from_numpy(img.copy()),\n",
        "                'gt_ill': torch.from_numpy(gt_ill)}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L9l_xV3TztcY"
      },
      "source": [
        "Loss function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qy6U3rx7zuSl"
      },
      "outputs": [],
      "source": [
        "def angular_loss(predicted, gt, shrink=True):\n",
        "  \"\"\"Computes angular error between predicted and gt illuminant color(s).\"\"\"\n",
        "  cossim = torch.clamp(torch.sum(predicted * gt, dim=1) / (\n",
        "      torch.norm(predicted, dim=1) * torch.norm(gt, dim=1) + EPS), -1, 1.)\n",
        "  if shrink:\n",
        "    angle = torch.acos(cossim * 0.999999)\n",
        "  else:\n",
        "    angle = torch.acos(cossim)\n",
        "  a_error = 180.0 / PI * angle\n",
        "  a_error = torch.sum(a_error) / a_error.shape[0]\n",
        "  return a_error"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "12R6L9H1z_qr"
      },
      "source": [
        "Network"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NzGn36M40AZT"
      },
      "outputs": [],
      "source": [
        "class Flatten(nn.Module):\n",
        "    def forward(self, x):\n",
        "        x = x.view(x.size()[0], -1)\n",
        "        return x\n",
        "\n",
        "class Net(nn.Module):\n",
        "    def __init__(self, output_chs=3, device='cude'):\n",
        "      \"\"\"Network (sensor mapping / illuminant estimation) constructor.\"\"\"\n",
        "      super(Net, self).__init__()\n",
        "      self._net = nn.Sequential(\n",
        "          nn.Conv2d(3, 92, kernel_size=5, padding=0, stride=2),\n",
        "          nn.ReLU(inplace=True),\n",
        "          nn.Conv2d(92, 128, kernel_size=3, padding=0, stride=2),\n",
        "          nn.ReLU(inplace=True),\n",
        "          nn.Conv2d(128, 256, kernel_size=2, padding=0, stride=1),\n",
        "          nn.ReLU(inplace=True),\n",
        "          Flatten(),\n",
        "          torch.nn.Linear(13*13*256, output_chs, bias=True))\n",
        "      self._hist_block = RGBuvHistBlock(insz=max_img_size, h=61)\n",
        "\n",
        "    def forward(self, img):\n",
        "      \"\"\"Forward function.\"\"\"\n",
        "      return self._net(self._hist_block(img))\n",
        "\n",
        "\n",
        "class SIIE(nn.Module):\n",
        "  def __init__(self, device='cuda', initial_est=True):\n",
        "    \"\"\"SIIE net constructor.\"\"\"\n",
        "    super(SIIE, self).__init__()\n",
        "    self.device = device\n",
        "    self._initial_est = initial_est\n",
        "\n",
        "    # Sensor mapping net\n",
        "    self._sensor_mapping_net = Net(output_chs=9)\n",
        "    self._sensor_mapping_net.to(device=device)\n",
        "\n",
        "    # Illuminant estimation net\n",
        "    self._illum_est_net = Net(output_chs=3)\n",
        "\n",
        "  def forward(self, img):\n",
        "    \"\"\"Forward function of SIIE net.\"\"\"\n",
        "    if self._initial_est:\n",
        "      initial_ill = torch.mean(img, dim=[2, 3])\n",
        "      gains = torch.unsqueeze(initial_ill[:, 1], dim=-1) / initial_ill\n",
        "      img[:, 0, ...] *= torch.unsqueeze(torch.unsqueeze(gains[:, 0], axis=-1),\n",
        "                                        axis=-1)\n",
        "      img[:, 2, ...] *= torch.unsqueeze(torch.unsqueeze(gains[:, 2], axis=-1),\n",
        "                                        axis=-1)\n",
        "\n",
        "    m = self._sensor_mapping_net(img)\n",
        "    m = m.reshape([-1, 3, 3])\n",
        "    m = torch.abs(m) / (\n",
        "        torch.unsqueeze(torch.sum(torch.abs(m), axis=-1), axis=-1) + EPS)\n",
        "    img_mapped = img.clone()\n",
        "    for i in range(m.shape[0]):\n",
        "      temp = torch.mm(torch.squeeze(m[i, :, :]),\n",
        "                      torch.reshape(torch.squeeze(\n",
        "                          img[i, :, :, :]), (3, -1)))\n",
        "      img_mapped[i, :, :, :] = torch.reshape(temp, (img.shape[1], img.shape[2],\n",
        "                                                    img.shape[3]))\n",
        "\n",
        "    ill_mapped = self._illum_est_net(img_mapped)\n",
        "    inds = (torch.linalg.matrix_rank(m) < 3).to(device=m.device)\n",
        "    offset = torch.zeros(m.shape[0], 1).to(device=m.device)\n",
        "    if torch.sum(inds) > 0:\n",
        "      offset[inds, ...] = torch.unsqueeze(\n",
        "          torch.rand(torch.sum(inds)), axis=-1).to(device=m.device) * 0.0001\n",
        "    inv_m = torch.linalg.inv(m + torch.unsqueeze(offset, axis=-1))\n",
        "    ill = ill_mapped.clone()\n",
        "    for i in range(inv_m.shape[0]):\n",
        "      temp = torch.mm(torch.squeeze(inv_m[i, :, :]),\n",
        "                      torch.unsqueeze(ill_mapped[i, :], axis=-1))\n",
        "      ill[i, :] = torch.reshape(temp, ill[i, :].shape)\n",
        "    if self._initial_est:\n",
        "      ill = initial_ill * ill\n",
        "    return ill, m"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "elapHVybFW_Y"
      },
      "source": [
        "Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kQtEqAGkFZRE"
      },
      "outputs": [],
      "source": [
        "def train_net(net, device, train_dir, val_dir, epochs=300,\n",
        "              batch_size=8, lr=0.00001, l2reg=0.0001,\n",
        "              chkpointperiod=20, save_cp=True):\n",
        "    retain_graph=True\n",
        "    train = BasicDataset(train_dir)\n",
        "    val = BasicDataset(val_dir)\n",
        "    train_loader = DataLoader(train, batch_size=batch_size, shuffle=True,\n",
        "                              num_workers=2, pin_memory=True)\n",
        "    val_loader = DataLoader(val, batch_size=batch_size, shuffle=False,\n",
        "                            num_workers=2, pin_memory=True, drop_last=True)\n",
        "    global_step = 0\n",
        "\n",
        "    logging.info(f'''Starting training:\n",
        "        Epochs:          {epochs} epochs\n",
        "        Batch size:      {batch_size}\n",
        "        Learning rate:   {lr}\n",
        "        Training size:   {len(train)}\n",
        "        Validation size: {len(val)}\n",
        "        Checkpoints:     {save_cp}\n",
        "        Device:          {device.type}\n",
        "    ''')\n",
        "\n",
        "    optimizer = optim.Adam(net.parameters(), lr=lr, betas=(0.85, 0.999),\n",
        "                           eps=1e-08, weight_decay=l2reg)\n",
        "    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs)\n",
        "\n",
        "    for epoch in range(epochs):\n",
        "        net.train()\n",
        "\n",
        "        epoch_loss = 0\n",
        "        with tqdm(total=len(train), desc=f'Epoch {epoch + 1}/{epochs}',\n",
        "                  unit='img') as pbar:\n",
        "            for batch in train_loader:\n",
        "                imgs = batch['image']\n",
        "                gt_ill = batch['gt_ill']\n",
        "                imgs = imgs.to(device=device, dtype=torch.float32)\n",
        "                gt_ill = gt_ill.to(device=device, dtype=torch.float32)\n",
        "                est_ill, _ = net(imgs)\n",
        "                loss = angular_loss(est_ill, gt_ill)\n",
        "                epoch_loss += loss.item()\n",
        "                pbar.set_postfix(**{'loss (batch)': loss.item()})\n",
        "                optimizer.zero_grad()\n",
        "                loss.backward()\n",
        "                optimizer.step()\n",
        "                pbar.update(np.ceil(imgs.shape[0]))\n",
        "                global_step += 1\n",
        "        if (epoch + 1) % validation_freq == 0:\n",
        "          val_score = vald_net(net, val_loader, device)\n",
        "          print(f'~~~~ Validation loss: {val_score} ~~~~~')\n",
        "\n",
        "        if save_cp and (epoch + 1) % chkpointperiod == 0:\n",
        "            if not os.path.exists(dir_checkpoint):\n",
        "                os.mkdir(dir_checkpoint)\n",
        "                logging.info('Created checkpoint directory')\n",
        "\n",
        "            torch.save(net.state_dict(), dir_checkpoint +\n",
        "                       f'siie_net_{epoch + 1}.pth')\n",
        "            logging.info(f'Checkpoint {epoch + 1} saved!')\n",
        "        scheduler.step()\n",
        "\n",
        "    if not os.path.exists(model_dir):\n",
        "        os.mkdir(model_dir)\n",
        "        logging.info('Created trained models directory')\n",
        "    torch.save(net.state_dict(), model_dir + 'siie_net.pth')\n",
        "    logging.info('Saved trained model!')\n",
        "    logging.info('End of training')\n",
        "\n",
        "\n",
        "def vald_net(net, loader, device):\n",
        "    \"\"\"Evaluation using MAE\"\"\"\n",
        "    net.eval()\n",
        "    n_val = len(loader) + 1\n",
        "    mae = 0\n",
        "\n",
        "    with tqdm(total=n_val, desc='Validation round', unit='batch',\n",
        "              leave=False) as pbar:\n",
        "        for batch in loader:\n",
        "            imgs = batch['image']\n",
        "            gt_ill = batch['gt_ill']\n",
        "            imgs = imgs.to(device=device, dtype=torch.float32)\n",
        "            gt_ill = gt_ill.to(device=device, dtype=torch.float32)\n",
        "\n",
        "            with torch.no_grad():\n",
        "                est_ill, m = net(imgs)\n",
        "                loss = angular_loss(gt_ill, est_ill)\n",
        "                mae = mae + loss\n",
        "\n",
        "            pbar.update(np.ceil(imgs.shape[0]))\n",
        "    net.train()\n",
        "    return mae / n_val\n",
        "\n",
        "\n",
        "logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')\n",
        "logging.info('Training of SIIE Net')\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "logging.info(f'Using device {device}')\n",
        "net = SIIE(device=device, initial_est=initial_est)\n",
        "net.to(device=device)\n",
        "\n",
        "try:\n",
        "  train_net(net=net, device=device, train_dir=trdir,\n",
        "            val_dir=valdir, epochs=epochs, l2reg=l2reg,\n",
        "            batch_size=batchsize, lr=lr, chkpointperiod=chkpointperiod)\n",
        "except KeyboardInterrupt:\n",
        "    torch.save(net.state_dict(), 'intrrupted_check_point.pth')\n",
        "    logging.info('Saved interrupt checkpoint backup')\n",
        "    try:\n",
        "        sys.exit(0)\n",
        "    except SystemExit:\n",
        "        os._exit(0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "brYED_jIIW4Y"
      },
      "source": [
        "Testing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F3QJkL1RIZuu"
      },
      "outputs": [],
      "source": [
        "def test_net(net, device, test_dir):\n",
        "    gt_ext = ('_metadata.json')\n",
        "    test = BasicDataset(test_dir)\n",
        "\n",
        "    imgfiles = [join(test_dir, file) for file in listdir(test_dir)\n",
        "                if file.endswith('.png') and not file.startswith('.')]\n",
        "    logging.info(f'Testing {len(imgfiles)} images')\n",
        "    ae = np.ones((len(imgfiles), 1)) * -1\n",
        "    estimated_ills = np.zeros((len(imgfiles), 3))\n",
        "    gt_ills = np.zeros((len(imgfiles), 3))\n",
        "    filenames = []\n",
        "    for i, f in enumerate(imgfiles):\n",
        "      filenames.append(f)\n",
        "      gt_file = f.replace('.png', gt_ext)\n",
        "      img = cv2.imread(f, -1) / ((2 ** 16) - 1)\n",
        "      img = img[..., ::-1]\n",
        "      if os.path.exists(gt_file):\n",
        "        with open(gt_file, 'r') as f:\n",
        "          data = json.load(f)\n",
        "          gt_ill = np.array(data['illuminant_color_raw'])\n",
        "          gt_ill = torch.from_numpy(gt_ill)\n",
        "      else:\n",
        "        gt_ill = None\n",
        "      img = cv2.resize(img, (max_img_size, max_img_size))\n",
        "      # HWC to CHW\n",
        "      img_chw = img.transpose((2, 0, 1))\n",
        "      in_img = torch.from_numpy(img_chw.copy())\n",
        "      in_img = in_img.to(device=device, dtype=torch.float32)\n",
        "      in_img = torch.unsqueeze(in_img, axis=0)\n",
        "      est_ill, _ = net(in_img)\n",
        "      if gt_ill is not None:\n",
        "        gt_ill = gt_ill.to(device=device, dtype=torch.float32)\n",
        "        gt_ill = torch.unsqueeze(gt_ill, axis=0)\n",
        "        loss = angular_loss(gt_ill, est_ill, shrink=False)\n",
        "        ae[i] = loss.detach().cpu().numpy()\n",
        "        gt = gt_ill.detach().cpu().numpy()\n",
        "        gt = gt / np.sqrt(np.sum(gt ** 2))\n",
        "        gt_ills[i, :] = gt.reshape(1, 3)\n",
        "      result = est_ill.detach().cpu().numpy()\n",
        "      result = result / np.sqrt(np.sum(result ** 2))\n",
        "      estimated_ills[i, :] = result\n",
        "    return estimated_ills, ae, filenames, gt_ills\n",
        "\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "logging.info(f'Using device {device}')\n",
        "net = SIIE(device=device, initial_est=initial_est)\n",
        "net.to(device=device)\n",
        "net.load_state_dict(torch.load(os.path.join(model_dir, 'siie_net.pth'),\n",
        "                                               map_location=device))\n",
        "estimated_illuminants, ae, filenames, gt_ill = test_net(net, device,\n",
        "                                                        testing_dir)\n",
        "data = {'filenames': filenames,\n",
        "        'estiamted_ills': estimated_illuminants.tolist(),\n",
        "        'gt_ills': gt_ill.tolist(), 'ae': ae.tolist()}\n",
        "json_object = json.dumps(data, indent=4)\n",
        "# Writing to sample.json\n",
        "if not os.path.exists(result_dir):\n",
        "  os.mkdir(result_dir)\n",
        "with open(os.path.join(result_dir, \"results.json\"), \"w\") as outfile:\n",
        "    outfile.write(json_object)\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}