demo_caffe.py

import numpy as np  
import sys,os  
import cv2
caffe_root = '/home/yaochuanqi/work/tmp/ssd/'
sys.path.insert(0, caffe_root + 'python')  
import caffe  


net_file= 'ssdlite/coco/deploy.prototxt'
caffe_model='ssdlite/deploy.caffemodel'  
test_dir = "images"

caffe.set_mode_cpu()
net = caffe.Net(net_file,caffe_model,caffe.TEST)  

COCO_CLASSES = ("background" , "person" , "bicycle" , "car" , "motorcycle" , 
     "airplane" , "bus" , "train" , "truck" , "boat" , "traffic light",
     "fire hydrant", "N/A" , "stop sign", "parking meter", "bench" , 
     "bird" , "cat" , "dog" , "horse" , "sheep" , "cow" , "elephant" , 
     "bear" , "zebra" , "giraffe" , "N/A" , "backpack" , "umbrella" , 
     "N/A" , "N/A" , "handbag" , "tie" , "suitcase" , "frisbee" , "skis" ,
     "snowboard" , "sports ball", "kite" , "baseball bat", "baseball glove",
     "skateboard" , "surfboard" , "tennis racket", "bottle" , "N/A" ,
     "wine glass", "cup" , "fork" , "knife" , "spoon" , "bowl" , "banana" ,
     "apple" , "sandwich" , "orange" , "broccoli" , "carrot" , "hot dog",
     "pizza" , "donut" , "cake" , "chair" , "couch" , "potted plant", 
     "bed" , "N/A" , "dining table", "N/A" , "N/A" , "toilet" , "N/A" ,
     "tv" , "laptop" , "mouse" , "remote" , "keyboard" , "cell phone",
     "microwave" , "oven" , "toaster" , "sink" , "refrigerator" , "N/A" ,
     "book" , "clock" , "vase" , "scissors" , "teddy bear", "hair drier",
     "toothbrush" )

def preprocess(src):
    img = cv2.resize(src, (300,300))
    img = img - 127.5
    img = img / 127.5
    return img

def postprocess(img, out):   
    h = img.shape[0]
    w = img.shape[1]
    box = out['detection_out'][0,0,:,3:7] * np.array([w, h, w, h])

    cls = out['detection_out'][0,0,:,1]
    conf = out['detection_out'][0,0,:,2]
    return (box.astype(np.int32), conf, cls)

def detect(imgfile):
    origimg = cv2.imread(imgfile)
    img = preprocess(origimg)
    
    img = img.astype(np.float32)
    img = img.transpose((2, 0, 1))

    net.blobs['data'].data[...] = img
    out = net.forward() 
    box, conf, cls = postprocess(origimg, out)

    for i in range(len(box)):
       p1 = (box[i][0], box[i][1])
       p2 = (box[i][2], box[i][3])
       cv2.rectangle(origimg, p1, p2, (0,255,0))
       p3 = (max(p1[0], 15), max(p1[1], 15))
       title = "%s:%.2f" % (COCO_CLASSES[int(cls[i])], conf[i])
       cv2.putText(origimg, title, p3, cv2.FONT_ITALIC, 0.6, (0, 255, 0), 1)
    cv2.imshow("SSD", origimg)
 
    k = cv2.waitKey(0) & 0xff
        #Exit if ESC pressed
    if k == 27 : return False
    return True

for f in os.listdir(test_dir):
    if detect(test_dir + "/" + f) == False:
       break