demo_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os, sys, cv2\n",
    "import argparse\n",
    "import os.path as osp\n",
    "import glob\n",
    "from easydict import EasyDict as edict\n",
    "\n",
    "\n",
    "from lib.networks.factory import get_network\n",
    "from lib.fast_rcnn.config import cfg\n",
    "from lib.fast_rcnn.test import im_detect\n",
    "from lib.fast_rcnn.nms_wrapper import nms_wrapper\n",
    "from lib.utils.timer import Timer\n",
    "\n",
    "CLASSES = ('__background__',\n",
    "           'aeroplane', 'bicycle', 'bird', 'boat',\n",
    "           'bottle', 'bus', 'car', 'cat', 'chair',\n",
    "           'cow', 'diningtable', 'dog', 'horse',\n",
    "           'motorbike', 'person', 'pottedplant',\n",
    "           'sheep', 'sofa', 'train', 'tvmonitor')\n",
    "\n",
    "\n",
    "# CLASSES = ('__background__','person','bike','motorbike','car','bus')\n",
    "\n",
    "def vis_detections(im, class_name, dets, ax, thresh=0.5):\n",
    "    \"\"\"Draw detected bounding boxes.\"\"\"\n",
    "    print(dets)\n",
    "    inds = np.where(dets[:, -1] >= thresh)[0]\n",
    "    if len(inds) == 0:\n",
    "        return\n",
    "\n",
    "    for i in inds:\n",
    "        bbox = dets[i, :4]\n",
    "        score = dets[i, -1]\n",
    "\n",
    "        ax.add_patch(\n",
    "            plt.Rectangle((bbox[0], bbox[1]),\n",
    "                          bbox[2] - bbox[0],\n",
    "                          bbox[3] - bbox[1], fill=False,\n",
    "                          edgecolor='red', linewidth=3.5)\n",
    "        )\n",
    "        ax.text(bbox[0], bbox[1] - 2,\n",
    "                '{:s} {:.3f}'.format(class_name, score),\n",
    "                bbox=dict(facecolor='blue', alpha=0.5),\n",
    "                fontsize=14, color='white')\n",
    "\n",
    "    ax.set_title(('{} detections with '\n",
    "                  'p({} | box) >= {:.1f}').format(class_name, class_name,\n",
    "                                                  thresh),\n",
    "                 fontsize=14)\n",
    "    plt.axis('off')\n",
    "    plt.tight_layout()\n",
    "    plt.draw()\n",
    "\n",
    "\n",
    "def demo(sess, net, image_name):\n",
    "    \"\"\"Detect object classes in an image using pre-computed object proposals.\"\"\"\n",
    "\n",
    "    # Load the demo image\n",
    "    im = cv2.imread(image_name)\n",
    "\n",
    "    # Detect all object classes and regress object bounds\n",
    "    timer = Timer()\n",
    "    timer.tic()\n",
    "    scores, boxes = im_detect(sess, net, im)\n",
    "\n",
    "    timer.toc()\n",
    "    print(('Detection took {:.3f}s for '\n",
    "           '{:d} object proposals').format(timer.total_time, boxes.shape[0]))\n",
    "\n",
    "    # Visualize detections for each class\n",
    "    im = im[:, :, (2, 1, 0)]\n",
    "    fig, ax = plt.subplots(figsize=(12, 12))\n",
    "    ax.imshow(im, aspect='equal')\n",
    "\n",
    "    CONF_THRESH = 0.7\n",
    "    NMS_THRESH = 0.3\n",
    "    res = nms_wrapper(scores, boxes, threshold=0.7)\n",
    "#     print(res)\n",
    "#     for cls_ind, cls in enumerate(CLASSES[1:]):\n",
    "#         cls_ind += 1  # because we skipped background\n",
    "#         cls_boxes = boxes[:, 4 * cls_ind:4 * (cls_ind + 1)]\n",
    "#         cls_scores = scores[:, cls_ind]\n",
    "#         dets = np.hstack((cls_boxes,\n",
    "#                           cls_scores[:, np.newaxis])).astype(np.float32)\n",
    "#         keep = nms(dets, NMS_THRESH)\n",
    "#         dets = dets[keep, :]\n",
    "#         vis_detections(im, cls, dets, ax, thresh=CONF_THRESH)\n",
    "    for ind, r in enumerate(res):\n",
    "        if r['dets'] is None: continue\n",
    "        dets = r['dets']\n",
    "        for i in range(0, dets.shape[0]):\n",
    "            vis_detections(im, r['class'], np.expand_dims(dets[i, :], 0), ax, thresh=CONF_THRESH)\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    cfg.TEST.HAS_RPN = True  # Use RPN for proposals\n",
    "    args = edict()\n",
    "    args.gpu_id=0\n",
    "    args.demo_net=\"Resnet50_test\"\n",
    "    args.model=\"./output/faster_rcnn_end2end_resnet_voc/voc_2007_trainval\"\n",
    "#     args = parse_args()\n",
    "\n",
    "    if args.model == ' ' or not os.path.exists(args.model):\n",
    "        print(('current path is ' + os.path.abspath(__file__)))\n",
    "        raise IOError(('Error: Model not found.\\n'))\n",
    "    # load network\n",
    "    device_name = '/gpu:{:d}'.format(args.gpu_id)\n",
    "    print(device_name)\n",
    "    with tf.device(device_name):\n",
    "        net = get_network(args.demo_net)\n",
    "    saver = tf.train.Saver()\n",
    "    # init session\n",
    "    c = tf.ConfigProto(allow_soft_placement=True)\n",
    "    c.gpu_options.visible_device_list=str(args.gpu_id)\n",
    "    sess = tf.Session(config=c)\n",
    "\n",
    "    # load model\n",
    "    print(('Loading network {:s}... '.format(args.demo_net)), end=' ')\n",
    "    ckpt = tf.train.latest_checkpoint(args.model)\n",
    "    if ckpt:\n",
    "        # the global_step will restore sa well\n",
    "        saver.restore(sess,ckpt)\n",
    "        print('restore from the checkpoint{0}'.format(ckpt))\n",
    "    #saver.restore(sess, args.model)\n",
    "    print (' done.')\n",
    "\n",
    "    # Warmup on a dummy image\n",
    "    im = 128 * np.ones((300, 300, 3), dtype=np.uint8)\n",
    "    for i in range(2):\n",
    "        _, _ = im_detect(sess, net, im)\n",
    "\n",
    "    im_names = glob.glob(os.path.join(cfg.DATA_DIR, 'demo', '*.png')) + \\\n",
    "               glob.glob(os.path.join(cfg.DATA_DIR, 'demo', '*.jpg'))\n",
    "#     im_names = [\"/home/antonio/tf_deformable_net/data/VOCdevkit2007/VOC2007/JPEGImages/00\"+name for name \\\n",
    "#                             in [\"7097.jpg\", \"8997.jpg\"]]\n",
    "    for im_name in im_names:\n",
    "        print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')\n",
    "        print('Demo for {:s}'.format(im_name))\n",
    "        demo(sess, net, im_name)\n",
    "\n",
    "    plt.show()\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [default]",
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}