Colorization_baseline_with_incpetionv3.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deep Colorization\n",
    "### Deep learning final project for conversion of gray scale images to rgb\n",
    "### Contributors: Bhumi Bhanushali, Avinash Hemaeshwara Raju, Kathan Nilesh Mehta, Atulya Ravishankar"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# wget -N images.cocodataset.org/zips/train2017.zip\n",
    "# wget -N images.cocodataset.org/zips/val2017.zip\n",
    "# wget -N images.cocodataset.org/zips/test2017.zip\n",
    "# pip3 install tensorboard\n",
    "# tensorboard --logdir=runs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import numpy as np \n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import torchvision.models as models\n",
    "from torchvision.utils import save_image\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Configuration:\n",
    "    model_file_name = 'checkpoint'\n",
    "    load_model_to_train = False\n",
    "    load_model_to_test = True\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    point_batches = 500"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyper Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HyperParameters:\n",
    "    epochs = 1\n",
    "    batch_size = 32\n",
    "    learning_rate = 0.001\n",
    "    num_workers = 16\n",
    "    learning_rate_decay = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = Configuration()\n",
    "hparams = HyperParameters()\n",
    "print('Device:',config.device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, root_dir, process_type):\n",
    "        self.root_dir = root_dir\n",
    "        self.files = [f for f in os.listdir(root_dir)]\n",
    "        self.process_type = process_type\n",
    "        print('File[0]:',self.files[0],'| Total Files:', len(self.files), '| Process:',self.process_type,)\n",
    "\n",
    "    def __len__(self):\n",
    "        return 100#len(self.files)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        try:\n",
    "            #*** Read the image from file ***\n",
    "            self.rgb_img = cv2.imread(os.path.join(self.root_dir,self.files[index]))\n",
    "            \n",
    "            if self.rgb_img is None:\n",
    "                raise Exception\n",
    "\n",
    "            self.rgb_img = self.rgb_img.astype(np.float32) \n",
    "            self.rgb_img /= 255.0 \n",
    "            \n",
    "            #*** Resize the color image to pass to encoder ***\n",
    "            rgb_encoder_img = cv2.resize(self.rgb_img, (224, 224))\n",
    "            \n",
    "            #*** Resize the color image to pass to decoder ***\n",
    "            rgb_inception_img = cv2.resize(self.rgb_img, (300, 300))\n",
    "            \n",
    "            ''' Encoder Images '''\n",
    "            #*** Convert the encoder color image to normalized lab space ***\n",
    "            self.lab_encoder_img = cv2.cvtColor(rgb_encoder_img,cv2.COLOR_BGR2Lab) \n",
    "            \n",
    "            #*** Splitting the lab images into l-channel, a-channel, b-channel ***\n",
    "            l_encoder_img, a_encoder_img, b_encoder_img = self.lab_encoder_img[:,:,0],self.lab_encoder_img[:,:,1],self.lab_encoder_img[:,:,2]\n",
    "            \n",
    "            #*** Normalizing l-channel between [-1,1] ***\n",
    "            l_encoder_img = l_encoder_img/50.0 - 1.0\n",
    "            \n",
    "            #*** Repeat the l-channel to 3 dimensions ***\n",
    "            l_encoder_img = torchvision.transforms.ToTensor()(l_encoder_img)\n",
    "            l_encoder_img = l_encoder_img.expand(3,-1,-1)\n",
    "            \n",
    "            #*** Normalize a and b channels and concatenate ***\n",
    "            a_encoder_img = (a_encoder_img/128.0)\n",
    "            b_encoder_img = (b_encoder_img/128.0)\n",
    "            a_encoder_img = torch.stack([torch.Tensor(a_encoder_img)])\n",
    "            b_encoder_img = torch.stack([torch.Tensor(b_encoder_img)])\n",
    "            ab_encoder_img = torch.cat([a_encoder_img, b_encoder_img], dim=0)\n",
    "            \n",
    "            ''' Inception Images '''\n",
    "            #*** Convert the inception color image to lab space ***\n",
    "            self.lab_inception_img = cv2.cvtColor(rgb_inception_img,cv2.COLOR_BGR2Lab)\n",
    "            \n",
    "            #*** Extract the l-channel of inception lab image *** \n",
    "            l_inception_img = self.lab_inception_img[:,:,0]/50.0 - 1.0\n",
    "             \n",
    "            #*** Convert the inception l-image to torch Tensor and stack it in 3 channels ***\n",
    "            l_inception_img = torchvision.transforms.ToTensor()(l_inception_img)\n",
    "            l_inception_img = l_inception_img.expand(3,-1,-1)\n",
    "            \n",
    "            ''' return images to data-loader '''\n",
    "            rgb_encoder_img = torchvision.transforms.ToTensor()(rgb_encoder_img)\n",
    "            return l_encoder_img, ab_encoder_img, l_inception_img, rgb_encoder_img, self.files[index]\n",
    "        \n",
    "        except Exception as e:\n",
    "            print('Exception at ',self.files[index], e)\n",
    "            return torch.tensor(-1), torch.tensor(-1), torch.tensor(-1), torch.tensor(-1), 'Error'\n",
    "\n",
    "    def show_rgb(self, index):\n",
    "        self.__getitem__(index)\n",
    "        print(\"RGB image size:\", self.rgb_img.shape)        \n",
    "        cv2.imshow('RGB image',self.rgb_img)\n",
    "        cv2.waitKey(0)\n",
    "        cv2.destroyAllWindows()\n",
    "\n",
    "    def show_lab_encoder(self, index):\n",
    "        self.__getitem__(index)\n",
    "        print(\"Encoder Lab image size:\", self.lab_encoder_img.shape)\n",
    "        cv2.imshow('Encoder Lab image',self.lab_encoder_img)\n",
    "        cv2.waitKey(0)\n",
    "        cv2.destroyAllWindows()\n",
    "\n",
    "    def show_lab_inception(self, index):\n",
    "        self.__getitem__(index)\n",
    "        print(\"Inception Lab image size:\", self.lab_inception_img.shape)\n",
    "        cv2.imshow('Inception Lab image',self.lab_inception_img)\n",
    "        cv2.waitKey(0)\n",
    "        cv2.destroyAllWindows()\n",
    "    \n",
    "    def show_other_images(self, index):\n",
    "        a,b,c,d,_ = self.__getitem__(index)\n",
    "        print(\"Encoder l channel image size:\",a.shape)\n",
    "        cv2.imshow('Encoder l channel image',a.detach().numpy().transpose(1,2,0))\n",
    "        cv2.waitKey(0)\n",
    "        cv2.destroyAllWindows()\n",
    "        print(\"Encoder ab channel image size:\",b.shape)\n",
    "        cv2.imshow('Encoder a channel image',b.detach().numpy().transpose(1,2,0)[:,:,0])\n",
    "        cv2.waitKey(0)\n",
    "        cv2.destroyAllWindows()\n",
    "        cv2.imshow('Encoder b channel image',b.detach().numpy().transpose(1,2,0)[:,:,1])\n",
    "        cv2.waitKey(0)\n",
    "        cv2.destroyAllWindows()\n",
    "        print(\"Inception l channel image size:\",c.shape)\n",
    "        cv2.imshow('Inception l channel image',c.detach().numpy().transpose(1,2,0))\n",
    "        cv2.waitKey(0)\n",
    "        cv2.destroyAllWindows()\n",
    "        print(\"Color resized image size:\",d.shape)\n",
    "        cv2.imshow('Color resized image',d.detach().numpy().transpose(1,2,0))\n",
    "        cv2.waitKey(0)\n",
    "        cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_dataset = CustomDataset('data/train','train')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_dataset.show_rgb(0)\n",
    "# train_dataset.show_lab_encoder(0)\n",
    "# train_dataset.show_lab_inception(0)\n",
    "# train_dataset.show_other_images(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Encoder,self).__init__()\n",
    "        self.model = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=2, padding=1), \n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=2, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=2, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.Conv2d(in_channels=512, out_channels=256, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.BatchNorm2d(256),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        self.model = self.model.float()\n",
    "        return self.model(x.float())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fusion Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FusionLayer(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(FusionLayer,self).__init__()\n",
    "\n",
    "    def forward(self, inputs, mask=None):\n",
    "        ip, emb = inputs\n",
    "        emb = torch.stack([torch.stack([emb],dim=2)],dim=3)\n",
    "        emb = emb.repeat(1,1,ip.shape[2],ip.shape[3])\n",
    "        fusion = torch.cat((ip,emb),1)\n",
    "        return fusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, input_depth):\n",
    "        super(Decoder,self).__init__()\n",
    "        self.model = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=input_depth, out_channels=128, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.Upsample(scale_factor=2.0),\n",
    "            nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.Upsample(scale_factor=2.0),\n",
    "            nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.Conv2d(in_channels=32, out_channels=2, kernel_size=3, stride=1, padding=1),\n",
    "            nn.Tanh(),\n",
    "            nn.Upsample(scale_factor=2.0),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Network Definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Colorization(nn.Module):\n",
    "    def __init__(self, depth_after_fusion):\n",
    "        super(Colorization,self).__init__()\n",
    "        self.encoder = Encoder()\n",
    "        self.fusion = FusionLayer()\n",
    "        self.after_fusion = nn.Conv2d(in_channels=1256, out_channels=depth_after_fusion,kernel_size=1, stride=1,padding=0)\n",
    "        self.bnorm = nn.BatchNorm2d(256)\n",
    "        self.decoder = Decoder(depth_after_fusion)\n",
    "\n",
    "    def forward(self, img_l, img_emb):\n",
    "        img_enc = self.encoder(img_l)\n",
    "        fusion = self.fusion([img_enc, img_emb])\n",
    "        fusion = self.after_fusion(fusion)\n",
    "        fusion = self.bnorm(fusion)\n",
    "        return self.decoder(fusion)\n",
    "\n",
    "def init_weights(m):\n",
    "    if type(m) == nn.Conv2d or type(m) == nn.Linear:\n",
    "        torch.nn.init.xavier_normal_(m.weight.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Architecture Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if config.load_model_to_train or config.load_model_to_test:\n",
    "    checkpoint = torch.load(config.model_file_name,map_location=torch.device(config.device))\n",
    "    model = checkpoint['model']\n",
    "    model.load_state_dict(checkpoint['model_state_dict'])\n",
    "    model = model.to(config.device) \n",
    "    optimizer = checkpoint['optimizer']\n",
    "    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "    print('Loaded pretrain model | Previous train loss:',checkpoint['train_loss'], '| Previous validation loss:',checkpoint['val_loss'])\n",
    "else:\n",
    "    model = Colorization(256).to(config.device) \n",
    "#     model.apply(init_weights)\n",
    "    optimizer = torch.optim.Adam(model.parameters(),lr=hparams.learning_rate, weight_decay=1e-6)\n",
    "\n",
    "inception_model = models.inception_v3(pretrained=True).float().to(config.device)\n",
    "inception_model = inception_model.float()\n",
    "inception_model.eval()\n",
    "loss_criterion = torch.nn.MSELoss(reduction='mean').to(config.device)\n",
    "milestone_list  = list(range(0,hparams.epochs,2))\n",
    "scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=milestone_list, gamma=hparams.learning_rate_decay)\n",
    "writer = SummaryWriter()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Loaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not config.load_model_to_test:\n",
    "    train_dataset = CustomDataset('data/train','train')\n",
    "    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=hparams.batch_size, shuffle=True, num_workers=hparams.num_workers)\n",
    "    \n",
    "\n",
    "    validataion_dataset = CustomDataset('data/validation','validation')\n",
    "    validation_dataloader = torch.utils.data.DataLoader(validataion_dataset, batch_size=hparams.batch_size, shuffle=False, num_workers=hparams.num_workers)\n",
    "    \n",
    "    print('Train:',len(train_dataloader), '| Total Images:',len(train_dataloader)*hparams.batch_size)\n",
    "    print('Valid:',len(validation_dataloader), '| Total Images:',len(validation_dataloader)*hparams.batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training & Validation Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if not config.load_model_to_test:\n",
    "    for epoch in range(hparams.epochs):\n",
    "        print('Starting epoch:',epoch+1)\n",
    "\n",
    "        #*** Training step ***\n",
    "        loop_start = time.time()\n",
    "        avg_loss = 0.0\n",
    "        batch_loss = 0.0\n",
    "        main_start = time.time()\n",
    "        model.train()\n",
    "\n",
    "        for idx,(img_l_encoder, img_ab_encoder, img_l_inception, img_rgb, file_name) in enumerate(train_dataloader):\n",
    "            #*** Skip bad data ***\n",
    "            if not img_l_encoder.ndim:\n",
    "                continue\n",
    "\n",
    "            #*** Move data to GPU if available ***\n",
    "            img_l_encoder = img_l_encoder.to(config.device)\n",
    "            img_ab_encoder = img_ab_encoder.to(config.device)\n",
    "            img_l_inception = img_l_inception.to(config.device)\n",
    "\n",
    "            #*** Initialize Optimizer ***\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            #*** Forward Propagation ***\n",
    "            img_embs = inception_model(img_l_inception.float())\n",
    "            output_ab = model(img_l_encoder,img_embs)\n",
    "\n",
    "            #*** Back propogation ***\n",
    "            loss = loss_criterion(output_ab, img_ab_encoder.float())\n",
    "            loss.backward()\n",
    "\n",
    "            #*** Weight Update ****\n",
    "            optimizer.step()\n",
    "\n",
    "            #*** Reduce Learning Rate ***\n",
    "            scheduler.step()\n",
    "\n",
    "            #*** Loss Calculation ***\n",
    "            avg_loss += loss.item()\n",
    "            batch_loss += loss.item()\n",
    "\n",
    "            #*** Print stats after every point_batches ***\n",
    "            if idx%config.point_batches==0: \n",
    "                loop_end = time.time()   \n",
    "                print('Batch:',idx, '| Processing time for',config.point_batches,':',str(loop_end-loop_start)+'s',' | Batch Loss:', batch_loss/config.point_batches)\n",
    "                loop_start = time.time()\n",
    "                batch_loss = 0.0\n",
    "\n",
    "        #*** Print Training Data Stats ***\n",
    "        train_loss = avg_loss/len(train_dataloader)*hparams.batch_size\n",
    "        writer.add_scalar('Loss/train', train_loss, epoch)\n",
    "        print('Training Loss:',train_loss,'| Processed in ',str(time.time()-main_start)+'s')\n",
    "\n",
    "        #*** Validation Step ***       \n",
    "        avg_loss = 0.0\n",
    "        loop_start = time.time()\n",
    "        #*** Intialize Model to Eval Mode for validation ***\n",
    "        model.eval()\n",
    "        for idx,(img_l_encoder, img_ab_encoder, img_l_inception, img_rgb, file_name) in enumerate(validation_dataloader):\n",
    "            #*** Skip bad data ***\n",
    "            if not img_l_encoder.ndim:\n",
    "                continue\n",
    "\n",
    "            #*** Move data to GPU if available ***\n",
    "            img_l_encoder = img_l_encoder.to(config.device)\n",
    "            img_ab_encoder = img_ab_encoder.to(config.device)\n",
    "            img_l_inception = img_l_inception.to(config.device)\n",
    "\n",
    "            #*** Forward Propagation ***\n",
    "            img_embs = inception_model(img_l_inception.float())\n",
    "            output_ab = model(img_l_encoder,img_embs)\n",
    "\n",
    "            #*** Loss Calculation ***\n",
    "            loss = loss_criterion(output_ab, img_ab_encoder.float())\n",
    "            avg_loss += loss.item()\n",
    "\n",
    "        val_loss = avg_loss/len(validation_dataloader)*hparams.batch_size\n",
    "        writer.add_scalar('Loss/validation', val_loss, epoch)\n",
    "        print('Validation Loss:', val_loss,'| Processed in ',str(time.time()-loop_start)+'s')\n",
    "\n",
    "        #*** Save the Model to disk ***\n",
    "        checkpoint = {'model': model,'model_state_dict': model.state_dict(),\\\n",
    "                      'optimizer' : optimizer,'optimizer_state_dict' : optimizer.state_dict(), \\\n",
    "                      'train_loss':train_loss, 'val_loss':val_loss}\n",
    "        torch.save(checkpoint, config.model_file_name+'.'+str(epoch+1))\n",
    "        print(\"Model saved at:\",os.getcwd()+'/'+str(config.model_file_name)+'.'+str(epoch+1))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset = CustomDataset('data/test','test')\n",
    "test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=hparams.num_workers)\n",
    "print('Test: ',len(test_dataloader), '| Total Image:',len(test_dataloader))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Convert Tensor Image -> Numpy Image -> Color  Image -> Tensor Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def concatente_and_colorize(im_lab, img_ab):\n",
    "    # Assumption is that im_lab is of size [1,3,224,224]\n",
    "    #print(im_lab.size(),img_ab.size())\n",
    "    np_img = im_lab[0].cpu().detach().numpy().transpose(1,2,0)\n",
    "    lab = np.empty([*np_img.shape[0:2], 3],dtype=np.float32)\n",
    "    lab[:, :, 0] = np.squeeze(((np_img + 1) * 50))\n",
    "    lab[:, :, 1:] = img_ab[0].cpu().detach().numpy().transpose(1,2,0) * 127\n",
    "    np_img = cv2.cvtColor(lab,cv2.COLOR_Lab2RGB) \n",
    "    color_im = torch.stack([torchvision.transforms.ToTensor()(np_img)],dim=0)\n",
    "    return color_im"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#*** Inference Step ***\n",
    "avg_loss = 0.0\n",
    "loop_start = time.time()\n",
    "batch_start = time.time()\n",
    "batch_loss = 0.0\n",
    "\n",
    "for idx,(img_l_encoder, img_ab_encoder, img_l_inception, img_rgb, file_name) in enumerate(test_dataloader):\n",
    "        #*** Skip bad data ***\n",
    "        if not img_l_encoder.ndim:\n",
    "            continue\n",
    "            \n",
    "        #*** Move data to GPU if available ***\n",
    "        img_l_encoder = img_l_encoder.to(config.device)\n",
    "        img_ab_encoder = img_ab_encoder.to(config.device)\n",
    "        img_l_inception = img_l_inception.to(config.device)\n",
    "        \n",
    "        #*** Intialize Model to Eval Mode ***\n",
    "        model.eval()\n",
    "        \n",
    "        #*** Forward Propagation ***\n",
    "        img_embs = inception_model(img_l_inception.float())\n",
    "        output_ab = model(img_l_encoder,img_embs)\n",
    "        \n",
    "        np_out = output_ab[0].detach().numpy().transpose(1,2,0)\n",
    "        print(np_out.shape)\n",
    "        n_bins = 100\n",
    "        plt.hist(np_out[:,:,0],bins = n_bins,label='a channel')\n",
    "        plt.hist(np_out[:,:,1],bins = n_bins,label='b channel')\n",
    "        plt.legend(loc='upper right')\n",
    "        plt.show()\n",
    "        \n",
    "        #*** Adding l channel to ab channels ***\n",
    "        color_img = concatente_and_colorize(torch.stack([img_l_encoder[:,0,:,:]],dim=1),output_ab)\n",
    "        color_img_jpg = color_img[0].detach().numpy().transpose(1,2,0)\n",
    "        # cv2.imshow(color_img_jpg)\n",
    "        # cv2.waitKey(0)\n",
    "        # cv2.destroyAllWindows()\n",
    "        # cv2.imwrite('outputs/'+file_name[0],color_img_jpg*255)\n",
    "        # save_image(color_img[0],'outputs/'+file_name[0])\n",
    "\n",
    "        #*** Printing to Tensor Board ***\n",
    "        grid = torchvision.utils.make_grid(color_img)\n",
    "        writer.add_image('Output Lab Images', grid, 0)\n",
    "        \n",
    "        #*** Loss Calculation ***\n",
    "        loss = loss_criterion(output_ab, img_ab_encoder.float())\n",
    "        avg_loss += loss.item()\n",
    "        batch_loss += loss.item()\n",
    "\n",
    "        if idx%config.point_batches==0: \n",
    "            batch_end = time.time()   \n",
    "            print('Batch:',idx, '| Processing time for',config.point_batches,':',str(batch_end-batch_start)+'s', '| Batch Loss:', batch_loss/config.point_batches)\n",
    "            batch_start = time.time()\n",
    "            batch_loss = 0.0\n",
    "        \n",
    "test_loss = avg_loss/len(test_dataloader)\n",
    "print('Test Loss:',avg_loss/len(test_dataloader),'| Processed in ',str(time.time()-loop_start)+'s')\n",
    "writer.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}