nas_seg_v2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Great! Good to go!\n"
     ]
    }
   ],
   "source": [
    "from fastai.vision.all import *\n",
    "from fastai.callback.tracker import SaveModelCallback\n",
    "# from fastprogress.fastprogress import master_bar, progress_bar\n",
    "\n",
    "from gepcore.utils import convolution\n",
    "from gepcore.utils import cell_graph\n",
    "from gepcore.entity import Gene, Chromosome\n",
    "from gepcore.symbol import PrimitiveSet\n",
    "from nas_seg.seg_model import get_net, arch_config, Network\n",
    "from nas_seg.utils import code_to_rgb\n",
    "from nas_seg.isprs_dataset import ISPRSDataset, img_to_mask, mask_to_img\n",
    "from pygraphviz import AGraph\n",
    "import glob\n",
    "\n",
    "#from tqdm import tqdm\n",
    "from skimage import io\n",
    "from sklearn.metrics import confusion_matrix\n",
    "\n",
    "torch.backends.cudnn.benchmark = True\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "  print(\"Great! Good to go!\")\n",
    "else:\n",
    "  print('CUDA is not up!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gepcore.utils import convolution\n",
    "from gepcore.utils import cell_graph\n",
    "from gepcore.entity import Gene, Chromosome, KExpressionGraph\n",
    "from gepcore.symbol import PrimitiveSet\n",
    "#from gepcore.operators import *\n",
    "from nas_seg.seg_model import get_net, arch_config\n",
    "from ptflops import get_model_complexity_info\n",
    "\n",
    "# import fastai deep learning tools\n",
    "from fastai.vision.all import *\n",
    "\n",
    "\n",
    "# imports and stuff\n",
    "import numpy as np\n",
    "from skimage import io\n",
    "from glob import glob\n",
    "from pathlib import Path\n",
    "\n",
    "from tqdm.notebook import tqdm\n",
    "from sklearn.metrics import confusion_matrix\n",
    "import random\n",
    "import itertools\n",
    "# Matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "# Torch imports\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.utils.data as data\n",
    "import torch.optim as optim\n",
    "import torch.optim.lr_scheduler\n",
    "import torch.nn.init\n",
    "from torch.autograd import Variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ISPRS color palette\n",
    "# Let's define the standard ISPRS color palette\n",
    "palette = {0 : (255, 255, 255), # Impervious surfaces (white)\n",
    "           1 : (0, 0, 255),     # Buildings (blue)\n",
    "           2 : (0, 255, 255),   # Low vegetation (cyan)\n",
    "           3 : (0, 255, 0),     # Trees (green)\n",
    "           4 : (255, 255, 0),   # Cars (yellow)\n",
    "           5 : (255, 0, 0),     # Clutter (red)\n",
    "           6 : (0, 0, 0)}       # Undefined (black)\n",
    "\n",
    "invert_palette = {v: k for k, v in palette.items()}\n",
    "\n",
    "def convert_to_color(arr_2d, palette=palette):\n",
    "    \"\"\" Numeric labels to RGB-color encoding \"\"\"\n",
    "    arr_3d = np.zeros((arr_2d.shape[0], arr_2d.shape[1], 3), dtype=np.uint8)\n",
    "\n",
    "    for c, i in palette.items():\n",
    "        m = arr_2d == c\n",
    "        arr_3d[m] = i\n",
    "\n",
    "    return arr_3d\n",
    "\n",
    "def convert_from_color(arr_3d, palette=invert_palette):\n",
    "    \"\"\" RGB-color encoding to grayscale labels \"\"\"\n",
    "    arr_2d = np.zeros((arr_3d.shape[0], arr_3d.shape[1]), dtype=np.uint8)\n",
    "\n",
    "    for c, i in palette.items():\n",
    "        m = np.all(arr_3d == np.array(c).reshape(1, 1, 3), axis=2)\n",
    "        arr_2d[m] = i\n",
    "\n",
    "    return arr_2d\n",
    "\n",
    "# # We load one tile from the dataset and we display it\n",
    "# img = io.imread('/home/cliff/rs_imagery/ISPRS-DATASETS/Vaihingen/top/top_mosaic_09cm_area11.tif')\n",
    "# fig = plt.figure()\n",
    "# fig.add_subplot(121)\n",
    "# plt.imshow(img)\n",
    "\n",
    "# # We load the ground truth\n",
    "# gt = io.imread('/home/cliff/rs_imagery/ISPRS-DATASETS/Vaihingen/gts_for_participants/top_mosaic_09cm_area11.tif')\n",
    "# fig.add_subplot(122)\n",
    "# plt.imshow(gt)\n",
    "# plt.show()\n",
    "\n",
    "# # We also check that we can convert the ground truth into an array format\n",
    "# array_gt = convert_from_color(gt)\n",
    "# print(\"Ground truth in numerical format has shape ({},{}) : \\n\".format(*array_gt.shape[:2]), array_gt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utils\n",
    "\n",
    "def get_random_pos(img, window_shape):\n",
    "    \"\"\" Extract of 2D random patch of shape window_shape in the image \"\"\"\n",
    "    w, h = window_shape\n",
    "    W, H = img.shape[:-1] #img.shape[-2:]\n",
    "    x1 = random.randint(0, W - w - 1)\n",
    "    x2 = x1 + w\n",
    "    y1 = random.randint(0, H - h - 1)\n",
    "    y2 = y1 + h\n",
    "    return x1, x2, y1, y2\n",
    "\n",
    "\n",
    "def accuracy(input, target):\n",
    "    return 100 * float(np.count_nonzero(input == target)) / target.size\n",
    "\n",
    "def sliding_window(top, step=10, window_size=(20,20)):\n",
    "    \"\"\" Slide a window_shape window across the image with a stride of step \"\"\"\n",
    "    for x in range(0, top.shape[0], step):\n",
    "        if x + window_size[0] > top.shape[0]:\n",
    "            x = top.shape[0] - window_size[0]\n",
    "        for y in range(0, top.shape[1], step):\n",
    "            if y + window_size[1] > top.shape[1]:\n",
    "                y = top.shape[1] - window_size[1]\n",
    "            yield x, y, window_size[0], window_size[1]\n",
    "            \n",
    "def count_sliding_window(top, step=10, window_size=(20,20)):\n",
    "    \"\"\" Count the number of windows in an image \"\"\"\n",
    "    c = 0\n",
    "    for x in range(0, top.shape[0], step):\n",
    "        if x + window_size[0] > top.shape[0]:\n",
    "            x = top.shape[0] - window_size[0]\n",
    "        for y in range(0, top.shape[1], step):\n",
    "            if y + window_size[1] > top.shape[1]:\n",
    "                y = top.shape[1] - window_size[1]\n",
    "            c += 1\n",
    "    return c\n",
    "\n",
    "def grouper(n, iterable):\n",
    "    \"\"\" Browse an iterator by chunk of n elements \"\"\"\n",
    "    it = iter(iterable)\n",
    "    while True:\n",
    "        chunk = tuple(itertools.islice(it, n))\n",
    "        if not chunk:\n",
    "            return\n",
    "        yield chunk\n",
    "\n",
    "def metrics(predictions, gts, label_values=LABELS):\n",
    "    cm = confusion_matrix(\n",
    "            gts,\n",
    "            predictions,\n",
    "            range(len(label_values)))\n",
    "    \n",
    "    print(\"Confusion matrix :\")\n",
    "    print(cm)\n",
    "    \n",
    "    print(\"---\")\n",
    "    \n",
    "    # Compute global accuracy\n",
    "    total = sum(sum(cm))\n",
    "    accuracy = sum([cm[x][x] for x in range(len(cm))])\n",
    "    accuracy *= 100 / float(total)\n",
    "    print(\"{} pixels processed\".format(total))\n",
    "    print(\"Total accuracy : {}%\".format(accuracy))\n",
    "    \n",
    "    print(\"---\")\n",
    "    \n",
    "    # Compute F1 score\n",
    "    F1Score = np.zeros(len(label_values))\n",
    "    for i in range(len(label_values)):\n",
    "        try:\n",
    "            F1Score[i] = 2. * cm[i,i] / (np.sum(cm[i,:]) + np.sum(cm[:,i]))\n",
    "        except:\n",
    "            # Ignore exception if there is no element in class i for test set\n",
    "            pass\n",
    "    print(\"F1Score :\")\n",
    "    for l_id, score in enumerate(F1Score):\n",
    "        print(\"{}: {}\".format(label_values[l_id], score))\n",
    "\n",
    "    print(\"---\")\n",
    "        \n",
    "    # Compute kappa coefficient\n",
    "    total = np.sum(cm)\n",
    "    pa = np.trace(cm) / float(total)\n",
    "    pe = np.sum(np.sum(cm, axis=0) * np.sum(cm, axis=1)) / float(total*total)\n",
    "    kappa = (pa - pe) / (1 - pe);\n",
    "    print(\"Kappa: \" + str(kappa))\n",
    "    return accuracy\n",
    "\n",
    "# IoU = TP / (TP + FP + FN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "   \n",
    "data_path = Path('/home/cliff/rs_imagery/ISPRS-DATASETS/Vaihingen/vaihingen_256')\n",
    "img_path = data_path/'images/train'\n",
    "msk_path = data_path/'masks/train'\n",
    "\n",
    "mask_labels = [\"roads\", \"buildings\", \"low veg.\", \"trees\", \"cars\", \"clutter\"] \n",
    "num_classes = len(mask_labels) \n",
    "print(img_path, '\\n', msk_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pygraphviz import AGraph\n",
    "#import glob\n",
    "\n",
    "graph = [AGraph(g) for g in glob('../graphs/*.dot')]\n",
    "_, comp_graphs = cell_graph.generate_comp_graph(graph)\n",
    "\n",
    "conf = arch_config(comp_graphs=comp_graphs, channels=64, classes=N_CLASSES)\n",
    "net = get_net(conf)\n",
    "net.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision.transforms as transforms\n",
    "\n",
    "tfms = transforms.Compose([transforms.ToTensor(),\n",
    "                           transforms.Normalize([0.4776, 0.3226, 0.3189], [0.1816, 0.1224, 0.1185])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Load the datasets\n",
    "if DATASET == 'Potsdam':\n",
    "    all_files = sorted(glob(LABEL_FOLDER.replace('{}', '*')))\n",
    "    all_ids = [\"\".join(f.split('')[5:7]) for f in all_files]\n",
    "elif DATASET == 'Vaihingen':\n",
    "    #all_ids = \n",
    "    all_files = sorted(glob(LABEL_FOLDER.replace('{}', '*')))\n",
    "    all_ids = [f.split('area')[-1].split('.')[0] for f in all_files]\n",
    "\n",
    "# Test split on Vaihingen :\n",
    "test_ids = ['5', '7', '23', '30']\n",
    "\n",
    "train_set = ISPRS_dataset(data_dir=img_path, label_dir=msk_path, transforms=tfms)\n",
    "train_loader = torch.utils.data.DataLoader(train_set, shuffle=True, batch_size=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parameters\n",
    "window_size = 256\n",
    "\n",
    "msk_labels = np.array([\"roads\", \"buildings\", \"low veg.\", \"trees\", \"cars\", \"clutter\"])\n",
    "num_classes = len(msk_labels) \n",
    "\n",
    "\n",
    "dataset = 'Vaihingen' #'Potsdam'\n",
    "dataset_dir = Path.home()/'rs_imagery/ISPRS_DATASETS/{}'.format(dataset)\n",
    "\n",
    "if dataset == 'Potsdam':\n",
    "    tiles = dataset_dir/'Ortho_IRRG/top_potsdam_{}_{}_IRRG.tif'\n",
    "    masks = dataset_dir/'Labels_for_participants/top_potsdam_{}_{}_label.tif'\n",
    "    e_masks = dataset_dir/'Labels_for_participants_no_Boundary/top_potsdam_{}_label_noBoundary.tif'\n",
    "    trainset_dir = dataset.lower() + '_{}'.format(window_size) \n",
    "    testset_ids = ['2_11', '2_12', '4_10', '5_11', '6_7', '7_8', '7_10'] # ['7_8', '4_10', 2 11, 5 11]\n",
    "elif dataset == 'Vaihingen':\n",
    "    tiles = dataset_dir/'top/top_mosaic_09cm_area{}.tif'\n",
    "    masks = dataset_dir/'gts_for_participants/top_mosaic_09cm_area{}.tif'\n",
    "    e_masks = dataset_dir/'gts_eroded_for_participants/top_mosaic_09cm_area{}_noBoundary.tif'\n",
    "    trainset_dir = dataset.lower() + '_{}'.format(window_size) \n",
    "    testset_ids = ['5', '7', '23', '30']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parameters\n",
    "WINDOW_SIZE = (128, 128) # Patch size\n",
    "STRIDE = 32 # Stride for testing\n",
    "IN_CHANNELS = 3 # Number of input channels (e.g. RGB)\n",
    "FOLDER = \"/home/cliff/rs_imagery/ISPRS-DATASETS/\" # Replace with your \"/path/to/the/ISPRS/dataset/folder/\"\n",
    "BATCH_SIZE = 10 # Number of samples in a mini-batch\n",
    "\n",
    "LABELS = [\"roads\", \"buildings\", \"low veg.\", \"trees\", \"cars\", \"clutter\"] # Label names\n",
    "N_CLASSES = len(LABELS) # Number of classes\n",
    "WEIGHTS = torch.ones(N_CLASSES) # Weights for class balancing\n",
    "CACHE = True # Store the dataset in-memory\n",
    "\n",
    "DATASET = 'Vaihingen'\n",
    "\n",
    "if DATASET == 'Potsdam':\n",
    "    MAIN_FOLDER = FOLDER + 'Potsdam/'\n",
    "    # Uncomment the next line for IRRG data\n",
    "    # DATA_FOLDER = MAIN_FOLDER + '3_Ortho_IRRG/top_potsdam_{}_IRRG.tif'\n",
    "    # For RGB data\n",
    "    DATA_FOLDER = MAIN_FOLDER + 'Ortho_RGB/top_potsdam_{}_RGB.tif'\n",
    "    LABEL_FOLDER = MAIN_FOLDER + 'Labels_for_participants/top_potsdam_{}_label.tif'\n",
    "    ERODED_FOLDER = MAIN_FOLDER + 'Labels_for_participants_no_Boundary/top_potsdam_{}_label_noBoundary.tif'    \n",
    "elif DATASET == 'Vaihingen':\n",
    "    MAIN_FOLDER = FOLDER + 'Vaihingen/'\n",
    "    DATA_FOLDER = MAIN_FOLDER + 'top/top_mosaic_09cm_area{}.tif'\n",
    "    LABEL_FOLDER = MAIN_FOLDER + 'gts_for_participants/top_mosaic_09cm_area{}.tif'\n",
    "    ERODED_FOLDER = MAIN_FOLDER + 'gts_eroded_for_participants/top_mosaic_09cm_area{}_noBoundary.tif'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "\n",
    "def test(net, test_ids, all=False, stride=WINDOW_SIZE[0], batch_size=BATCH_SIZE, window_size=WINDOW_SIZE):\n",
    "    # Use the network on the test set\n",
    "    test_images = (np.asarray(io.imread(DATA_FOLDER.format(id))) for id in test_ids)\n",
    "    test_labels = (np.asarray(io.imread(LABEL_FOLDER.format(id)), dtype='uint8') for id in test_ids)\n",
    "    eroded_labels = (convert_from_color(io.imread(ERODED_FOLDER.format(id))) for id in test_ids)\n",
    "    all_preds = []\n",
    "    all_gts = []\n",
    "    \n",
    "    # Switch the network to inference mode\n",
    "    net.eval()\n",
    "\n",
    "    for img, gt, gt_e in tqdm(zip(test_images, test_labels, eroded_labels), total=len(test_ids), leave=False):\n",
    "        pred = np.zeros(img.shape[:2] + (N_CLASSES,))\n",
    "\n",
    "        total = count_sliding_window(img, step=stride, window_size=window_size) // batch_size\n",
    "        for i, coords in enumerate(tqdm(grouper(batch_size, \n",
    "                                                sliding_window(img, step=stride, window_size=window_size)), \n",
    "                                                total=total, leave=False)):\n",
    "            # Display in progress results\n",
    "            if i > 0 and total > 10 and i % int(10 * total / 100) == 0:\n",
    "                    _pred = np.argmax(pred, axis=-1)\n",
    "                    fig = plt.figure()\n",
    "                    fig.add_subplot(1,3,1)\n",
    "                    plt.imshow(np.asarray(img, dtype='uint8'))\n",
    "                    fig.add_subplot(1,3,2)\n",
    "                    plt.imshow(convert_to_color(_pred))\n",
    "                    fig.add_subplot(1,3,3)\n",
    "                    plt.imshow(gt)\n",
    "                    clear_output()\n",
    "                    plt.show()\n",
    "                    \n",
    "            # Build the tensor\n",
    "            #image_patches = [np.copy(img[x:x+w, y:y+h]).transpose((2,0,1)) for x,y,w,h in coords]\n",
    "            #image_patches = np.asarray(image_patches)\n",
    "            #image_patches = Variable(torch.from_numpy(image_patches).cuda(), volatile=True)\n",
    "\n",
    "            image_patches = [torch.clone(tf.normalize(tf.to_tensor(img[x:x+w, y:y+h]), \n",
    "                                        [0.4776, 0.3226, 0.3189], [0.1816, 0.1224, 0.1185]))\n",
    "                             for x,y,w,h in coords]\n",
    "            image_patches = torch.stack(image_patches).cuda()\n",
    "            #print(image_patches)\n",
    "            \n",
    "            #image_patches = Variable(torch.from_numpy(image_patches).cuda(), volatile=True)\n",
    "            \n",
    "            # Do the inference\n",
    "            outs = net(image_patches)\n",
    "            outs = outs.data.cpu().numpy()\n",
    "            \n",
    "            # Fill in the results array\n",
    "            for out, (x, y, w, h) in zip(outs, coords):\n",
    "                out = out.transpose((1,2,0))\n",
    "                pred[x:x+w, y:y+h] += out\n",
    "            del(outs)\n",
    "\n",
    "        pred = np.argmax(pred, axis=-1)\n",
    "\n",
    "        # Display the result\n",
    "        clear_output()\n",
    "        fig = plt.figure()\n",
    "        fig.add_subplot(1,3,1)\n",
    "        plt.imshow(np.asarray(img, dtype='uint8'))\n",
    "        fig.add_subplot(1,3,2)\n",
    "        plt.imshow(convert_to_color(pred))\n",
    "        fig.add_subplot(1,3,3)\n",
    "        plt.imshow(gt)\n",
    "        plt.show()\n",
    "\n",
    "        all_preds.append(pred)\n",
    "        all_gts.append(gt_e)\n",
    "\n",
    "        clear_output()\n",
    "        # Compute some metrics\n",
    "        metrics(pred.ravel(), gt_e.ravel())\n",
    "        accuracy = metrics(np.concatenate([p.ravel() for p in all_preds]), \n",
    "                           np.concatenate([p.ravel() for p in all_gts]).ravel())\n",
    "    if all:\n",
    "        return accuracy, all_preds, all_gts\n",
    "    else:\n",
    "        return accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "net.load_state_dict(torch.load('./model.pth'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, all_preds, all_gts = test(net, test_ids, all=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for p, id_ in zip(all_preds, test_ids):\n",
    "    img = convert_to_color(p)\n",
    "    plt.imshow(img) and plt.show()\n",
    "    io.imsave('./inference_tile_{}.png'.format(id_), img)"
   ]
  },
  {
   "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}