demo

tools/demo.py

+#!/usr/bin/env python
+
+# --------------------------------------------------------
+# Fast R-CNN
+# Copyright (c) 2015 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ross Girshick
+# --------------------------------------------------------
+
+import _init_paths
+from fast_rcnn.config import cfg
+from fast_rcnn.test import im_detect
+from utils.cython_nms import nms
+from utils.timer import Timer
+import matplotlib.pyplot as plt
+import numpy as np
+import caffe, cPickle, os, cv2
+
+CLASSES = ('__background__',
+           'aeroplane', 'bicycle', 'bird', 'boat',
+           'bottle', 'bus', 'car', 'cat', 'chair',
+           'cow', 'diningtable', 'dog', 'horse',
+           'motorbike', 'person', 'pottedplant',
+           'sheep', 'sofa', 'train', 'tvmonitor')
+
+def vis_detections(im, class_name, dets, thresh=0.5):
+    """Draw detected bounding boxes."""
+    inds = np.where(dets[:, -1] >= thresh)[0]
+    if len(inds) == 0:
+        return
+
+    im = im[:, :, (2, 1, 0)]
+    fig, ax = plt.subplots(figsize=(14, 14))
+    ax.imshow(im, aspect='equal')
+    for i in inds:
+        bbox = dets[i, :4]
+        score = dets[i, -1]
+
+        ax.add_patch(
+            plt.Rectangle((bbox[0], bbox[1]),
+                          bbox[2] - bbox[0],
+                          bbox[3] - bbox[1], fill=False,
+                          edgecolor='red', linewidth=3.5)
+            )
+        ax.text(bbox[0], bbox[1] - 2,
+                '{:s} {:.3f}'.format(class_name, score),
+                bbox=dict(facecolor='blue', alpha=0.5),
+                fontsize=14, color='white')
+
+    ax.set_title(('All {} detections with '
+                  'score >= {:.1f}').format(class_name, thresh),
+                  fontsize=18)
+    plt.axis('off')
+    plt.tight_layout()
+    plt.show()
+
+def demo(net, image_name, classes):
+    """Detect object classes in an image using pre-computed object proposals."""
+
+    # Load pre-computed Selected Search object proposals
+    box_file = os.path.join(cfg.ROOT_DIR, 'data', 'demo',
+                            image_name + '_boxes.pkl')
+    with open(box_file, 'rb') as f:
+        obj_proposals = cPickle.load(f)
+
+    # Load the demo image
+    im_file = os.path.join(cfg.ROOT_DIR, 'data', 'demo', image_name + '.jpg')
+    im = cv2.imread(im_file)
+
+    # Detect all object classes and regress object bounds
+    timer = Timer()
+    timer.tic()
+    scores, boxes = im_detect(net, im, obj_proposals)
+    timer.toc()
+    print ('Detection took {:.3f}s for '
+           '{:d} object proposals').format(timer.total_time, boxes.shape[0])
+
+    # Visualize detections for each class
+    CONF_THRESH = 0.8
+    NMS_THRESH = 0.3
+    for cls in classes:
+        cls_ind = CLASSES.index(cls)
+        cls_boxes = boxes[:, 4*cls_ind:4*(cls_ind + 1)]
+        cls_scores = scores[:, cls_ind]
+        dets = np.hstack((cls_boxes,
+                          cls_scores[:, np.newaxis])).astype(np.float32)
+        keep = nms(dets, NMS_THRESH)
+        dets = dets[keep, :]
+        print 'All {} detections with score >= {:.1f}'.format(cls, CONF_THRESH)
+        print 'Close image window (ctrl-w) to continue'
+        vis_detections(im, cls, dets, thresh=CONF_THRESH)
+
+if __name__ == '__main__':
+    gpu_id = 0
+    prototxt = 'models/VGG16/test.prototxt'
+    caffemodel = ('data/fast_rcnn_models/'
+                  'vgg16_fast_rcnn_iter_40000.caffemodel')
+    # prototxt = 'models/VGG_CNN_M_1024/test.prototxt'
+    # caffemodel = ('data/fast_rcnn_models/'
+    #               'vgg_cnn_m_1024_fast_rcnn_iter_40000.caffemodel')
+
+    if not os.path.isfile(caffemodel):
+        raise IOError(('{:s} not found.\nDid you run ./data/script/'
+                       'fetch_fast_rcnn_models.sh?').format(caffemodel))
+
+    caffe.set_mode_gpu()
+    caffe.set_device(gpu_id)
+    net = caffe.Net(prototxt, caffemodel, caffe.TEST)
+
+    print '\n\nLoaded network {:s}'.format(caffemodel)
+
+    print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
+    print 'Demo for data/demo/000004.jpg'
+    demo(net, '000004', ('car',))
+
+    print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
+    print 'Demo for data/demo/001551.jpg'
+    demo(net, '001551', ('sofa', 'tvmonitor'))