segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Imports\n",
    "from scipy.misc import imsave\n",
    "import os\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from IPython import embed\n",
    "from model import get_frontend, add_softmax, add_context\n",
    "from utils import interp_map, pascal_palette\n",
    "from keras.preprocessing.image import load_img, img_to_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#CHANGE THIS TO YOUR GIT REPO LOCATION\n",
    "INPUT_DIR='/home/bfortuner/workplace/deephacks/'\n",
    "\n",
    "IMG_EXTEN='.jpg'\n",
    "STYLE_FILENAME='starrynight.jpg'\n",
    "TEST_PHOTO_FILENAME='cat.jpg' #'good_dog_img_for_testing.jpg'\n",
    "\n",
    "OUTPUT_DIR='images/segmentations/'\n",
    "IMAGES_DIR='images/'\n",
    "INPUT_FILE=IMAGES_DIR+TEST_PHOTO_FILENAME\n",
    "OUTPUT_FILE=OUTPUT_DIR+TEST_PHOTO_FILENAME+'_seg.png'\n",
    "WEIGHTS_PATH=INPUT_DIR+'conversion/dilation8_pascal_voc.npy'\n",
    "ZOOM=8\n",
    "MEAN=[102.93, 111.36, 116.52]\n",
    "PASCAL_PALETTE = {\n",
    "    0: (0, 0, 0),\n",
    "    1: (128, 0, 0),\n",
    "    2: (0, 128, 0),\n",
    "    3: (128, 128, 0),\n",
    "    4: (0, 0, 128),\n",
    "    5: (128, 0, 128),\n",
    "    6: (0, 128, 128),\n",
    "    7: (128, 128, 128),\n",
    "    8: (64, 0, 0),\n",
    "    9: (192, 0, 0),\n",
    "    10: (64, 128, 0),\n",
    "    11: (192, 128, 0),\n",
    "    12: (64, 0, 128),\n",
    "    13: (192, 0, 128),\n",
    "    14: (64, 128, 128),\n",
    "    15: (192, 128, 128),\n",
    "    16: (0, 64, 0),\n",
    "    17: (128, 64, 0),\n",
    "    18: (0, 192, 0),\n",
    "    19: (128, 192, 0),\n",
    "    20: (0, 64, 128),\n",
    "}\n",
    "'''\n",
    "categories={\n",
    "        'aeroplane'; %1\n",
    "        'bicycle'; %2\n",
    "        'bird'; %3\n",
    "        'boat'; %4\n",
    "        'bottle'; %5\n",
    "        'bus'; %6\n",
    "        'car'; %7\n",
    "        'cat'; %8\n",
    "        'chair'; %9\n",
    "        'cow'; %10\n",
    "        'diningtable';%11\n",
    "        'dog';%12\n",
    "        'horse';%13\n",
    "        'motorbike';%14\n",
    "        'person';%15\n",
    "        'pottedplant'; %16\n",
    "        'sheep'; %17\n",
    "        'sofa'; %18\n",
    "        'train'; %19\n",
    "        'tvmonitor'; %20\n",
    "};\n",
    "'''\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Settings for the Pascal dataset\n",
    "input_width, input_height = 900, 900\n",
    "label_margin = 186\n",
    "has_context_module = False\n",
    "\n",
    "def get_trained_model():\n",
    "    \"\"\" Returns a model with loaded weights. \"\"\"\n",
    "\n",
    "    model = get_frontend(input_width, input_height)\n",
    "\n",
    "    if has_context_module:\n",
    "        model = add_context(model)\n",
    "\n",
    "    model = add_softmax(model)\n",
    "\n",
    "    def load_tf_weights():\n",
    "        \"\"\" Load pretrained weights converted from Caffe to TF. \"\"\"\n",
    "\n",
    "        # 'latin1' enables loading .npy files created with python2\n",
    "        weights_data = np.load(WEIGHTS_PATH, encoding='latin1').item()\n",
    "\n",
    "        for layer in model.layers:\n",
    "            if layer.name in weights_data.keys():\n",
    "                layer_weights = weights_data[layer.name]\n",
    "                layer.set_weights((layer_weights['weights'],\n",
    "                                   layer_weights['biases']))\n",
    "\n",
    "    def load_keras_weights():\n",
    "        \"\"\" Load a Keras checkpoint. \"\"\"\n",
    "        model.load_weights(WEIGHTS_PATH)\n",
    "\n",
    "    if WEIGHTS_PATH.endswith('.npy'):\n",
    "        load_tf_weights()\n",
    "    elif WEIGHTS_PATH.endswith('.hdf5'):\n",
    "        load_keras_weights()\n",
    "    else:\n",
    "        raise Exception(\"Unknown weights format.\")\n",
    "\n",
    "    return model\n",
    "\n",
    "\n",
    "def forward_pass():\n",
    "    ''' Runs a forward pass to segment the image. '''\n",
    "\n",
    "    model = get_trained_model()\n",
    "\n",
    "    # Load image and swap RGB -> BGR to match the trained weights\n",
    "    image_rgb = np.array(Image.open(INPUT_FILE)).astype(np.float32)\n",
    "    image = image_rgb[:, :, ::-1] - MEAN\n",
    "    image_size = image.shape\n",
    "\n",
    "    # Network input shape (batch_size=1)\n",
    "    net_in = np.zeros((1, input_height, input_width, 3), dtype=np.float32)\n",
    "\n",
    "    output_height = input_height - 2 * label_margin\n",
    "    output_width = input_width - 2 * label_margin\n",
    "\n",
    "    # This simplified prediction code is correct only if the output\n",
    "    # size is large enough to cover the input without tiling\n",
    "    assert image_size[0] < output_height\n",
    "    assert image_size[1] < output_width\n",
    "\n",
    "    # Center pad the original image by label_margin.\n",
    "    # This initial pad adds the context required for the prediction\n",
    "    # according to the preprocessing during training.\n",
    "    image = np.pad(image,\n",
    "                   ((label_margin, label_margin),\n",
    "                    (label_margin, label_margin),\n",
    "                    (0, 0)), 'reflect')\n",
    "\n",
    "    # Add the remaining margin to fill the network input width. This\n",
    "    # time the image is aligned to the upper left corner though.\n",
    "    margins_h = (0, input_height - image.shape[0])\n",
    "    margins_w = (0, input_width - image.shape[1])\n",
    "    image = np.pad(image,\n",
    "                   (margins_h,\n",
    "                    margins_w,\n",
    "                    (0, 0)), 'reflect')\n",
    "\n",
    "    # Run inference\n",
    "    net_in[0] = image\n",
    "    prob = model.predict(net_in)[0]\n",
    "\n",
    "    # Reshape to 2d here since the networks outputs a flat array per channel\n",
    "    prob_edge = np.sqrt(prob.shape[0]).astype(np.int)\n",
    "    prob = prob.reshape((prob_edge, prob_edge, 21))\n",
    "\n",
    "    # Upsample\n",
    "    if ZOOM > 1:\n",
    "        prob = interp_map(prob, ZOOM, image_size[1], image_size[0])\n",
    "\n",
    "    # Recover the most likely prediction (actual segment class)\n",
    "    prediction = np.argmax(prob, axis=2)\n",
    "    # Apply the color palette to the segmented image\n",
    "    color_image = np.array(pascal_palette)[prediction.ravel()].reshape(\n",
    "        prediction.shape + (3,))\n",
    "    print('Saving results to: ', OUTPUT_FILE)\n",
    "    with open(OUTPUT_FILE, 'wb') as out_file:\n",
    "        Image.fromarray(color_image).save(out_file)\n",
    "    return prediction, color_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "pred,color_image = forward_pass()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "TARGET_CATEGORY=15 #human\n",
    "TARGET_CATEGORY=12 #dog\n",
    "TARGET_CATEGORY=8 #cat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#Test some things\n",
    "print(np.max(pred))\n",
    "print(np.min(pred))\n",
    "print(np.unique(pred))\n",
    "print(pred.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#Reshape np array back to image format with 1 channel\n",
    "out = pred.reshape(pred.shape + (1,))\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#Change all pixels != our target category to 0 (black)\n",
    "pred[pred != TARGET_CATEGORY] = 0\n",
    "\n",
    "#Change all pixels == our target category to 1 (white)\n",
    "pred[pred == TARGET_CATEGORY] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#Save image locally\n",
    "imsave(OUTPUT_FILE, pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#View segmentation image\n",
    "segimg = load_img(OUTPUT_FILE)\n",
    "segimg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}