Building Detector using SSD in FastAI on Sentinel-2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#  Import deep learning libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from arcgis.gis import GIS\n",
    "\n",
    "import torch\n",
    "from fastai.conv_learner import *\n",
    "from fastai.dataset import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use training data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import json\n",
    "\n",
    "from PIL import ImageDraw, ImageFont\n",
    "from matplotlib import patches, patheffects\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm as cmx\n",
    "import matplotlib.colors as mcolors\n",
    "from cycler import cycler\n",
    "\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# defining the PATH variable which will be used to save and load files\n",
    "PATH = Path('/home/splash/buildings/data/')\n",
    "\n",
    "# tilemapping is a dictionary which contains mapping from image to bounding boxes\n",
    "tilemapping = pickle.load(open(PATH/'tilemapping.pkl', 'rb'))\n",
    "len(tilemapping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This cells defines some utility functions for plotting\n",
    "\n",
    "def get_color(text):\n",
    "    if text=='bg':\n",
    "        c = 'gray'\n",
    "    elif text=='building ':\n",
    "        c = 'white'\n",
    "    else:\n",
    "        c = 'yellow'\n",
    "    return c\n",
    "\n",
    "def draw_outline(obj, width):\n",
    "    obj.set_path_effects([patheffects.Stroke(linewidth=width, foreground='black'), patheffects.Normal()])\n",
    "\n",
    "def draw_bbox(ax, bbox, text=None, i=None, pr=None):    \n",
    "    c = get_color(text)\n",
    "    patch = ax.add_patch(patches.Rectangle(xy=(bbox[1], bbox[0]), width=bbox[3]-bbox[1],\n",
    "                                           height=bbox[2]-bbox[0], \n",
    "                                           fill=False, edgecolor=c, lw=2))\n",
    "    draw_outline(patch, 4)\n",
    "    if i is not None:\n",
    "        text = text\n",
    "    if pr is not None:\n",
    "        text = text + f' {pr:.2f}'\n",
    "    if text is not None:\n",
    "        t = ax.text(bbox[1], bbox[0], s=text, color='white', fontsize=12,\n",
    "                    verticalalignment='bottom', weight='bold') #,bbox=dict(facecolor=c, pad=2))\n",
    "        draw_outline(t, 1)\n",
    "\n",
    "def draw_img_bboxes(ax, img, bboxes, classes=None, prs=None, thresh=0.3):\n",
    "    \n",
    "    if prs is None:  \n",
    "        prs  = [None]*len(bboxes)\n",
    "    if classes is None: \n",
    "        classes = [1]*len(bboxes)\n",
    "    \n",
    "    ax.set_yticklabels([])\n",
    "    ax.set_xticklabels([])\n",
    "    ax.imshow(img);\n",
    "    for i, (bbox, c, pr) in enumerate(zip(bboxes, classes, prs)):\n",
    "        if((bbox[3] - bbox[1] > 0) and (pr is None or pr > thresh)):\n",
    "            try: cat = categories[c]\n",
    "            except: cat = 'bg'\n",
    "            draw_bbox(ax, bbox, cat, i, pr)\n",
    "            \n",
    "def draw_image(filename):\n",
    "    fig, ax = plt.subplots(figsize=(4,4))\n",
    "    img = open_image(PATH / 'images' / filename)\n",
    "    \n",
    "    bboxes=[]\n",
    "    for center in tilemapping[filename]:\n",
    "        bboxes.append(get_bbox(center))\n",
    "        \n",
    "    draw_img_bboxes(ax, img, bboxes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define categories which the image contains\n",
    "categories = {0: 'ground', 1: 'building'}\n",
    "\n",
    "sz = 224\n",
    "objsize = 30\n",
    "\n",
    "\n",
    "# In dataframe the bounding boxes coordinates are passed as string which are joined using spaces\n",
    "# The following function convert bounding boxes to string\n",
    "def get_bboxes_str(filename):\n",
    "    bboxes = ''\n",
    "    for center in tilemapping[filename]:\n",
    "        bboxes = bboxes + ' '.join(str(c) for c in (get_bbox(center)))\n",
    "        bboxes = bboxes + ' '\n",
    "    return bboxes\n",
    "\n",
    "# a function to convert bounding center location coordinates to boxxes\n",
    "def get_bbox(center):\n",
    "    xpix, ypix = int(sz*center[0]), int(sz*(1.0-center[1]))\n",
    "        \n",
    "    xmin, xmax = xpix - objsize//2, xpix + objsize//2 + 2\n",
    "    ymin, ymax = ypix - objsize//2, ypix + objsize//2 + 2\n",
    "    xmin = xmin # if xmin > 0 else 0\n",
    "    ymin = ymin # if ymin > 0 else 0\n",
    "    xmax = xmax # if xmax < sz else sz\n",
    "    ymax = ymax # if ymax < sz else sz\n",
    "    return [ymin, xmin, ymax, xmax]\n",
    "\n",
    "bboxes_dict = {name: get_bboxes_str(name).strip()\n",
    "                   for name, centers in tilemapping.items()}\n",
    "\n",
    "# creating a pandas dataframe in the specific format\n",
    "bboxes_df = pd.DataFrame.from_dict(bboxes_dict, orient='index').reset_index()\n",
    "bboxes_df.columns = ['fn', 'bbox']\n",
    "\n",
    "MULTIBOX_CSV = PATH/'bboxes.csv' # CSV file to use\n",
    "bboxes_df.sort_values(by='fn', inplace=True)\n",
    "bboxes_df.to_csv(MULTIBOX_CSV, index=False)\n",
    "bboxes_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = resnet34                 # Base model to be used for transfer learning\n",
    "sz = 224                         # size of images to be used\n",
    "bs = 64                          # batch size\n",
    "JPEGS = 'images'                 # folder name which contains images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reading the CSV from the disk\n",
    "bboxes_df = pd.read_csv(MULTIBOX_CSV)\n",
    "\n",
    "# Calculating the number of objects column for every row in the Dataframe\n",
    "bboxes_df['n'] = bboxes_df['bbox'].apply(lambda x: len(x.split(' ')) // 4)\n",
    "\n",
    "object_classes = np.array([np.ones(row['n']).astype(int) for index, row in bboxes_df.iterrows()])\n",
    "\n",
    "# The following line splits the training and validation sets based on their index\n",
    "val_idxs = get_cv_idxs(len(object_classes))\n",
    "((val_classes, train_classes),) = split_by_idx(val_idxs, object_classes)\n",
    "\n",
    "len(val_classes), len(train_classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Defining the augmentation and getting the model data object which will be used to create the model\n",
    "\n",
    "aug_tfms = [RandomRotate(3, p=0.5, tfm_y=TfmType.COORD),\n",
    "            RandomLighting(0.05, 0.5, tfm_y=TfmType.COORD),\n",
    "            RandomFlip(tfm_y=TfmType.COORD)]\n",
    "\n",
    "tfms = tfms_from_model(model, sz, crop_type=CropType.NO, tfm_y=TfmType.COORD, aug_tfms=aug_tfms)\n",
    "\n",
    "model_data = ImageClassifierData.from_csv(PATH, JPEGS, MULTIBOX_CSV, tfms=tfms, continuous=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating a python class to merge classes of the bounding boxes into the model_data\n",
    "\n",
    "class ConcatLblDataset(Dataset):\n",
    "    def __init__(self, ds, y2):\n",
    "        self.ds, self.y2 = ds, y2\n",
    "        self.sz = ds.sz\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.ds)\n",
    "    \n",
    "    def __getitem__(self, i):\n",
    "        x, y1 = self.ds[i]\n",
    "        y2 = self.y2[i]\n",
    "        return (x, (y1, y2)) # y1 is bboxes y2 is classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# merging the classes into the model_data object with the class of the label\n",
    "\n",
    "train_dataset = ConcatLblDataset(model_data.trn_ds, train_classes)\n",
    "val_dataset   = ConcatLblDataset(model_data.val_ds, val_classes)\n",
    "\n",
    "model_data.trn_dl.dataset = train_dataset\n",
    "model_data.val_dl.dataset = val_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = to_np(next(iter(model_data.val_dl)))\n",
    "x    = model_data.val_ds.ds.denorm(x)\n",
    "\n",
    "def show_training_data():\n",
    "    fig, axes = plt.subplots(3, 4, figsize=(16, 12))\n",
    "    for i,ax in enumerate(axes.flat):\n",
    "        draw_img_bboxes(ax, x[i], y[0][i].reshape(-1, 4), y[1][i])\n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "#### Visualize image chips with bounding boxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "show_training_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train AI model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Single Shot Detector (SSD) model architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class StdConv(nn.Module):\n",
    "    def __init__(self, nin, nout, stride=2, drop=0.1):\n",
    "        super().__init__()\n",
    "        self.conv = nn.Conv2d(nin, nout, 3, stride=stride, padding=1)\n",
    "        self.bn = nn.BatchNorm2d(nout)\n",
    "        self.drop = nn.Dropout(drop)\n",
    "        \n",
    "    def forward(self, x): return self.drop(self.bn(F.relu(self.conv(x))))\n",
    "        \n",
    "def flatten_conv(x, k):\n",
    "    bs, nf, gx, gy = x.size()                # batch size, num filters, width, height\n",
    "    x = x.permute(0, 3, 2, 1).contiguous()   # batch_size, \n",
    "    return x.view(bs, -1, nf//k)\n",
    "\n",
    "class OutConv(nn.Module):\n",
    "    def __init__(self, k, nin, bias):\n",
    "        super().__init__()\n",
    "        self.k = k\n",
    "        self.oconv1 = nn.Conv2d(nin, (len(id2cat)+1)*k, 3, padding=1) # nclasses\n",
    "        self.oconv2 = nn.Conv2d(nin, 4*k, 3, padding=1) # bboxes\n",
    "        self.oconv1.bias.data.zero_().add_(bias)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return [flatten_conv(self.oconv1(x), self.k),\n",
    "                flatten_conv(self.oconv2(x), self.k)]\n",
    "\n",
    "\n",
    "class SSD_MultiHead(nn.Module):\n",
    "    def __init__(self, k, bias):\n",
    "        super().__init__()\n",
    "        self.drop = nn.Dropout(drop)\n",
    "        self.sconv0 = StdConv(512,256, stride=1, drop=drop)\n",
    "        self.sconv1 = StdConv(256,256, drop=drop)\n",
    "        self.sconv2 = StdConv(256,256, drop=drop)\n",
    "        self.sconv3 = StdConv(256,256, drop=drop)\n",
    "        self.out0 = OutConv(k, 256, bias)\n",
    "        self.out1 = OutConv(k, 256, bias)\n",
    "        self.out2 = OutConv(k, 256, bias)\n",
    "        self.out3 = OutConv(k, 256, bias)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.drop(F.relu(x))\n",
    "        x = self.sconv0(x)\n",
    "        x = self.sconv1(x)\n",
    "        o1c,o1l = self.out1(x)\n",
    "        x = self.sconv2(x)\n",
    "        o2c,o2l = self.out2(x)\n",
    "        x = self.sconv3(x)\n",
    "        o3c,o3l = self.out3(x)\n",
    "        return [torch.cat([o1c, o2c, o3c], dim=1),\n",
    "                torch.cat([o1l, o2l, o3l], dim=1)]\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# converts height width to corners\n",
    "def hw2corners(ctr, hw): \n",
    "    return torch.cat([ctr-hw/2, ctr+hw/2], dim=1)\n",
    "\n",
    "anc_grids = [4,2,1]\n",
    "anc_zooms = [0.7, 1., 1.3]\n",
    "anc_ratios = [(1.,1.), (1.,0.5), (0.5,1.)]\n",
    "anchor_scales = [(anz*i,anz*j) for anz in anc_zooms for (i,j) in anc_ratios]\n",
    "k = len(anchor_scales)\n",
    "anc_offsets = [1/(o*2) for o in anc_grids]\n",
    "\n",
    "anc_x = np.concatenate([np.repeat(np.linspace(ao, 1-ao, ag), ag)\n",
    "                        for ao,ag in zip(anc_offsets,anc_grids)])\n",
    "anc_y = np.concatenate([np.tile(np.linspace(ao, 1-ao, ag), ag)\n",
    "                        for ao,ag in zip(anc_offsets,anc_grids)])\n",
    "anc_ctrs = np.repeat(np.stack([anc_x,anc_y], axis=1), k, axis=0)\n",
    "\n",
    "anc_sizes  =   np.concatenate([np.array([[o/ag,p/ag] for i in range(ag*ag) for o,p in anchor_scales])\n",
    "               for ag in anc_grids])\n",
    "grid_sizes = V(np.concatenate([np.array([ 1/ag       for i in range(ag*ag) for o,p in anchor_scales])\n",
    "               for ag in anc_grids]), requires_grad=False).unsqueeze(1)\n",
    "anchors = V(np.concatenate([anc_ctrs, anc_sizes], axis=1), requires_grad=False).float()\n",
    "anchor_cnr = hw2corners(anchors[:,:2], anchors[:,2:])\n",
    "\n",
    "id2cat = list(categories.values())\n",
    "cat2id = {v:k for k,v in enumerate(id2cat)}\n",
    "\n",
    "n_clas = len(id2cat)+1\n",
    "n_act = k*(4+n_clas)\n",
    "\n",
    "n_clas, n_act"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f_model = model # defining the base model for transfer learning\n",
    "md = model_data # definig model data\n",
    "\n",
    "drop=0.4\n",
    "\n",
    "head_reg4 = SSD_MultiHead(k, -4.)\n",
    "models = ConvnetBuilder(f_model, 0, 0, 0, custom_head=head_reg4)\n",
    "learn = ConvLearner(md, models)\n",
    "learn.opt_fn = optim.Adam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# defining the loss functions\n",
    "\n",
    "def one_hot_embedding(labels, num_classes):\n",
    "    return torch.eye(num_classes)[labels.data.cpu()]\n",
    "\n",
    "class BCE_Loss(nn.Module):\n",
    "    def __init__(self, num_classes):\n",
    "        super().__init__()\n",
    "        self.num_classes = num_classes\n",
    "\n",
    "    def forward(self, pred, targ):\n",
    "        t = one_hot_embedding(targ, self.num_classes+1)\n",
    "        t = V(t[:,:-1].contiguous())#.cpu()\n",
    "        x = pred[:,:-1]\n",
    "        w = self.get_weight(x,t)\n",
    "        return F.binary_cross_entropy_with_logits(x, t, w, size_average=False)/self.num_classes\n",
    "    \n",
    "    def get_weight(self,x,t): return None\n",
    "\n",
    "loss_f = BCE_Loss(len(id2cat))\n",
    "\n",
    "def intersect(box_a, box_b):\n",
    "    max_xy = torch.min(box_a[:, None, 2:], box_b[None, :, 2:])\n",
    "    min_xy = torch.max(box_a[:, None, :2], box_b[None, :, :2])\n",
    "    inter = torch.clamp((max_xy - min_xy), min=0)\n",
    "    return inter[:, :, 0] * inter[:, :, 1]\n",
    "\n",
    "def box_sz(b): return ((b[:, 2]-b[:, 0]) * (b[:, 3]-b[:, 1]))\n",
    "\n",
    "def jaccard(box_a, box_b):\n",
    "    try:\n",
    "        inter = intersect(box_a, box_b)\n",
    "        union = box_sz(box_a).unsqueeze(1) + box_sz(box_b).unsqueeze(0) - inter\n",
    "        return inter / union\n",
    "    except:\n",
    "        return 0.\n",
    "\n",
    "def get_y(bbox,clas):\n",
    "    bbox = bbox.view(-1,4)/sz\n",
    "    try:\n",
    "        xx = ((bbox[:,2]-bbox[:,0])>0).nonzero()\n",
    "        #print(xx)\n",
    "        bb_keep = xx[:,0]\n",
    "    except:\n",
    "        # print('No bboxes')\n",
    "        return None, None\n",
    "    \n",
    "    #xx = ((bbox[:,2]-bbox[:,0])>0).nonzero()\n",
    "    #bb_keep = xx[:,0]\n",
    "    \n",
    "    return bbox[bb_keep], clas[bb_keep]\n",
    "\n",
    "def actn_to_bb(actn, anchors):\n",
    "    actn_bbs = torch.tanh(actn)\n",
    "    actn_centers = (actn_bbs[:,:2]/2 * grid_sizes) + anchors[:,:2]\n",
    "    actn_hw = (actn_bbs[:,2:]/2+1) * anchors[:,2:] # [-1, 1]/2 + 1 = [0.5, 1.5] * anchor ht, w\n",
    "    return hw2corners(actn_centers, actn_hw)\n",
    "\n",
    "def map_to_ground_truth(overlaps, print_it=False):\n",
    "    prior_overlap, prior_idx = overlaps.max(1)\n",
    "    if print_it: print(prior_overlap)\n",
    "#     pdb.set_trace()\n",
    "    gt_overlap, gt_idx = overlaps.max(0)\n",
    "    gt_overlap[prior_idx] = 1.99\n",
    "    for i,o in enumerate(prior_idx): gt_idx[o] = i\n",
    "    return gt_overlap,gt_idx\n",
    "\n",
    "def ssd_1_loss(b_c, b_bb, bbox, clas, print_it=False):\n",
    "    bbox, clas = get_y(bbox, clas)             # bbox, class from target\n",
    "    if bbox is None:\n",
    "        return 0.0, 0.0\n",
    "    a_ic = actn_to_bb(b_bb, anchors)           # predicted bbox\n",
    "    overlaps = jaccard(bbox.data, anchor_cnr.data)\n",
    "    gt_overlap,gt_idx = map_to_ground_truth(overlaps,print_it)\n",
    "    gt_clas = clas[gt_idx]\n",
    "    pos = gt_overlap > 0.4\n",
    "    pos_idx = torch.nonzero(pos)[:,0]\n",
    "    gt_clas[1-pos] = len(id2cat)\n",
    "    gt_bbox = bbox[gt_idx]\n",
    "    loc_loss = ((a_ic[pos_idx] - gt_bbox[pos_idx]).abs()).mean()\n",
    "    clas_loss  = loss_f(b_c, gt_clas)  # binary cross entropy\n",
    "    return loc_loss, clas_loss\n",
    "\n",
    "def ssd_loss(pred, targ, print_it=False):\n",
    "    lcs, lls = 0.,0. # class loss, location loss\n",
    "    for b_c, b_bb, bbox, clas in zip(*pred, *targ): \n",
    "        loc_loss, clas_loss = ssd_1_loss(b_c, b_bb, bbox, clas, print_it)\n",
    "        lls += loc_loss\n",
    "        lcs += clas_loss\n",
    "    if print_it: print(f'loc: {lls.data[0]}, clas: {lcs.data[0]}')\n",
    "    return lls+lcs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anchors = anchors.cuda();\n",
    "grid_sizes = grid_sizes.cuda(); \n",
    "anchor_cnr = anchor_cnr.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.crit = ssd_loss\n",
    "lr = 3e-3\n",
    "lrs = np.array([lr/100,lr/10,lr])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_torch_img(ax, ima, bbox, clas, prs=None, thresh=0.4):\n",
    "    return draw_img_bboxes(ax, ima, to_np((bbox*224).long()),\n",
    "         to_np(clas), to_np(prs) if prs is not None else None, thresh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def map_to_ground_truth(overlaps, print_it=False):\n",
    "    prior_overlap, prior_idx = overlaps.max(1)\n",
    "\n",
    "    gt_overlap, gt_idx = overlaps.max(0)\n",
    "    gt_overlap[prior_idx] = 1.99\n",
    "    for i,o in enumerate(prior_idx): \n",
    "        gt_idx[o] = i\n",
    "    return gt_overlap, gt_idx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train model using labeled training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.lr_find(lrs/1000,1.)\n",
    "learn.sched.plot(n_skip_end=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.fit(lrs, 1, cycle_len=4, use_clr=(20,8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.fit(lrs, 1, cycle_len=4, use_clr=(20,8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Freeze all layers except last two(run it twice)\n",
    "learn.freeze_to(-2)\n",
    "learn.fit(lrs/2, 1, cycle_len=4, use_clr=(20,8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_results(thresh):\n",
    "    x,y = next(iter(md.val_dl))\n",
    "    y = V(y)\n",
    "    batch = learn.model(V(x))\n",
    "    b_clas,b_bb = batch\n",
    "\n",
    "    x = to_np(x)\n",
    "    fig, axes = plt.subplots(3, 4, figsize=(16, 12))\n",
    "    for idx,ax in enumerate(axes.flat):\n",
    "        ima=md.val_ds.ds.denorm(x)[idx]\n",
    "        bbox,clas = get_y(y[0][idx], y[1][idx])\n",
    "        a_ic = actn_to_bb(b_bb[idx], anchors)\n",
    "        clas_pr, clas_ids = b_clas[idx].max(1)\n",
    "        clas_pr = clas_pr.sigmoid()\n",
    "        draw_torch_img(ax, ima, a_ic, clas_ids, clas_pr, clas_pr.max().data[0]*thresh)\n",
    "    plt.tight_layout()\n",
    "\n",
    "#changing loss function to focal loss    \n",
    "class FocalLoss(BCE_Loss):\n",
    "    def get_weight(self,x,t):\n",
    "        alpha,gamma = 0.25,1\n",
    "        p = x.sigmoid()\n",
    "        pt = p*t + (1-p)*(1-t)\n",
    "        w = alpha*t + (1-alpha)*(1-t)\n",
    "        return w * (1-pt).pow(gamma)\n",
    "\n",
    "loss_f = FocalLoss(len(id2cat))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.lr_find(lrs/1000,1.)\n",
    "learn.sched.plot(n_skip_end=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.fit(lrs/10, 1, cycle_len=5, use_clr=(20,10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.freeze_to(-2)\n",
    "learn.fit(lrs/50, 1, cycle_len=4, use_clr=(20,10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.save('8')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# unfreeze all layers and fine tune the complete model\n",
    "learn.unfreeze()\n",
    "learn.fit(lrs/250, 1, cycle_len=4, use_clr=(20,10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.save('9')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Detect swimming buildings in validation set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.load('9')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the non maximal suppression function\n",
    "def nms(boxes, scores, overlap=0.5, top_k=100):\n",
    "    keep = scores.new(scores.size(0)).zero_().long()\n",
    "    if boxes.numel() == 0:\n",
    "        return keep\n",
    "    x1 = boxes[:, 0]\n",
    "    y1 = boxes[:, 1]\n",
    "    x2 = boxes[:, 2]\n",
    "    y2 = boxes[:, 3]\n",
    "    area = torch.mul(x2 - x1, y2 - y1)\n",
    "    v, idx = scores.sort(0)  # sort in ascending order\n",
    "    idx = idx[-top_k:]  # indices of the top-k largest vals\n",
    "    xx1 = boxes.new()\n",
    "    yy1 = boxes.new()\n",
    "    xx2 = boxes.new()\n",
    "    yy2 = boxes.new()\n",
    "    w = boxes.new()\n",
    "    h = boxes.new()\n",
    "\n",
    "    count = 0\n",
    "    while idx.numel() > 0:\n",
    "        i = idx[-1]  # index of current largest val\n",
    "        keep[count] = i\n",
    "        count += 1\n",
    "        if idx.size(0) == 1: break\n",
    "        idx = idx[:-1]  # remove kept element from view\n",
    "        # load bboxes of next highest vals\n",
    "        torch.index_select(x1, 0, idx, out=xx1)\n",
    "        torch.index_select(y1, 0, idx, out=yy1)\n",
    "        torch.index_select(x2, 0, idx, out=xx2)\n",
    "        torch.index_select(y2, 0, idx, out=yy2)\n",
    "        # store element-wise max with next highest score\n",
    "        xx1 = torch.clamp(xx1, min=x1[i])\n",
    "        yy1 = torch.clamp(yy1, min=y1[i])\n",
    "        xx2 = torch.clamp(xx2, max=x2[i])\n",
    "        yy2 = torch.clamp(yy2, max=y2[i])\n",
    "        w.resize_as_(xx2)\n",
    "        h.resize_as_(yy2)\n",
    "        w = xx2 - xx1\n",
    "        h = yy2 - yy1\n",
    "        # check sizes of xx1 and xx2.. after each iteration\n",
    "        w = torch.clamp(w, min=0.0)\n",
    "        h = torch.clamp(h, min=0.0)\n",
    "        inter = w*h\n",
    "        # IoU = i / (area(a) + area(b) - i)\n",
    "        rem_areas = torch.index_select(area, 0, idx)  # load remaining areas)\n",
    "        union = (rem_areas - inter) + area[i]\n",
    "        IoU = inter/union  # store result in iou\n",
    "        # keep only elements with an IoU <= overlap\n",
    "        idx = idx[IoU.le(overlap)]\n",
    "    return keep, count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# take out a batch\n",
    "x,y = next(iter(md.val_dl))\n",
    "y = V(y)\n",
    "batch = learn.model(V(x))\n",
    "b_clas,b_bb = batch\n",
    "x = to_np(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# non maximal suppression with comparisons\n",
    "def show_results(idx, thresh=0.25, overlap=0.1):\n",
    "    ima=md.val_ds.ds.denorm(x)[idx]\n",
    "    bbox,clas = get_y(y[0][idx], y[1][idx])\n",
    "    a_ic = actn_to_bb(b_bb[idx], anchors)\n",
    "    clas_pr, clas_ids = b_clas[idx].max(1)\n",
    "    clas_pr = clas_pr.sigmoid()\n",
    "\n",
    "    conf_scores = b_clas[idx].sigmoid().t().data\n",
    "\n",
    "    out1,out2,cc = [],[],[]\n",
    "    for cl in range(0, len(conf_scores)-1):\n",
    "        c_mask = conf_scores[cl] > thresh\n",
    "        if c_mask.sum() == 0: continue\n",
    "        scores = conf_scores[cl][c_mask]\n",
    "        l_mask = c_mask.unsqueeze(1).expand_as(a_ic)\n",
    "        boxes = a_ic[l_mask].view(-1, 4)\n",
    "        ids, count = nms(boxes.data, scores, overlap, 50)\n",
    "        ids = ids[:count]\n",
    "        out1.append(scores[ids])\n",
    "        out2.append(boxes.data[ids])\n",
    "        cc.append([cl]*count)\n",
    "    if cc == []:\n",
    "        cc = [[0]]\n",
    "    if out1 == []:\n",
    "        out1 = [torch.Tensor()]\n",
    "    if out2 == []:\n",
    "        out2 = [torch.Tensor()]\n",
    "    cc = T(np.concatenate(cc))\n",
    "    out1 = torch.cat(out1)\n",
    "    out2 = torch.cat(out2)\n",
    "\n",
    "    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(15, 45))\n",
    "    \n",
    "    torch_gt(ax2, ima, out2, cc, out1, 0.)\n",
    "    title1_str = '# of predicted objects: '+ str(len([x for x in out1 if x > 0]))\n",
    "    ax2.text(0,-10,title1_str)\n",
    "    \n",
    "    if bbox is None:\n",
    "        bbox = torch.Tensor([0,0,0, 0])\n",
    "    if clas is None:\n",
    "        clas = torch.Tensor(0)\n",
    "    torch_gt(ax3, ima, bbox, clas, None, 0.)\n",
    "    title2_str = '# of ground truth objects: '+ str(len(clas))\n",
    "    ax3.text(0,-10,title2_str)\n",
    "    \n",
    "    show_img(ima, ax=ax1)\n",
    "    title3_str = md.val_ds.ds.fnames[i]\n",
    "    ax1.text(0,-10,title3_str)\n",
    "    ax1.grid(False)    \n",
    "    \n",
    "def torch_gt(ax, ima, bbox, clas, prs=None, thresh=0.4):\n",
    "    return show_ground_truth(ax, ima, to_np((bbox*224).long()),\n",
    "         to_np(clas), to_np(prs) if prs is not None else None, thresh)\n",
    "\n",
    "def show_ground_truth(ax, im, bbox, clas=None, prs=None, thresh=0.3):\n",
    "    bb = [bb_hw(o) for o in bbox.reshape(-1,4)]\n",
    "    if prs is None:  prs  = [None]*len(bb)\n",
    "    if clas is None: clas = [None]*len(bb)\n",
    "    ax = show_img(im, ax=ax)\n",
    "    for i,(b,c,pr) in enumerate(zip(bb, clas, prs)):\n",
    "        if((b[2]>0) and (pr is None or pr > thresh)):\n",
    "            draw_rect(ax, b, color=colr_list[i%num_colr])\n",
    "            txt = f''\n",
    "#             print(c)\n",
    "            if c is not None: txt += ('bg' if c==len(id2cat) else id2cat[c])\n",
    "            if pr is not None: txt += f' {pr:.2f}'\n",
    "            draw_text(ax, b[:2], txt, color=colr_list[i%num_colr])    \n",
    "            \n",
    "def bb_hw(a): return np.array([a[1],a[0],a[3]-a[1]+1,a[2]-a[0]+1])\n",
    "def show_img(im, figsize=None, ax=None):\n",
    "    if not ax: fig,ax = plt.subplots(figsize=figsize)\n",
    "    ax.imshow(im)\n",
    "    ax.set_xticks(np.linspace(0, 224, 8))\n",
    "    ax.set_yticks(np.linspace(0, 224, 8))\n",
    "#     ax.grid()\n",
    "    ax.set_yticklabels([])\n",
    "    ax.set_xticklabels([])\n",
    "    return ax\n",
    "\n",
    "def draw_outline(o, lw):\n",
    "    o.set_path_effects([patheffects.Stroke(\n",
    "        linewidth=lw, foreground='black'), patheffects.Normal()])\n",
    "\n",
    "def draw_rect(ax, b, color='white'):\n",
    "    patch = ax.add_patch(patches.Rectangle(b[:2], *b[-2:], fill=False, edgecolor=color, lw=2))\n",
    "    draw_outline(patch, 4)\n",
    "\n",
    "def draw_text(ax, xy, txt, sz=14, color='white'):\n",
    "    text = ax.text(*xy, txt,\n",
    "        verticalalignment='top', color=color, fontsize=sz, weight='bold')\n",
    "    draw_outline(text, 1)\n",
    "    \n",
    "def get_cmap(N):\n",
    "    color_norm  = mcolors.Normalize(vmin=0, vmax=N-1)\n",
    "    return cmx.ScalarMappable(norm=color_norm, cmap='Set3').to_rgba\n",
    "\n",
    "num_colr = 12\n",
    "cmap = get_cmap(num_colr)\n",
    "colr_list = [cmap(float(x)) for x in range(num_colr)]   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare image, predictions and ground truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for i in range(3):\n",
    "    show_results(i, overlap=0.05)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define area of interest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Have to define AoI here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_img(filename):\n",
    "    return np.array(Image.open(filename))\n",
    "\n",
    "def get_tile_images(image, width=224, height=224):\n",
    "    _nrows, _ncols, depth = image.shape\n",
    "    _size = image.size\n",
    "    _strides = image.strides\n",
    "\n",
    "    nrows, _m = divmod(_nrows, height)\n",
    "    ncols, _n = divmod(_ncols, width)\n",
    "    if _m != 0 or _n != 0:\n",
    "        return None\n",
    "\n",
    "    return np.lib.stride_tricks.as_strided(\n",
    "        np.ravel(image),\n",
    "        shape=(nrows, ncols, height, width, depth),\n",
    "        strides=(height * _strides[0], width * _strides[1], *_strides),\n",
    "        writeable=False)\n",
    "\n",
    "def predict_(model,images):\n",
    "    images = V(images)\n",
    "    if not len(images.size()) > 3:\n",
    "        images.data = images.data.view(1,*images.size())\n",
    "    bbox,clas = model(images)\n",
    "    return bbox,clas\n",
    "\n",
    "imagenet_stats = A([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "\n",
    "mean = 255* np.array(imagenet_stats[0], dtype=np.float32)\n",
    "std  = 255* np.array(imagenet_stats[1], dtype=np.float32)\n",
    "\n",
    "norm = lambda x: (x-mean)/ std\n",
    "denorm = lambda x: x * std + mean\n",
    "\n",
    "learn.model.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = get_img(PATH/filename)\n",
    "_nrows = int(img.shape[0] / 224)\n",
    "_ncols = int(img.shape[1] / 224)\n",
    "img = get_tile_images(img)\n",
    "img = img.reshape(64,224,224,3)\n",
    "img_normed = norm(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clas, bbox = predict_(learn.model,img_normed.transpose(0,3,1,2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_results(images,b_clas,b_bb,idx,anchors,ax1):\n",
    "    ax1.set_axis_off()\n",
    "    if not len(images.shape) > 3:\n",
    "        images = images[np.newaxis,:]\n",
    "    ima = denorm(images[idx]).astype(int)\n",
    "    \n",
    "#     bbox,clas = get_y(y[0][idx], y[1][idx])\n",
    "    a_ic = actn_to_bb(b_bb[idx], anchors)\n",
    "    clas_pr, clas_ids = b_clas[idx].max(1)\n",
    "    clas_pr = clas_pr.sigmoid()\n",
    "\n",
    "    conf_scores = b_clas[idx].sigmoid().t().data\n",
    "\n",
    "    out1, out2, cc = [], [], []\n",
    "    for cl in range(0, len(conf_scores)-1):\n",
    "        c_mask = conf_scores[cl] > 0.1\n",
    "        if c_mask.sum() == 0: continue\n",
    "        scores = conf_scores[cl][c_mask]\n",
    "        l_mask = c_mask.unsqueeze(1).expand_as(a_ic)\n",
    "        boxes = a_ic[l_mask].view(-1, 4)\n",
    "        ids, count = nms(boxes.data, scores, 0.1, 50)\n",
    "        ids = ids[:count]\n",
    "        out1.append(scores[ids])\n",
    "        out2.append(boxes.data[ids])\n",
    "        cc.append([cl]*count)\n",
    "#         print(cc)\n",
    "#     print(cc)\n",
    "    if cc == []:\n",
    "        cc = [[0]]\n",
    "    cc = T(np.concatenate(cc))\n",
    "#     print(out1)\n",
    "    if out1 == []:\n",
    "        out1 = [torch.Tensor()]\n",
    "    out1 = torch.cat(out1)\n",
    "    if out2 == []:\n",
    "        out2 = [torch.Tensor()]\n",
    "    out2 = torch.cat(out2)\n",
    "    \n",
    "    torch_gt(ax1, ima, out2, cc, out1, 0.)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Detect buildings in area of interest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(nrows=_nrows, ncols=_ncols, figsize=(24, 24))\n",
    "plt.subplots_adjust(hspace=0.01, wspace=0.01)\n",
    "\n",
    "for idx in range(bbox.size()[0]): \n",
    "    \n",
    "    visualize_results(img_normed, clas, bbox, idx, anchors, ax[idx//8, idx%8])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (fastai)",
   "language": "python",
   "name": "myenv"
  },
  "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}