模型
Git 基于SSD的端到端模型
模型查看
测试代码
使用python直接运行onnx验证效果
import onnx
import onnxruntime as ort
import numpy as np
import cv2
import time
# 原始onnx再netron无法显示shape需要转存一下
from onnx import shape_inference
path = "version-RFB-640.onnx"
onnx.save(onnx.shape_inference.infer_shapes(onnx.load(path)), "version-RFB-640-info.onnx")
onnx_model = onnx.load("version-RFB-640.onnx")
onnx.checker.check_model(onnx_model)
sess = ort.InferenceSession('version-RFB-640.onnx')
input_name = sess.get_inputs()[0].name
# 读取测试图片
image = cv2.imread('19.jpg')
# 保留原始图片后续绘制
origin = image.copy()
print(image.shape)
# 缩放至模型能够输入尺寸
image = cv2.resize(image, (640, 480))
print(image.shape)
# 颜色转换
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
# 归一化
image_mean = np.array([127, 127, 127])
image = (image - image_mean) / 128
# rgbrgb转为rrggbb
image = np.transpose(image, [2, 0, 1])
# 增加batch纬度
image = np.expand_dims(image, axis=0)
# 修改数据类型
image = image.astype(np.float32)
time_time = time.time()
# confidences (1, 17640, 2)
# boxes (1, 17640, 4)
confidences, boxes = sess.run(None, {input_name: image})
print("cost time:{}".format(time.time() - time_time))
confidences = confidences[0]
boxes = boxes[0]
# 绘制置信度大于0.7的人脸框
for class_index in range(1, confidences.shape[1]):
probs = confidences[:, class_index];
mask = probs > 0.7
probs = probs[mask]
if probs.shape[0] == 0:
continue
subset_boxes = boxes[mask, :]
print(probs)
print(subset_boxes)
for i in range(0, len(subset_boxes)):
box = subset_boxes[i]
w = origin.shape[1]
h = origin.shape[0]
cv2.rectangle(origin, (int(box[0]*w), int(box[1]*h)), (int(box[2]*w), int(box[3]*h)), (255,255,0), thickness=4)
cv2.imwrite("test.jpg", origin)
备注
- 17640是怎么计算出来的
17640=(60*80*12+30*40*8+15*20*8+80*12)/4
17640是边框的个数,除4是因为一个边框需要4个点描述
