DeepHomographyEstimation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai.vision import *\n",
    "from glob import glob\n",
    "from matplotlib import pyplot as plt\n",
    "import cv2\n",
    "import random\n",
    "import numpy as np\n",
    "from numpy.linalg import inv\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torchsummary import summary\n",
    "import torch.optim as optim\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Model,self).__init__()\n",
    "        self.layer1 = nn.Sequential(nn.Conv2d(2,64,3,padding=1),\n",
    "                                    nn.BatchNorm2d(64),\n",
    "                                    nn.ReLU())\n",
    "                                    \n",
    "        self.layer2 = nn.Sequential(nn.Conv2d(64,64,3,padding=1),\n",
    "                                    nn.BatchNorm2d(64),\n",
    "                                    nn.ReLU(),\n",
    "                                    nn.MaxPool2d(2))\n",
    "        self.layer3 = nn.Sequential(nn.Conv2d(64,64,3,padding=1),\n",
    "                                    nn.BatchNorm2d(64),\n",
    "                                    nn.ReLU())\n",
    "        self.layer4 = nn.Sequential(nn.Conv2d(64,64,3,padding=1),\n",
    "                                    nn.BatchNorm2d(64),\n",
    "                                    nn.ReLU(),\n",
    "                                    nn.MaxPool2d(2))\n",
    "        self.layer5 = nn.Sequential(nn.Conv2d(64,128,3,padding=1),\n",
    "                                    nn.BatchNorm2d(128),\n",
    "                                    nn.ReLU())        \n",
    "        self.layer6 = nn.Sequential(nn.Conv2d(128,128,3,padding=1),\n",
    "                                    nn.BatchNorm2d(128),\n",
    "                                    nn.ReLU(),\n",
    "                                    nn.MaxPool2d(2))\n",
    "        self.layer7 = nn.Sequential(nn.Conv2d(128,128,3,padding=1),\n",
    "                                    nn.BatchNorm2d(128),\n",
    "                                    nn.ReLU())\n",
    "        self.layer8 = nn.Sequential(nn.Conv2d(128,128,3,padding=1),\n",
    "                                    nn.BatchNorm2d(128),\n",
    "                                    nn.ReLU())\n",
    "        self.fc1 = nn.Linear(128*16*16,1024)\n",
    "        self.fc2 = nn.Linear(1024,8)\n",
    "        \n",
    "    def forward(self,x):\n",
    "        out = self.layer1(x)\n",
    "        out = self.layer2(out)\n",
    "        out = self.layer3(out)\n",
    "        out = self.layer4(out)\n",
    "        out = self.layer5(out)\n",
    "        out = self.layer6(out)\n",
    "        out = self.layer7(out)\n",
    "        out = self.layer8(out)\n",
    "        out = out.view(-1,128* 16* 16)\n",
    "        out = self.fc1(out)\n",
    "        out = self.fc2(out)\n",
    "        return out\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating a customized dataset class in pytorch\n",
    " \n",
    "class CocoDdataset(Dataset):\n",
    "    def __init__(self,path):\n",
    "        X=()\n",
    "        Y=()\n",
    "        lst = os.listdir(path)\n",
    "        it=0\n",
    "        for i in lst:\n",
    "            array = np.load(path+'%s'%i)\n",
    "            x = torch.from_numpy((array[0].astype(float)-127.5)/127.5)\n",
    "            X = X+(x,)\n",
    "            y = torch.from_numpy(array[1].astype(float) / 32.)\n",
    "            Y = Y+(y,)\n",
    "            it+=1\n",
    "        self.len = it\n",
    "        self.X_data = X\n",
    "        self.Y_data = Y\n",
    "    def __getitem__(self,index):\n",
    "        return self.X_data[index], self.Y_data[index] \n",
    "    def __len__(self):\n",
    "        return self.len\n",
    "    \n",
    "\n",
    "\n",
    "\n",
    "train_path = './train2017processed/'\n",
    "validation_path = './val2017processed/'\n",
    "test_path = './test2017processed/'\n",
    "\n",
    "TrainingData = CocoDdataset(train_path)\n",
    "ValidationData = CocoDdataset(validation_path)\n",
    "TestData = CocoDdataset(test_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 64, 128, 128]           1,216\n",
      "       BatchNorm2d-2         [-1, 64, 128, 128]             128\n",
      "              ReLU-3         [-1, 64, 128, 128]               0\n",
      "            Conv2d-4         [-1, 64, 128, 128]          36,928\n",
      "       BatchNorm2d-5         [-1, 64, 128, 128]             128\n",
      "              ReLU-6         [-1, 64, 128, 128]               0\n",
      "         MaxPool2d-7           [-1, 64, 64, 64]               0\n",
      "            Conv2d-8           [-1, 64, 64, 64]          36,928\n",
      "       BatchNorm2d-9           [-1, 64, 64, 64]             128\n",
      "             ReLU-10           [-1, 64, 64, 64]               0\n",
      "           Conv2d-11           [-1, 64, 64, 64]          36,928\n",
      "      BatchNorm2d-12           [-1, 64, 64, 64]             128\n",
      "             ReLU-13           [-1, 64, 64, 64]               0\n",
      "        MaxPool2d-14           [-1, 64, 32, 32]               0\n",
      "           Conv2d-15          [-1, 128, 32, 32]          73,856\n",
      "      BatchNorm2d-16          [-1, 128, 32, 32]             256\n",
      "             ReLU-17          [-1, 128, 32, 32]               0\n",
      "           Conv2d-18          [-1, 128, 32, 32]         147,584\n",
      "      BatchNorm2d-19          [-1, 128, 32, 32]             256\n",
      "             ReLU-20          [-1, 128, 32, 32]               0\n",
      "        MaxPool2d-21          [-1, 128, 16, 16]               0\n",
      "           Conv2d-22          [-1, 128, 16, 16]         147,584\n",
      "      BatchNorm2d-23          [-1, 128, 16, 16]             256\n",
      "             ReLU-24          [-1, 128, 16, 16]               0\n",
      "           Conv2d-25          [-1, 128, 16, 16]         147,584\n",
      "      BatchNorm2d-26          [-1, 128, 16, 16]             256\n",
      "             ReLU-27          [-1, 128, 16, 16]               0\n",
      "           Linear-28                 [-1, 1024]      33,555,456\n",
      "           Linear-29                    [-1, 8]           8,200\n",
      "================================================================\n",
      "Total params: 34,193,800\n",
      "Trainable params: 34,193,800\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.12\n",
      "Forward/backward pass size (MB): 70.26\n",
      "Params size (MB): 130.44\n",
      "Estimated Total Size (MB): 200.82\n",
      "----------------------------------------------------------------\n",
      "Train Epoch: [1/48] [1848/1849 (100%)]\\Mean Squared Error: 0.116034\n",
      "Train Epoch: [2/48] [1848/1849 (100%)]\\Mean Squared Error: 0.087828\n",
      "Train Epoch: [3/48] [1848/1849 (100%)]\\Mean Squared Error: 0.080873\n",
      "Train Epoch: [4/48] [1848/1849 (100%)]\\Mean Squared Error: 0.068192\n",
      "Train Epoch: [5/48] [1848/1849 (100%)]\\Mean Squared Error: 0.063107\n",
      "Train Epoch: [6/48] [1848/1849 (100%)]\\Mean Squared Error: 0.056298\n",
      "Train Epoch: [7/48] [1848/1849 (100%)]\\Mean Squared Error: 0.050319\n",
      "Train Epoch: [8/48] [1848/1849 (100%)]\\Mean Squared Error: 0.046902\n",
      "Train Epoch: [9/48] [1848/1849 (100%)]\\Mean Squared Error: 0.040800\n",
      "Train Epoch: [10/48] [1848/1849 (100%)]\\Mean Squared Error: 0.037759\n",
      "Train Epoch: [11/48] [1848/1849 (100%)]\\Mean Squared Error: 0.035209\n",
      "Train Epoch: [12/48] [1848/1849 (100%)]\\Mean Squared Error: 0.032453\n",
      "Train Epoch: [13/48] [1848/1849 (100%)]\\Mean Squared Error: 0.030337\n",
      "Train Epoch: [14/48] [1848/1849 (100%)]\\Mean Squared Error: 0.026058\n",
      "Train Epoch: [15/48] [1848/1849 (100%)]\\Mean Squared Error: 0.024556\n",
      "Train Epoch: [16/48] [1848/1849 (100%)]\\Mean Squared Error: 0.022692\n",
      "Train Epoch: [17/48] [1848/1849 (100%)]\\Mean Squared Error: 0.020574\n",
      "Train Epoch: [18/48] [1848/1849 (100%)]\\Mean Squared Error: 0.017353\n",
      "Train Epoch: [19/48] [1848/1849 (100%)]\\Mean Squared Error: 0.016512\n",
      "Train Epoch: [20/48] [1848/1849 (100%)]\\Mean Squared Error: 0.015141\n",
      "Train Epoch: [21/48] [1848/1849 (100%)]\\Mean Squared Error: 0.014694\n",
      "Train Epoch: [22/48] [1848/1849 (100%)]\\Mean Squared Error: 0.013501\n",
      "Train Epoch: [23/48] [1848/1849 (100%)]\\Mean Squared Error: 0.012297\n",
      "Train Epoch: [24/48] [1848/1849 (100%)]\\Mean Squared Error: 0.011153\n",
      "Train Epoch: [25/48] [1848/1849 (100%)]\\Mean Squared Error: 0.010395\n",
      "Train Epoch: [26/48] [1848/1849 (100%)]\\Mean Squared Error: 0.009637\n",
      "Train Epoch: [27/48] [1848/1849 (100%)]\\Mean Squared Error: 0.009614\n",
      "Train Epoch: [28/48] [1848/1849 (100%)]\\Mean Squared Error: 0.008482\n",
      "Train Epoch: [29/48] [1848/1849 (100%)]\\Mean Squared Error: 0.007813\n",
      "Train Epoch: [30/48] [1848/1849 (100%)]\\Mean Squared Error: 0.007763\n",
      "Train Epoch: [31/48] [1848/1849 (100%)]\\Mean Squared Error: 0.006555\n",
      "Train Epoch: [32/48] [1848/1849 (100%)]\\Mean Squared Error: 0.006307\n",
      "Train Epoch: [33/48] [1848/1849 (100%)]\\Mean Squared Error: 0.005889\n",
      "Train Epoch: [34/48] [1848/1849 (100%)]\\Mean Squared Error: 0.005683\n",
      "Train Epoch: [35/48] [1848/1849 (100%)]\\Mean Squared Error: 0.005530\n",
      "Train Epoch: [36/48] [1848/1849 (100%)]\\Mean Squared Error: 0.005131\n",
      "Train Epoch: [37/48] [1848/1849 (100%)]\\Mean Squared Error: 0.005221\n",
      "Train Epoch: [38/48] [1848/1849 (100%)]\\Mean Squared Error: 0.004674\n",
      "Train Epoch: [39/48] [1848/1849 (100%)]\\Mean Squared Error: 0.004853\n",
      "Train Epoch: [40/48] [1848/1849 (100%)]\\Mean Squared Error: 0.004739\n",
      "Train Epoch: [41/48] [1848/1849 (100%)]\\Mean Squared Error: 0.005890\n",
      "Train Epoch: [42/48] [1848/1849 (100%)]\\Mean Squared Error: 0.006754\n",
      "Train Epoch: [43/48] [1848/1849 (100%)]\\Mean Squared Error: 0.008164\n",
      "Train Epoch: [44/48] [1848/1849 (100%)]\\Mean Squared Error: 0.011123\n",
      "Train Epoch: [45/48] [1848/1849 (100%)]\\Mean Squared Error: 0.010395\n",
      "Train Epoch: [46/48] [1848/1849 (100%)]\\Mean Squared Error: 0.007917\n",
      "Train Epoch: [47/48] [1848/1849 (100%)]\\Mean Squared Error: 0.006665\n",
      "Train Epoch: [48/48] [1848/1849 (100%)]\\Mean Squared Error: 0.005879\n"
     ]
    }
   ],
   "source": [
    "batch_size = 64\n",
    "TrainLoader = DataLoader(TrainingData,batch_size)\n",
    "ValidationLoader = DataLoader(ValidationData,batch_size)\n",
    "TestLoader = DataLoader(TestData,batch_size)\n",
    "criterion = nn.MSELoss()\n",
    "num_samples = 118287\n",
    "total_iteration = 90000\n",
    "steps_per_epoch = num_samples / batch_size\n",
    "epochs = int(total_iteration / steps_per_epoch)\n",
    "model = Model().to(device)\n",
    "summary(model,(2,128,128))\n",
    "optimizer = optim.SGD(model.parameters(),lr=0.005, momentum=0.9)\n",
    "for epoch in range(epochs):\n",
    "    \n",
    "    for i, (images, target) in enumerate(TrainLoader):\n",
    "        optimizer.zero_grad()\n",
    "        images = images.to(device); target = target.to(device)\n",
    "        images = images.permute(0,3,1,2).float(); target = target.float()\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, target.view(-1,8))\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        if (i+1) % len(TrainLoader) == 0:\n",
    "            print('Train Epoch: [{}/{}] [{}/{} ({:.0f}%)]\\Mean Squared Error: {:.6f}'.format(\n",
    "                epoch+1,epochs, i , len(TrainLoader),\n",
    "                100. * i / len(TrainLoader), loss))\n",
    "\n",
    "state = {'epoch': epochs, 'state_dict': model.state_dict(),\n",
    "             'optimizer': optimizer.state_dict() }\n",
    "torch.save(state, 'DeepHomographyEstimation.pth')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    for i,(images, target) in enumerate(ValidationLoader):\n",
    "        images = images.to(device)\n",
    "        target = target.to(device)\n",
    "        images = images.permute(0,3,1,2).float()\n",
    "        target = target.float()\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, target.view(-1,8))\n",
    "    print('\\Mean Squared Error: {:.6f}'.format(np.mean(loss)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}