如果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