用python实现caffe上ssd的检测结果输出

如果deploy.prototxt没有写错的话，可以跑出相应的检测结果（坐标和分类）
但目前此代码只能输出top1

import numpy as np  
import sys,os  
import cv2
caffe_root = '/work/caffessd'     #caffe目录
sys.path.insert(0, caffe_root + 'python')  
import caffe  
import time

test_txt = 'imglist.txt'   #测试图片txt和目录
jpg_fold = './image/'

net_file= 'deploy.prototxt'  #模型描述
caffe_model='deploy.caffemodel'  #训练出现的模型

save_txt = 'result.txt'  #检测结果保存的txt

vis = 0
cacu_time_forward = 0
tic = time.time()
if not os.path.exists(caffe_model):
    print("deploy.caffemodel does not exist,")
    exit()
#caffe.set_device(2)
#caffe.set_mode_gpu()
caffe.set_mode_cpu()
net = caffe.Net(net_file,caffe_model,caffe.TEST)  

CLASSES = ('background',
           'gray')                             #物体分类数和名称


def preprocess(src):
    img = cv2.resize(src, (300,300))       #数据归一化参数
    img = img - 127.5
    img = img * 0.007843
    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):
    global cacu_time_forward
    origimg = cv2.imread(imgfile)
    img = preprocess(origimg)
    
    img = img.astype(np.float32)
    img = img.transpose((2, 0, 1))

    net.blobs['data'].data[...] = img

    tic_forward = time.time()
    out = net.forward()
    toc_forward = time.time()
    cacu_time_forward += toc_forward - tic_forward

    box, conf, cls = postprocess(origimg, out)
    if vis == 1:

        p1 = (box[0][0],box[0][1])
        p2 = (box[0][2],box[0][3])
        cv2.rectangle(origimg, p1, p2,(0,255,0))
        p3 = (max(p1[0],15),max(p1[1],15))
        title = "%s:%.2f" % (CLASSES[int(cls)], conf)
        cv2.putText(origimg, title, p3, cv2.FONT_ITALIC, 0.6, (0, 255, 0), 1)
        cv2.imshow("ssd",origimg)
        k = cv2.waitKey(0) & 0xff
        if k == 27: return False

    return box, conf, CLASSES[int(cls)]
   

with open(test_txt) as f:
    lines = f.readlines()

with open(save_txt,"a") as f_write:
    for item in range(len(lines)):
        if item%1000 == 0:
            toc = time.time()
            print toc - tic
            tic = time.time()
            print '{}\\{}'.format(str(item),len(lines))
        path_jpg = jpg_fold + str(lines[item].strip()) + '.jpg'
        box,conf,cls = detect(path_jpg)     
        #print box[0][1]
        #print conf[0]
        #print cls
        f_write.write('{} {} {} {} {} {}\n'.format(lines[item].strip(), str(box[0][0]),str(box[0][1]),str(box[0][2]),str(box[0][3]),cls))
    print 'forward time:'
    print cacu_time_forward/1000

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