caffe-python测试代码块

#coding=utf-8
import cv2
import os
caffe_root = '/disk3/python/caffe_s3fd-ssd'
import sys
os.chdir(caffe_root)
sys.path.insert(0, 'python')
import caffe
import time
caffe.set_device(0)
caffe.set_mode_gpu()
model_weights = '/disk3/face_detect/beijing/sfd_head_train/test/models/new/_iter_117000.caffemodel'
model_def = '/disk3/face_detect/beijing/sfd_head_train/test/models/new/deploy_vgg_half.prototxt'
net = caffe.Net(model_def, model_weights, caffe.TEST)
for image_name in os.listdir("/disk3/dataset/head_test/baitian"):
    image_path=os.path.join("/disk3/dataset/head_test/baitian",image_name)
    image= cv2.imread(image_path)
    heigh = image.shape[0]
    width = image.shape[1]
    org_image=image
    re_size =640.0
    if max(image.shape[0], image.shape[1]) > int(re_size):
        re_shrink=re_size/max(image.shape[0],image.shape[1])
        y_shrink=1.5*re_size/image.shape[0]
        x_shrink=1.5*re_size/image.shape[1]
                #im_shrink =re_size/max(image.shape[0],image.shape[1])
        im_shrink=0.9
                #image = cv2.resize(image, None, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_LINEAR)
        #image = cv2.resize(image, None, None, fx=x_shrink, fy=y_shrink, interpolation=cv2.INTER_LINEAR)
        image = cv2.resize(image, None, None, fx=im_shrink, fy=im_shrink, interpolation=cv2.INTER_LINEAR)
                #image = cv2.resize(image, None, None, fx=re_shrink, fy=re_shrink, interpolation=cv2.INTER_LINEAR)
        time_start=time.time()
            # print("after height is {},width is{}".format(image.shape[0],image.shape[1]))
            # width=image.shape[1]
            # heigh=image.shape[0]
        net.blobs['data'].reshape(1, 3, image.shape[0], image.shape[1])
        transformer = caffe.io.Transformer({'data': net. blobs['data'].data.shape})
        transformer.set_transpose('data', (2, 0, 1))
        transformed_image = transformer.preprocess('data', image)
        net.blobs['data'].data[...] = transformed_image
        time_end=time.time()
        process_t = (time_end-time_start)*1000
        time_start=time.time()
        output= net.forward()
        detections=output['detection_out']
        score=net.blobs['mbox_conf_reshape'].data[0]

            #scores=net.forward()['mbox_conf_softmax']
            #print(scores)
        time_end=time.time()
        forward_t = (time_end-time_start)*1000
        print ("forward_time is {} ms".format(forward_t))
        time_start=time.time()
        class_name=detections[0,0,:,1]
        det_conf = detections[0, 0, :, 2]
        det_xmin = detections[0, 0, :, 3]
        det_ymin = detections[0, 0, :, 4]
        det_xmax = detections[0, 0, :, 5]
        det_ymax = detections[0, 0, :, 6]
        print("detection num is {}".format(len(class_name)))
        # print("using time is{}" .format(res_t))
        #print("nms before is{}".format(det_conf.shape[0]))
        a=0
        for i in range(det_conf.shape[0]):
            a=a+1
            score = det_conf[i]
            if score < 0.75:                                #0.34
                continue
            #print("accurary is {}".format(score))
            xmin = int(round(det_xmin[i] *width))
            ymin = int(round(det_ymin[i] * heigh))
            xmax = int(round(det_xmax[i] * width))
            ymax = int(round(det_ymax[i] * heigh))
            # xmin = int(round(det_xmin[i] * width))
            # ymin = int(round(det_ymin[i] * heigh)+280)
            # xmax = int(round(det_xmax[i] * width))
            # ymax = int(round(det_ymax[i] * heigh)+280)
            cv2.rectangle(org_image, (xmin, ymin), (xmax, ymax), (0, 0, 255), 3)
            #print("image_sise:%dx%d"%(image.shape[1],image.shape[0]))
            # print("face_size:%dx%d"%(xmax-xmin,ymax-ymin))
            # print(class_name[i].dtype)
            font_size = max((xmax - xmin)*0.01, 1)
            cv2.putText(org_image, '%.03f'%score, (xmin, ymin+int((ymax - ymin)*0.15)), cv2.FONT_HERSHEY_PLAIN, font_size, (0, 255, 0))
        #print("detect num is {}".format(n))
        #print("num sis {}".format(a))
        time_end = time.time()
        res_t = (time_end - time_start) * 1000
        print("process time is {}".format(res_t))
        save_name="/disk3/face_detect/beijing/sfd_head_train/test/test_result7/"+image_name
            # save1_name="/disk3/face_detect/beijing/sfd_head_train/test/result1/"+str(n)+"_real"+".jpg"

        cv2.imwrite(save_name,org_image)
            # out.write(org_image)
        cv2.namedWindow("video",0)
        cv2.imshow("video",org_image)