[Caffe-SSD][原创]利用caffe的python API进行caffe ssd模型GPU加速预测

网上很多都是opencv写的，没办法gpu加速，我开始也用opencv写过，但是不加速运行很慢。没办法，亲自操刀写了一个利用caffe的python API进行caffe ssd模型预测的代码，代码在ubuntu16.04 caffe python2.7 cuda10.0+cudnn7.6.5运行成功。利用这个代码可以轻松扩展为摄像头。视频文件预测，代码写的有有点乱，嘿嘿，我故意的，请食用，祝您用餐愉快

import os
import numpy as np
import cv2

import sys

sys.path.append("/home/fut/caffe-ssd/python")
import caffe

caffe.set_device(0)
caffe.set_mode_gpu()

model_def = '/home/fut/caffe-ssd/myproj/prototxt/deploy.prototxt'
model_weights = '/home/fut/caffe-ssd/myproj/model/model_SSD_300x300_iter_97648.caffemodel'
net = caffe.Net(model_def,  # defines the structure of the model
                model_weights,  # contains the trained weights
                caffe.TEST)  # use test mode (e.g., don't perform dropout)
# input preprocessing: 'data' is the name of the input blob == net.inputs[0]
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
transformer.set_transpose('data', (2, 0, 1))
transformer.set_mean('data', np.array([104, 117, 123]))  # mean pixel
transformer.set_raw_scale('data', 255)  # the reference model operates on images in [0,255] range instead of [0,1]
transformer.set_channel_swap('data', (2, 1, 0))  # the reference model has channels in BGR order instead of RGB

# set net to batch size of 1
image_resize = 300
net.blobs['data'].reshape(1, 3, image_resize, image_resize)
image_file = '/home/fut/Desktop/share/Femal-732/JPEGImages/20201010_100742_287718.jpg'
# image = caffe.io.load_image(image_file)
# print("image type is:", type(image))
# print ("image shape:",image.shape)
image = cv2.imread('/home/fut/Desktop/share/Femal-732/JPEGImages/20201010_100742_287718.jpg')
image1 = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image1 = image1 / 255.
transformed_image = transformer.preprocess('data', image1)
net.blobs['data'].data[...] = transformed_image
# Forward pass.
detections = net.forward()['detection_out']
print(type(detections))
print (detections.shape)


def load_labels(file_name):
    labels = []
    with open(file_name, "r") as f:
        labels = f.read().split('\n').strip('\n')
    return labels


def SSDPostProcess(inference_result, image_resolution, confidence_threshold, labels=[]):
    detection_result_list = []
    lines = inference_result.shape[2]
    print("row count is:", lines)
    for i in range(lines):
        if inference_result[0][0][i][2] < confidence_threshold:
            continue
        class_index = int(inference_result[0][0][i][1])
        confidence = float(inference_result[0][0][i][2])
        xmin = int(max(min(float(inference_result[0][0][i][3]), 1), 0) * image_resolution[1])
        ymin = int(max(min(float(inference_result[0][0][i][4]), 1), 0) * image_resolution[0])
        xmax = int(max(min(float(inference_result[0][0][i][5]), 1), 0) * image_resolution[1])
        ymax = int(max(min(float(inference_result[0][0][i][6]), 1), 0) * image_resolution[0])
        detection_result_list.append([labels[class_index], confidence, xmin, ymin, xmax, ymax])
    print(detection_result_list)
    return detection_result_list


def imshow(detection_result_list,img):
    font = cv2.FONT_HERSHEY_SIMPLEX
    for result in detection_result_list:
        img = cv2.putText(img, result[0] + "," + str(result[1]), (result[2] - 5, result[3]), font, 2, (0, 0, 255), 1)
        cv2.rectangle(image, (result[2], result[3]), (result[4], result[5]), (0, 0, 255), 4)
    cv2.imshow("result", img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


result_list = SSDPostProcess(detections, (image.shape[0], image.shape[1]), 0.2, ["1", "2", "3"])
imshow(result_list,image)