DeepMatting.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import numpy as np\n",
    "import argparse\n",
    "import os\n",
    "from random import  shuffle\n",
    "from tqdm import *\n",
    "\n",
    "##########\n",
    "# TORCH\n",
    "import torch\n",
    "import torchvision.models as models\n",
    "from torch.autograd import Variable\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "#########\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "from scipy import misc,ndimage\n",
    "import copy\n",
    "import itertools\n",
    "import os\n",
    "from sys import getrefcount\n",
    "import gc\n",
    "trimap_kernel = [val for val in range(20,40)]\n",
    "g_mean = np.array(([123.11,111.78,100.59])).reshape([1,1,3])\n",
    "\n",
    "def UR_center(trimap):\n",
    "\n",
    "    target = np.where(trimap==128)\n",
    "    index = random.choice([i for i in range(len(target[0]))])\n",
    "    return  np.array(target)[:,index][:2]\n",
    "\n",
    "def load_path(alpha,eps,BG,hard_mode = False):\n",
    "    #print(alpha, eps, BG)\n",
    "    images_alpha = sorted(os.listdir(alpha))\n",
    "    #images_alpha = list(np.repeat(images_alpha,100))\n",
    "    images_merged = sorted(os.listdir(eps))\n",
    "    images_bg = sorted(os.listdir(BG))\n",
    "    #for i in images_merged:\n",
    "    f_ix = [images_alpha.index('_'.join(i.split('_')[:-1])+'.jpg') for i in images_merged]\n",
    "    m_ix = [int(i.split('_')[-1].split('.')[0]) for i in images_merged]\n",
    "    #print(f_ix[:100], m_ix[:100])\n",
    "    b_ix = list(100*np.array(f_ix) + np.array(m_ix))\n",
    "    #BGs_abspath = [os.path.join(BG,common_path)[:-3] + 'jpg' for common_path in images_bg]\n",
    "    #print(b_ix[:100])\n",
    "    #b_ix = [i+1 for i in b_ix]\n",
    "    #folders = os.listdir(alpha)\n",
    "    #common_paths = []\n",
    "    #if hard_mode:\n",
    "    #    for folder in folders:\n",
    "    #        if int(folder) in hard_samples: \n",
    "    #            images = os.listdir(os.path.join(alpha,folder))\n",
    "    #            common_paths.extend([os.path.join(folder,image) for image in images])\n",
    "    #else:\n",
    "    #    for folder in folders:\n",
    "    #        #if int(folder)==137:\n",
    "    #        images = os.listdir(os.path.join(alpha,folder))\n",
    "    #        common_paths.extend([os.path.join(folder,image) for image in images])\n",
    "    #print(len(images_alpha))\n",
    "    images_alpha_1 = list(np.array(images_alpha)[f_ix])\n",
    "    alphas_abspath = [os.path.join(alpha,common_path) for common_path in images_alpha_1]\n",
    "    #print(len(images_bg), b_ix)\n",
    "    images_bg_1 = list(np.array(images_bg)[b_ix])\n",
    "    BGs_abspath = [os.path.join(BG,common_path)[:-3] + 'jpg' for common_path in images_bg_1]\n",
    "    epses_abspath = images_merged\n",
    "    epses_abspath = [os.path.join(eps,common_path) for common_path in images_merged]\n",
    "    #print(alphas_abspath[:10], BGs_abspath[:10])\n",
    "    return np.array(alphas_abspath),np.array(epses_abspath),np.array(BGs_abspath)\n",
    "\n",
    "def load_data(batch_alpha_paths,batch_eps_paths,batch_BG_paths):\n",
    "    batch_size = batch_alpha_paths.shape[0]\n",
    "    train_batch = []\n",
    "    images_without_mean_reduction = []\n",
    "    batch = 0\n",
    "    while batch < batch_size:\n",
    "        i = batch\n",
    "        #for i in range(batch_size):\t\n",
    "        alpha = misc.imread(batch_alpha_paths[i],'L').astype(np.float32)\n",
    "        #print(alpha.shape)\n",
    "        eps = misc.imread(batch_eps_paths[i]).astype(np.float32)\n",
    "        #print(eps.shape)\n",
    "        BG = misc.imread(batch_BG_paths[i]).astype(np.float32)\n",
    "        bbox = alpha.shape\n",
    "        #print(bbox)\n",
    "        BG = misc.imresize(BG, bbox) \n",
    "        batch_i,raw_RGB = preprocessing_single(alpha, BG, eps,batch_alpha_paths[i])\t\n",
    "        train_batch.append(batch_i)\n",
    "        images_without_mean_reduction.append(raw_RGB)\n",
    "        batch += 1\n",
    "    train_batch = np.stack(train_batch).astype(np.float64)\n",
    "    images_without_mean_reduction = np.asarray(images_without_mean_reduction).astype(np.float64)\n",
    "    return train_batch[:,:,:,:3],np.expand_dims(train_batch[:,:,:,3],3),np.expand_dims(train_batch[:,:,:,4],3),train_batch[:,:,:,5:8],train_batch[:,:,:,8:],images_without_mean_reduction\n",
    "\n",
    "def generate_trimap(trimap,alpha):\n",
    "\n",
    "    k_size = random.choice(trimap_kernel)\n",
    "    trimap[np.where((ndimage.grey_dilation(alpha[:,:,0],size=(k_size,k_size)) - ndimage.grey_erosion(alpha[:,:,0],size=(k_size,k_size)))!=0)] = 128\n",
    "    #trimap[np.where((ndimage.grey_dilation(alpha[:,:,0],size=(k_size,k_size)) - alpha[:,:,0]!=0))] = 128\n",
    "    return trimap\n",
    "\n",
    "def preprocessing_single(alpha, BG, eps,name,image_size=320):\n",
    "\n",
    "    alpha = np.expand_dims(alpha,2)\n",
    "    trimap = np.copy(alpha)\n",
    "    trimap = generate_trimap(trimap,alpha)\n",
    "\n",
    "    train_data = np.zeros([image_size,image_size,8])\n",
    "    crop_size = random.choice([320,480,620])\n",
    "#    crop_size = 320   \n",
    "    flip = random.choice([0,1])\n",
    "    i_UR_center = UR_center(trimap)\n",
    "    #i_UR_center = [int(alpha.shape[0]/2),int(alpha.shape[1]/2)]\n",
    "    #print(trimap.shape,alpha.shape,BG.shape,eps.shape)\n",
    "    train_pre = np.concatenate([trimap,alpha,BG,eps],2)\n",
    "\n",
    "    if crop_size == 320:\n",
    "        h_start_index = i_UR_center[0] - 159\n",
    "        if h_start_index<0:\n",
    "            h_start_index = 0\n",
    "        w_start_index = i_UR_center[1] - 159\n",
    "        if w_start_index<0:\n",
    "            w_start_index = 0\n",
    "        tmp = train_pre[h_start_index:h_start_index+320, w_start_index:w_start_index+320, :]\n",
    "        if flip:\n",
    "            tmp = tmp[:,::-1,:]\n",
    "        tmp1 = np.zeros([image_size,image_size,8]).astype(np.float32)\n",
    "        tmp1[:,:,0] = misc.imresize(tmp[:,:,0].astype(np.uint8),[image_size,image_size],interp = 'nearest',mode='L').astype(np.float32)\n",
    "        tmp1[:,:,1] = misc.imresize(tmp[:,:,1].astype(np.uint8),[image_size,image_size]).astype(np.float32) / 255.0\n",
    "        tmp1[:,:,2:5] = misc.imresize(tmp[:,:,2:5].astype(np.uint8),[image_size,image_size,3]).astype(np.float32)\n",
    "        tmp1[:,:,5:] = misc.imresize(tmp[:,:,5:].astype(np.uint8),[image_size,image_size,3]).astype(np.float32)\n",
    "        tmp1[:,:,5:] = np.expand_dims(tmp1[:,:,1],2)  * tmp1[:,:,5:]  # here replace eps with FG        \n",
    "        #tmp[:,:,1] = tmp[:,:,1] / 255.0\n",
    "        #tmp[:,:,5:] = np.expand_dims(tmp[:,:,1],2)  * tmp[:,:,5:]  # here replace eps with FG\n",
    "        raw_RGB = np.expand_dims(tmp1[:,:,1],2)  * tmp1[:,:,5:] + np.expand_dims((1. - tmp1[:,:,1]),2) * tmp1[:,:,2:5]\n",
    "        reduced_RGB = raw_RGB - g_mean\n",
    "        tmp1 = np.concatenate([reduced_RGB,tmp1],2)\n",
    "        train_data = tmp1\n",
    "\n",
    "    if crop_size == 480:\n",
    "        h_start_index = i_UR_center[0] - 239\n",
    "        if h_start_index<0:\n",
    "            h_start_index = 0\n",
    "        w_start_index = i_UR_center[1] - 239\n",
    "        if w_start_index<0:\n",
    "            w_start_index = 0\n",
    "        tmp = train_pre[h_start_index:h_start_index+480, w_start_index:w_start_index+480, :]\n",
    "        if flip:\n",
    "            tmp = tmp[:,::-1,:]\n",
    "        tmp1 = np.zeros([image_size,image_size,8]).astype(np.float32)\n",
    "        tmp1[:,:,0] = misc.imresize(tmp[:,:,0].astype(np.uint8),[image_size,image_size],interp = 'nearest',mode='L').astype(np.float32)\n",
    "        tmp1[:,:,1] = misc.imresize(tmp[:,:,1].astype(np.uint8),[image_size,image_size]).astype(np.float32) / 255.0\n",
    "        tmp1[:,:,2:5] = misc.imresize(tmp[:,:,2:5].astype(np.uint8),[image_size,image_size,3]).astype(np.float32)\n",
    "        tmp1[:,:,5:] = misc.imresize(tmp[:,:,5:].astype(np.uint8),[image_size,image_size,3]).astype(np.float32)\n",
    "        tmp1[:,:,5:] = np.expand_dims(tmp1[:,:,1],2)  * tmp1[:,:,5:]  # here replace eps with FG        \n",
    "        raw_RGB = np.expand_dims(tmp1[:,:,1],2)  * tmp1[:,:,5:] + np.expand_dims((1. - tmp1[:,:,1]),2) * tmp1[:,:,2:5]\n",
    "        reduced_RGB = raw_RGB - g_mean      \n",
    "        tmp1 = np.concatenate([reduced_RGB,tmp1],2)\n",
    "        train_data = tmp1\n",
    "\n",
    "    if crop_size == 620:\n",
    "        h_start_index = i_UR_center[0] - 309\n",
    "        #boundary security\n",
    "        if h_start_index<0:\n",
    "            h_start_index = 0\n",
    "        w_start_index = i_UR_center[1] - 309\n",
    "        if w_start_index<0:\n",
    "            w_start_index = 0\n",
    "        tmp = train_pre[h_start_index:h_start_index+620, w_start_index:w_start_index+620, :]\n",
    "        if flip:\n",
    "            tmp = tmp[:,::-1,:]\n",
    "        tmp1 = np.zeros([image_size,image_size,8]).astype(np.float32)\n",
    "        tmp1[:,:,0] = misc.imresize(tmp[:,:,0].astype(np.uint8),[image_size,image_size],interp = 'nearest',mode='L').astype(np.float32)\n",
    "        tmp1[:,:,1] = misc.imresize(tmp[:,:,1].astype(np.uint8),[image_size,image_size]).astype(np.float32) / 255.0\n",
    "        tmp1[:,:,2:5] = misc.imresize(tmp[:,:,2:5].astype(np.uint8),[image_size,image_size,3]).astype(np.float32)\n",
    "        tmp1[:,:,5:] = misc.imresize(tmp[:,:,5:].astype(np.uint8),[image_size,image_size,3]).astype(np.float32)\n",
    "        tmp1[:,:,5:] = np.expand_dims(tmp1[:,:,1],2)  * tmp1[:,:,5:]  # here replace eps with FG        \n",
    "        raw_RGB = np.expand_dims(tmp1[:,:,1],2)  * tmp1[:,:,5:] + np.expand_dims((1. - tmp1[:,:,1]),2) * tmp1[:,:,2:5]\n",
    "        reduced_RGB = raw_RGB - g_mean      \n",
    "        tmp1 = np.concatenate([reduced_RGB,tmp1],2)\n",
    "        train_data = tmp1\n",
    "    train_data = train_data.astype(np.float32)\n",
    "#    misc.imsave('./train_alpha.png',train_data[:,:,4])\n",
    "    return train_data,raw_RGB\n",
    "\n",
    "def load_alphamatting_path(test_alpha):\n",
    "\trgb_path = os.path.join(test_alpha,'merged')\n",
    "\ttrimap_path = os.path.join(test_alpha,'trimaps/')\n",
    "\talpha_path = os.path.join(test_alpha,'mask')\t\n",
    "\timages = [os.path.join(rgb_path, i) for i in sorted(os.listdir(rgb_path))]\n",
    "        tri_images = [os.path.join(trimap_path, i) for i in sorted(os.listdir(trimap_path))]\n",
    "        alpha_images = [os.path.join(alpha_path, i) for i in list(np.repeat(np.array(sorted(os.listdir(alpha_path))),20))]\n",
    "        return images, tri_images, alpha_images, sorted(os.listdir(trimap_path)) \n",
    "\n",
    "def load_alphamatting_data(rgb_path,trimap_path, alpha_path ):\n",
    "        rgb = misc.imread(rgb_path)\n",
    "        trimap = misc.imread(trimap_path,'L')\n",
    "        alpha = misc.imread(alpha_path,'L')/255.0\n",
    "        all_shape = trimap.shape\n",
    "        rgb = misc.imresize(rgb,[320,320,3])-g_mean\n",
    "        trimap = misc.imresize(trimap,[320,320],interp = 'nearest').astype(np.float32)\n",
    "        trimap = np.expand_dims(trimap,2)\n",
    "        trimap_size = trimap.shape\n",
    "\treturn np.array(rgb), np.array(trimap), np.array(alpha), all_shape, trimap_size\n",
    "\n",
    "\n",
    "def load_validation_data(vali_root):\n",
    "\talpha_dir = os.path.join(vali_root,'alpha')\n",
    "\tRGB_dir = os.path.join(vali_root,'RGB')\n",
    "\timages = os.listdir(alpha_dir)\n",
    "\ttest_num = len(images)\n",
    "\t\n",
    "\tall_shape = []\n",
    "\trgb_batch = []\n",
    "\ttri_batch = []\n",
    "\talp_batch = []\n",
    "\n",
    "\tfor i in range(test_num):\n",
    "\t\trgb = misc.imread(os.path.join(RGB_dir,images[i]))\n",
    "\t\talpha = misc.imread(os.path.join(alpha_dir,images[i]),'L') \n",
    "\t\ttrimap = generate_trimap(np.expand_dims(np.copy(alpha),2),np.expand_dims(alpha,2))[:,:,0]\n",
    "\t\talpha = alpha / 255.0\n",
    "\t\tall_shape.append(trimap.shape)\n",
    "\t\trgb_batch.append(misc.imresize(rgb,[320,320,3])-g_mean)\n",
    "\t\ttrimap = misc.imresize(trimap,[320,320],interp = 'nearest').astype(np.float32)\n",
    "\t\ttri_batch.append(np.expand_dims(trimap,2))\n",
    "\t\talp_batch.append(alpha)\n",
    "\treturn np.array(rgb_batch),np.array(tri_batch),np.array(alp_batch),all_shape,images\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "range_size is 43100\n",
      "=> no checkpoint found at 'False'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:242: DeprecationWarning: `imread` is deprecated!\n",
      "`imread` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``imageio.imread`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:63: DeprecationWarning: `imread` is deprecated!\n",
      "`imread` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``imageio.imread`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:65: DeprecationWarning: `imread` is deprecated!\n",
      "`imread` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``imageio.imread`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:67: DeprecationWarning: `imread` is deprecated!\n",
      "`imread` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``imageio.imread`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:70: DeprecationWarning: `imresize` is deprecated!\n",
      "`imresize` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``skimage.transform.resize`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:112: DeprecationWarning: `imresize` is deprecated!\n",
      "`imresize` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``skimage.transform.resize`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:113: DeprecationWarning: `imresize` is deprecated!\n",
      "`imresize` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``skimage.transform.resize`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:114: DeprecationWarning: `imresize` is deprecated!\n",
      "`imresize` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``skimage.transform.resize`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:115: DeprecationWarning: `imresize` is deprecated!\n",
      "`imresize` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``skimage.transform.resize`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:135: DeprecationWarning: `imresize` is deprecated!\n",
      "`imresize` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``skimage.transform.resize`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:136: DeprecationWarning: `imresize` is deprecated!\n",
      "`imresize` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``skimage.transform.resize`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:137: DeprecationWarning: `imresize` is deprecated!\n",
      "`imresize` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``skimage.transform.resize`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:138: DeprecationWarning: `imresize` is deprecated!\n",
      "`imresize` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``skimage.transform.resize`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:157: DeprecationWarning: `imresize` is deprecated!\n",
      "`imresize` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``skimage.transform.resize`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:158: DeprecationWarning: `imresize` is deprecated!\n",
      "`imresize` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``skimage.transform.resize`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:159: DeprecationWarning: `imresize` is deprecated!\n",
      "`imresize` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``skimage.transform.resize`` instead.\n",
      "/home/manisha.padala/env/lib/python2.7/site-packages/ipykernel_launcher.py:160: DeprecationWarning: `imresize` is deprecated!\n",
      "`imresize` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``skimage.transform.resize`` instead.\n"
     ]
    }
   ],
   "source": [
    "image_size = 320\n",
    "train_batch_size = 1\n",
    "max_epochs = 1000000\n",
    "\n",
    "#checkpoint file path\n",
    "pretrained_model = False\n",
    "test_dir = '/ssd_scratch/cvit/vamsi.muthireddy/Test_set'\n",
    "test_outdir = '/ssd_scratch/cvit/vamsi.muthireddy/test_predict'\n",
    "log_dir = 'matting_log'\n",
    "\n",
    "dataset_alpha = '/ssd_scratch/cvit/manisha/Training_set/mask'\n",
    "dataset_merged = dataset_eps= '/ssd_scratch/cvit/manisha/Training_set/merged'\n",
    "dataset_BG = '/ssd_scratch/cvit/manisha/Training_set/bg'\n",
    "\n",
    "paths_alpha,paths_eps,paths_BG = load_path(dataset_alpha,dataset_eps,dataset_BG,hard_mode = False)\n",
    "\n",
    "range_size = len(paths_alpha)\n",
    "print('range_size is %d' % range_size)\n",
    "#range_size/batch_size has to be int\n",
    "batchs_per_epoch = int(range_size/train_batch_size) \n",
    "def weights_init(m):\n",
    "    if isinstance(m, nn.Conv2d):\n",
    "        #print(m.weight.data.shape, m.bias.data.shape)\n",
    "        nn.init.xavier_normal(m.weight.data)\n",
    "\n",
    "class DeepMatting(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(DeepMatting, self).__init__()\n",
    "        batchNorm_momentum = 0.1\n",
    "        self.conv1_1 = nn.Conv2d(4, 64, kernel_size=3,stride = 1, padding=1,bias=True)\n",
    "        self.bn11 = nn.BatchNorm2d(64, momentum= batchNorm_momentum)\n",
    "        self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3,stride = 1, padding=1,bias=True)\n",
    "        self.bn12 = nn.BatchNorm2d(64, momentum= batchNorm_momentum)\n",
    "        self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, padding=1,bias=True)\n",
    "        self.bn21 = nn.BatchNorm2d(128, momentum= batchNorm_momentum)\n",
    "        self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, padding=1,bias=True)\n",
    "        self.bn22 = nn.BatchNorm2d(128, momentum= batchNorm_momentum)\n",
    "        self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, padding=1,bias=True)\n",
    "        self.bn31 = nn.BatchNorm2d(256, momentum= batchNorm_momentum)\n",
    "        self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, padding=1,bias=True)\n",
    "        self.bn32 = nn.BatchNorm2d(256, momentum= batchNorm_momentum)\n",
    "        self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, padding=1,bias=True)\n",
    "        self.bn33 = nn.BatchNorm2d(256, momentum= batchNorm_momentum)\n",
    "        self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, padding=1,bias=True)\n",
    "        self.bn41 = nn.BatchNorm2d(512, momentum= batchNorm_momentum)\n",
    "        self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1,bias=True)\n",
    "        self.bn42 = nn.BatchNorm2d(512, momentum= batchNorm_momentum)\n",
    "        self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1,bias=True)\n",
    "        self.bn43 = nn.BatchNorm2d(512, momentum= batchNorm_momentum)\n",
    "        self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, padding=1,bias=True)\n",
    "        self.bn51 = nn.BatchNorm2d(512, momentum= batchNorm_momentum)\n",
    "        self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1,bias=True)\n",
    "        self.bn52 = nn.BatchNorm2d(512, momentum= batchNorm_momentum)\n",
    "        self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1,bias=True)\n",
    "        self.bn53 = nn.BatchNorm2d(512, momentum= batchNorm_momentum)\n",
    "        self.conv6_1 = nn.Conv2d(512, 4096, kernel_size=7, padding=3,bias=True)\n",
    "        self.bn61 = nn.BatchNorm2d(4096, momentum= batchNorm_momentum)\n",
    "        \n",
    "        self.deconv6_1 = nn.Conv2d(4096, 512, kernel_size=1,bias=True)\n",
    "        self.bn61d = nn.BatchNorm2d(512, momentum= batchNorm_momentum)\n",
    "        self.deconv5_1 = nn.Conv2d(512, 512, kernel_size=5, padding=2,bias=True)\n",
    "        self.bn51d = nn.BatchNorm2d(512, momentum= batchNorm_momentum)\n",
    "        self.deconv4_1 = nn.Conv2d(512, 256, kernel_size=5, padding=2,bias=True)\n",
    "        self.bn41d = nn.BatchNorm2d(256, momentum= batchNorm_momentum)\n",
    "        self.deconv3_1 = nn.Conv2d(256, 128, kernel_size=5, padding=2,bias=True)\n",
    "        self.bn31d = nn.BatchNorm2d(128, momentum= batchNorm_momentum)\n",
    "        self.deconv2_1 = nn.Conv2d(128, 64, kernel_size=5, padding=2,bias=True)\n",
    "        self.bn21d = nn.BatchNorm2d(64, momentum= batchNorm_momentum)\n",
    "        self.deconv1_1 = nn.Conv2d(64, 64, kernel_size=5, padding=2,bias=True)\n",
    "        self.bn11d = nn.BatchNorm2d(64, momentum= batchNorm_momentum)\n",
    "        \n",
    "        self.deconv1 = nn.Conv2d(64, 1, kernel_size=5, padding=2,bias=True)\n",
    "        \n",
    "        \n",
    "    def forward(self,x):\n",
    "              # Stage 1\n",
    "        x11 = F.relu(self.bn11(self.conv1_1(x)))\n",
    "        x12 = F.relu(self.bn12(self.conv1_2(x11)))\n",
    "        x1p, id1 = F.max_pool2d(x12,kernel_size=(2,2), stride=(2,2),return_indices=True)\n",
    "\n",
    "        # Stage 2\n",
    "        x21 = F.relu(self.bn21(self.conv2_1(x1p)))\n",
    "        x22 = F.relu(self.bn22(self.conv2_2(x21)))\n",
    "        x2p, id2 = F.max_pool2d(x22,kernel_size=(2,2), stride=(2,2),return_indices=True)\n",
    "\n",
    "        # Stage 3\n",
    "        x31 = F.relu(self.bn31(self.conv3_1(x2p)))\n",
    "        x32 = F.relu(self.bn32(self.conv3_2(x31)))\n",
    "        x33 = F.relu(self.bn33(self.conv3_3(x32)))\n",
    "        x3p, id3 = F.max_pool2d(x33,kernel_size=(2,2), stride=(2,2),return_indices=True)\n",
    "\n",
    "        # Stage 4\n",
    "        x41 = F.relu(self.bn41(self.conv4_1(x3p)))\n",
    "        x42 = F.relu(self.bn42(self.conv4_2(x41)))\n",
    "        x43 = F.relu(self.bn43(self.conv4_3(x42)))\n",
    "        x4p, id4 = F.max_pool2d(x43,kernel_size=(2,2), stride=(2,2),return_indices=True)\n",
    "\n",
    "        # Stage 5\n",
    "        x51 = F.relu(self.bn51(self.conv5_1(x4p)))\n",
    "        x52 = F.relu(self.bn52(self.conv5_2(x51)))\n",
    "        x53 = F.relu(self.bn53(self.conv5_3(x52)))\n",
    "        x5p, id5 = F.max_pool2d(x53,kernel_size=(2,2), stride=(2,2),return_indices=True)\n",
    "\n",
    "        # Stage 6\n",
    "        x61 = F.relu(self.bn61(self.conv6_1(x5p)))\n",
    "\n",
    "        # Stage 6d\n",
    "\n",
    "        x61d = F.relu(self.bn61d(self.deconv6_1(x61)))\n",
    "\n",
    "\n",
    "        # Stage 5d\n",
    "        x5d = F.max_unpool2d(x61d,id5, kernel_size=2, stride=2)\n",
    "        x51d = F.relu(self.bn51d(self.deconv5_1(x5d)))\n",
    "\n",
    "\n",
    "\n",
    "        # Stage 4d\n",
    "        x4d = F.max_unpool2d(x51d, id4, kernel_size=2, stride=2)\n",
    "        x41d = F.relu(self.bn41d(self.deconv4_1(x4d)))\n",
    "\n",
    "        # Stage 3d\n",
    "        x3d = F.max_unpool2d(x41d, id3, kernel_size=2, stride=2)\n",
    "        x31d = F.relu(self.bn31d(self.deconv3_1(x3d)))\n",
    "\n",
    "        # Stage 2d\n",
    "        x2d = F.max_unpool2d(x31d, id2, kernel_size=2, stride=2)\n",
    "        x21d = F.relu(self.bn21d(self.deconv2_1(x2d)))\n",
    "\n",
    "        # Stage 1d\n",
    "        x1d = F.max_unpool2d(x21d, id1, kernel_size=2, stride=2)\n",
    "        x12d = F.relu(self.bn11d(self.deconv1_1(x1d)))\n",
    "        x11d = F.sigmoid(self.deconv1(x12d))\n",
    "\n",
    "        return x11d\n",
    "\n",
    "    def load_my_state_dict(self, model_dict):\n",
    "\n",
    "        own_state = self.state_dict()\n",
    "        #print(own_state.keys())\n",
    "        own_state_keys = self.state_dict().keys()\n",
    "        model_state = model_dict\n",
    "        model_p = 0\n",
    "        for count, name in enumerate(model_state.keys()):\n",
    "\n",
    "            if(count % 2 == 0 and not count==0):\n",
    "                model_p+=4\n",
    "            #print(count, model_p)\n",
    "            if count == 28:\n",
    "                break\n",
    "            if count == 0:\n",
    "                #print(model_state[name].shape)\n",
    "                own_state[own_state_keys[model_p]].copy_(torch.cat((model_state[name], torch.zeros(64,1,3,3)),1))\n",
    "            else:\n",
    "                if count == 26:\n",
    "                    own_state[own_state_keys[model_p]].copy_(model_state[name].view((4096,512,7,7)))\n",
    "                else:\n",
    "                    #print(count, name)\n",
    "                    #print(own_state_keys[model_p], name)\n",
    "                    #print(own_state[own_state_keys[model_p]].shape, model_state[name].shape)\n",
    "                    own_state[own_state_keys[model_p]].copy_(model_state[name])\n",
    "            model_p+=1\n",
    "                \n",
    "                    \n",
    "                                                           \n",
    "                \n",
    "\n",
    "\n",
    "# input_nbr = 3\n",
    "# imsize = 224\n",
    "\n",
    "# # Training settings\n",
    "# parser = argparse.ArgumentParser(description='PyTorch SegNet example')\n",
    "# parser.add_argument('--batch-size', type=int, default=32, metavar='N',\n",
    "#                     help='input batch size for training (default: 64)')\n",
    "# parser.add_argument('--epochs', type=int, default=300, metavar='N',\n",
    "#                     help='number of epochs to train (default: 10)')\n",
    "# parser.add_argument('--lr', type=float, default=0.01, metavar='LR',\n",
    "#                     help='learning rate (default: 0.01)')\n",
    "# parser.add_argument('--momentum', type=float, default=0.5, metavar='M',\n",
    "#                     help='SGD momentum (default: 0.5)')\n",
    "# parser.add_argument('--no-cuda', action='store_true', default=False,\n",
    "#                     help='enables CUDA training')\n",
    "# parser.add_argument('--seed', type=int, default=1, metavar='S',\n",
    "#                     help='random seed (default: 1)')\n",
    "\n",
    "args = {}\n",
    "\n",
    "args['cuda'] = True\n",
    "args['resume'] = False\n",
    "args['seed'] = 1\n",
    "# cuda\n",
    "\n",
    "args['cuda'] = torch.cuda.is_available()\n",
    "USE_CUDA = True\n",
    "# set the seed\n",
    "torch.manual_seed(args['seed'])\n",
    "if args['cuda']:\n",
    "    torch.cuda.manual_seed(args['seed'])\n",
    "\n",
    "\n",
    "model = DeepMatting()\n",
    "if USE_CUDA:# convert to cuda if needed\n",
    "    model.cuda()\n",
    "else:\n",
    "    model.float()\n",
    "\n",
    "# define the optimizer\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-5)\n",
    "\n",
    "\n",
    "def train():\n",
    "    paths_alpha, paths_merged, paths_BG = load_path(dataset_alpha, dataset_merged, dataset_BG, hard_mode = False)\n",
    "    model.train()\n",
    "    model.apply(weights_init)\n",
    "    initial_epoch = 0\n",
    "    if args['resume']:\n",
    "        if os.path.isfile(args['resume']):\n",
    "            print(\"=> loading checkpoint '{}'\".format(args['resume']))\n",
    "            checkpoint = torch.load(args['resume'])\n",
    "            args['start_epoch'] = checkpoint['epoch']\n",
    "            initial_epoch = args['start_epoch']\n",
    "            best_prec1 = checkpoint['best_prec1']\n",
    "            model.load_state_dict(checkpoint['state_dict'])\n",
    "            optimizer.load_state_dict(checkpoint['optimizer'])\n",
    "            print(\"=> loaded checkpoint '{}' (epoch {})\"\n",
    "                  .format(args['resume'], checkpoint['epoch']))\n",
    "    else:\n",
    "        print(\"=> no checkpoint found at '{}'\".format(args['resume']))\n",
    "        #TODO_1 load_weights according to the deep matting code + match layers for loading\n",
    "        model.load_my_state_dict( models.vgg16(pretrained=True).state_dict())\n",
    "        #model.load_weights(\"vgg16-00b39a1b.pth\") \n",
    "\n",
    "    batch_size = 1\n",
    "    for epoch in range(initial_epoch, max_epochs):\n",
    "        batch = 0\n",
    "        total_batches = len(paths_merged)/batch_size\n",
    "        while batch < total_batches:\n",
    "            batch_alpha_paths = paths_alpha[batch*batch_size:(batch+1)*batch_size]\n",
    "            batch_merged_paths = paths_merged[batch*batch_size:(batch+1)*batch_size]\n",
    "            batch_BG_paths = paths_BG[batch*batch_size:(batch+1)*batch_size]\n",
    "            if(len(misc.imread(batch_BG_paths[0]).shape) < 3):\n",
    "                continue\n",
    "                \n",
    "            batch_RGBs, batch_trimaps, batch_alphas, batch_BGs, batch_FGs, RGBs_with_mean = load_data(batch_alpha_paths, batch_merged_paths,batch_BG_paths)\n",
    "    \n",
    "            batch_RGBsT, batch_trimapsT, batch_alphasT, batch_BGsT, batch_FGsT, RGBs_with_meanT = [Variable(torch.Tensor(batch_RGBs.astype(np.float64))),Variable(torch.Tensor(batch_trimaps.astype(np.float64))),Variable(torch.Tensor(batch_alphas.astype(np.float64))),Variable(torch.Tensor(batch_BGs.astype(np.float64))),Variable(torch.Tensor(batch_FGs.astype(np.float64))),Variable(torch.Tensor(RGBs_with_mean))]\n",
    "            \n",
    "            batch_RGBsT, batch_trimapsT, batch_alphasT, batch_BGsT, batch_FGsT, RGBs_with_meanT = [batch_RGBsT.permute(0,3,1,2), batch_trimapsT.permute(0,3,1,2), batch_alphasT.permute(0,3,1,2), batch_BGsT.permute(0,3,1,2), batch_FGsT.permute(0,3,1,2), RGBs_with_meanT.permute(0,3,1,2)]\n",
    "            if USE_CUDA:\n",
    "                batch_RGBsT, batch_trimapsT, batch_alphasT, batch_BGsT, batch_FGsT, RGBs_with_meanT = [batch_RGBsT.cuda(), batch_trimapsT.cuda(), batch_alphasT.cuda(), batch_BGsT.cuda(), batch_FGsT.cuda(), RGBs_with_meanT.cuda()]\n",
    "\n",
    "\n",
    "            # initilize gradients\n",
    "            optimizer.zero_grad()\n",
    "            b_input = torch.cat((batch_RGBsT,batch_trimapsT),1)\n",
    "            \n",
    "\n",
    "            # predictions\n",
    "            pred_mattes = model(b_input)\n",
    "            \n",
    "\n",
    "            alpha_diff = torch.sqrt((pred_mattes - batch_alphasT)**2 +1e-12)\n",
    "            #c_diff = torch.sqrt(batch_RGBsT - raw)\n",
    "\n",
    "            cond = torch.eq(batch_trimapsT, 128)\n",
    "            #print(cond.is_cuda)\n",
    "            cond = cond.type(torch.cuda.FloatTensor)\n",
    "            #print(type(cond))\n",
    "            wl =  cond * Variable(torch.ones([batch_size, image_size, image_size, 1]).cuda()) + ((1-cond) *  Variable(torch.zeros([batch_size, image_size, image_size, 1]).cuda())) \n",
    "            unknown_region_size = wl.sum()\n",
    "\n",
    "            pred_final = cond * (pred_mattes) + (1 - cond)*(batch_trimapsT/255.0)\n",
    "\n",
    "            alpha_loss = (alpha_diff * wl).sum()/unknown_region_size\n",
    "\n",
    "            \n",
    "            \n",
    "\n",
    "            total_loss = alpha_loss\n",
    "            total_loss.cuda()\n",
    "            total_loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "            if(batch % 1000 == 0 and not batch==0):\n",
    "                print('Epoch:',epoch,'Batch:', batch, 'Loss:',total_loss)\n",
    "                test()\n",
    "            batch+=1\n",
    "            \n",
    "def test():\n",
    "    model.eval()\n",
    "    test_RGBs, test_trimaps, test_alphas, image_paths = load_alphamatting_path(test_dir)\n",
    "    vali_diff = []\n",
    "    # iteration over the batches\n",
    "    for i in range(100):\n",
    "        #print(test_RGBs[i])\n",
    "        RGB, trimap, test_alpha, shape_i, trimap_size = load_alphamatting_data(test_RGBs[i], test_trimaps[i], test_alphas[i])\n",
    "\n",
    "        test_RGB = Variable(torch.Tensor(np.expand_dims(RGB,0).astype(np.float64))).permute(0,3,1,2).cuda()\n",
    "        test_trimap = Variable(torch.Tensor(np.expand_dims(trimap,0).astype(np.float64))).permute(0,3,1,2).cuda()\n",
    "\n",
    "        b_input = torch.cat((test_RGB, test_trimap),1)\n",
    "\n",
    "\n",
    "        # predictions\n",
    "        test_out = model(b_input)\n",
    "        pred_mattes = misc.imresize(test_out[0,0,:,:].data.cpu().numpy(),shape_i)\n",
    "        x = np.abs(pred_mattes - test_alpha)\n",
    "        y = np.sum(x)\n",
    "        z = y/trimap_size\n",
    "        vali_diff.append(z)\n",
    "    vali_loss = np.mean(vali_diff)\n",
    "    print(\"validation Loss:\",vali_loss)\n",
    "\n",
    "train()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "language": "python",
   "name": "env"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}