webcam_inference.py

#!/opt/intel/intelpython2/bin/python
import numpy as np
import tensorflow as tf
from yolo_model import YOLO
import cv2
from matplotlib import pyplot as plt
import time
import os, sys
import argparse

parser = argparse.ArgumentParser(description='Sample program for YOLO inference in TensroFlow')

parser.add_argument('--v2',action='store_true', default=False, help='Type of the yolo mode Tiny or V2')
parser.add_argument('--par',action='store_true', default=False, help='Enable parallel session with INTER/INTRA TensorFlow threads')
parser.add_argument('--image',action='store', default='sample/dog.jpg', help='Select image for object detection')

args = parser.parse_args()

#select the checkpoint
if not args.v2:
    path="utils/ckpt/tiny_yolo"
    _type = 'TINY_VOC'
else:
    path="utils/ckpt/yolov2"
    _type = 'V2_VOC'



#Helper function to draw boxes deduced from the feature map
def draw_boxes(img,box_preds):

    batch_addr = box_preds['batch_addr']
    boxes =box_preds['boxes']
    indices = box_preds['indices']
    class_names = box_preds['class_names']

    boxes = [boxes[i] for i in indices]
    class_names = [class_names[i] for i in indices]
    for i,b in enumerate(boxes):
        idx  = batch_addr[i]
        left = int(max(0,b[0]))
        bot  = int(max(0,b[1]))
        right= int(min(415,b[2]))
        top  = int(min(415,b[3]))
        cv2.rectangle(img[idx],(left,bot),(right,top),(0,255,255),2)
        cv2.putText(img[idx], class_names[i], (int(left), int(bot)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2);


#TensorFlow graph and Session
with tf.Graph().as_default():

    batch = 1
    img_in = tf.placeholder(tf.float32,[None,416,416,3])
    clf = YOLO(img_in,yolotype=_type)
    saver = tf.train.Saver()

    #read and preprocess the image
    cam = cv2.VideoCapture(0)

    #select the session Type
    if not args.par:
        sess_type = tf.Session()
    else:
        sess_type = tf.Session(config=tf.ConfigProto(inter_op_parallelism_threads=int(os.environ['NUM_INTER_THREADS']),
                                                    intra_op_parallelism_threads=int(os.environ['NUM_INTRA_THREADS'])))

    with sess_type as sess:
        saver.restore(sess, path)
        while(True):
            t0 = time.time()
            ret,img=cam.read()
            assert ret,'error reading webcam'
            img2=[cv2.resize(img,(416,416))]*batch
            image = [im*0.003921569 for im in img2]
            box_preds=sess.run(clf.preds,{img_in: image})
            t1 = time.time()
            ftime = 1/(t1-t0)
            print("Frame rate : %f fps"  % ftime)
            draw_boxes(img2,box_preds)

            cv2.imshow('frame',img2[0])
            if cv2.waitKey(1) == 27:
                break
        sess.close()