成功运行onnx推理模型；
parent
dfd9569120
commit
29c2033053
|
|
@ -3,7 +3,7 @@
|
||||||
#batch=1
|
#batch=1
|
||||||
#subdivisions=1
|
#subdivisions=1
|
||||||
# Training
|
# Training
|
||||||
batch=16
|
batch=1
|
||||||
subdivisions=1
|
subdivisions=1
|
||||||
width=416
|
width=416
|
||||||
height=416
|
height=416
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,165 @@
|
||||||
|
# coding: utf-8
|
||||||
|
# 2019-12-10
|
||||||
|
"""
|
||||||
|
YOlo相关的预处理api；
|
||||||
|
"""
|
||||||
|
import cv2
|
||||||
|
import time
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
|
||||||
|
# 加载label names；
|
||||||
|
def get_labels(names_file):
|
||||||
|
names = list()
|
||||||
|
with open(names_file, 'r') as f:
|
||||||
|
lines = f.read()
|
||||||
|
for name in lines.splitlines():
|
||||||
|
names.append(name)
|
||||||
|
f.close()
|
||||||
|
return names
|
||||||
|
|
||||||
|
|
||||||
|
# 照片预处理
|
||||||
|
def process_img(img_path, input_shape):
|
||||||
|
ori_img = cv2.imread(img_path)
|
||||||
|
img = cv2.resize(ori_img, input_shape)
|
||||||
|
image = img[:, :, ::-1].transpose((2, 0, 1))
|
||||||
|
image = image[np.newaxis, :, :, :] / 255
|
||||||
|
image = np.array(image, dtype=np.float32)
|
||||||
|
return ori_img, ori_img.shape, image
|
||||||
|
|
||||||
|
|
||||||
|
# 视频预处理
|
||||||
|
def frame_process(frame, input_shape):
|
||||||
|
image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||||
|
image = cv2.resize(image, input_shape)
|
||||||
|
# image = cv2.resize(image, (640, 480))
|
||||||
|
image_mean = np.array([127, 127, 127])
|
||||||
|
image = (image - image_mean) / 128
|
||||||
|
image = np.transpose(image, [2, 0, 1])
|
||||||
|
image = np.expand_dims(image, axis=0)
|
||||||
|
image = image.astype(np.float32)
|
||||||
|
return image
|
||||||
|
|
||||||
|
|
||||||
|
# sigmoid函数
|
||||||
|
def sigmoid(x):
|
||||||
|
s = 1 / (1 + np.exp(-1 * x))
|
||||||
|
return s
|
||||||
|
|
||||||
|
|
||||||
|
# 获取预测正确的类别，以及概率和索引;
|
||||||
|
def get_result(class_scores):
|
||||||
|
class_score = 0
|
||||||
|
class_index = 0
|
||||||
|
for i in range(len(class_scores)):
|
||||||
|
if class_scores[i] > class_score:
|
||||||
|
class_index += 1
|
||||||
|
class_score = class_scores[i]
|
||||||
|
return class_score, class_index
|
||||||
|
|
||||||
|
|
||||||
|
# 通过置信度筛选得到bboxs
|
||||||
|
def get_bbox(feat, anchors, image_shape, confidence_threshold=0.25):
|
||||||
|
box = list()
|
||||||
|
for i in range(len(anchors)):
|
||||||
|
for cx in range(feat.shape[0]):
|
||||||
|
for cy in range(feat.shape[1]):
|
||||||
|
tx = feat[cx][cy][0 + 85 * i]
|
||||||
|
ty = feat[cx][cy][1 + 85 * i]
|
||||||
|
tw = feat[cx][cy][2 + 85 * i]
|
||||||
|
th = feat[cx][cy][3 + 85 * i]
|
||||||
|
cf = feat[cx][cy][4 + 85 * i]
|
||||||
|
cp = feat[cx][cy][5 + 85 * i:85 + 85 * i]
|
||||||
|
|
||||||
|
bx = (sigmoid(tx) + cx) / feat.shape[0]
|
||||||
|
by = (sigmoid(ty) + cy) / feat.shape[1]
|
||||||
|
bw = anchors[i][0] * np.exp(tw) / image_shape[0]
|
||||||
|
bh = anchors[i][1] * np.exp(th) / image_shape[1]
|
||||||
|
b_confidence = sigmoid(cf)
|
||||||
|
b_class_prob = sigmoid(cp)
|
||||||
|
b_scores = b_confidence * b_class_prob
|
||||||
|
b_class_score, b_class_index = get_result(b_scores)
|
||||||
|
|
||||||
|
if b_class_score >= confidence_threshold:
|
||||||
|
box.append([bx, by, bw, bh, b_class_score, b_class_index])
|
||||||
|
return box
|
||||||
|
|
||||||
|
|
||||||
|
# 采用nms算法筛选获取到的bbox
|
||||||
|
def nms(boxes, nms_threshold=0.6):
|
||||||
|
l = len(boxes)
|
||||||
|
if l == 0:
|
||||||
|
return []
|
||||||
|
else:
|
||||||
|
b_x = boxes[:, 0]
|
||||||
|
b_y = boxes[:, 1]
|
||||||
|
b_w = boxes[:, 2]
|
||||||
|
b_h = boxes[:, 3]
|
||||||
|
scores = boxes[:, 4]
|
||||||
|
areas = (b_w + 1) * (b_h + 1)
|
||||||
|
order = scores.argsort()[::-1]
|
||||||
|
keep = list()
|
||||||
|
while order.size > 0:
|
||||||
|
i = order[0]
|
||||||
|
keep.append(i)
|
||||||
|
xx1 = np.maximum(b_x[i], b_x[order[1:]])
|
||||||
|
yy1 = np.maximum(b_y[i], b_y[order[1:]])
|
||||||
|
xx2 = np.minimum(b_x[i] + b_w[i], b_x[order[1:]] + b_w[order[1:]])
|
||||||
|
yy2 = np.minimum(b_y[i] + b_h[i], b_y[order[1:]] + b_h[order[1:]])
|
||||||
|
|
||||||
|
# 相交面积,不重叠时面积为0
|
||||||
|
w = np.maximum(0.0, xx2 - xx1 + 1)
|
||||||
|
h = np.maximum(0.0, yy2 - yy1 + 1)
|
||||||
|
inter = w * h
|
||||||
|
# 相并面积,面积1+面积2-相交面积
|
||||||
|
union = areas[i] + areas[order[1:]] - inter
|
||||||
|
# 计算IoU：交 /（面积1+面积2-交）
|
||||||
|
IoU = inter / union
|
||||||
|
# 保留IoU小于阈值的box
|
||||||
|
inds = np.where(IoU <= nms_threshold)[0]
|
||||||
|
order = order[inds + 1] # 因为IoU数组的长度比order数组少一个,所以这里要将所有下标后移一位
|
||||||
|
|
||||||
|
final_boxes = [boxes[i] for i in keep]
|
||||||
|
return final_boxes
|
||||||
|
|
||||||
|
|
||||||
|
# 绘制预测框
|
||||||
|
def draw_box(boxes, img, img_shape):
|
||||||
|
label = ["background", "person",
|
||||||
|
"bicycle", "car", "motorbike", "aeroplane",
|
||||||
|
"bus", "train", "truck", "boat", "traffic light",
|
||||||
|
"fire hydrant", "stop sign", "parking meter", "bench",
|
||||||
|
"bird", "cat", "dog", "horse", "sheep", "cow", "elephant",
|
||||||
|
"bear", "zebra", "giraffe", "backpack", "umbrella", "handbag",
|
||||||
|
"tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
|
||||||
|
"kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
|
||||||
|
"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon",
|
||||||
|
"bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog",
|
||||||
|
"pizza", "donut", "cake", "chair", "sofa", "potted plant", "bed", "dining table",
|
||||||
|
"toilet", "TV monitor", "laptop", "mouse", "remote", "keyboard", "cell phone",
|
||||||
|
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
|
||||||
|
"scissors", "teddy bear", "hair drier", "toothbrush"]
|
||||||
|
for box in boxes:
|
||||||
|
x1 = int((box[0] - box[2] / 2) * img_shape[1])
|
||||||
|
y1 = int((box[1] - box[3] / 2) * img_shape[0])
|
||||||
|
x2 = int((box[0] + box[2] / 2) * img_shape[1])
|
||||||
|
y2 = int((box[1] + box[3] / 2) * img_shape[0])
|
||||||
|
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
||||||
|
cv2.putText(img, label[int(box[5])] + ":" + str(round(box[4], 3)), (x1 + 5, y1 + 10), cv2.FONT_HERSHEY_SIMPLEX,
|
||||||
|
0.5, (0, 0, 255), 1)
|
||||||
|
print(label[int(box[5])] + ":" + "概率值：%.3f" % box[4])
|
||||||
|
cv2.imshow('image', img)
|
||||||
|
cv2.waitKey(10)
|
||||||
|
cv2.destroyAllWindows()
|
||||||
|
|
||||||
|
|
||||||
|
# 获取预测框
|
||||||
|
def get_boxes(prediction, anchors, img_shape, confidence_threshold=0.25, nms_threshold=0.6):
|
||||||
|
boxes = []
|
||||||
|
for i in range(len(prediction)):
|
||||||
|
feature_map = prediction[i][0].transpose((2, 1, 0))
|
||||||
|
box = get_bbox(feature_map, anchors[i], img_shape, confidence_threshold)
|
||||||
|
boxes.extend(box)
|
||||||
|
Boxes = nms(np.array(boxes), nms_threshold)
|
||||||
|
return Boxes
|
||||||
166
detect_onnx.py
166
detect_onnx.py
|
|
@ -1,133 +1,61 @@
|
||||||
from __future__ import division
|
# coding: utf-8
|
||||||
|
# author: hxy
|
||||||
from models import *
|
# 2019-12-10
|
||||||
from utils.utils import *
|
|
||||||
from utils.datasets import *
|
|
||||||
|
|
||||||
|
"""
|
||||||
|
照片的inference；
|
||||||
|
默认推理过程在CPU上；
|
||||||
|
"""
|
||||||
import os
|
import os
|
||||||
import sys
|
|
||||||
import time
|
import time
|
||||||
import datetime
|
import logging
|
||||||
import argparse
|
|
||||||
|
|
||||||
from PIL import Image
|
|
||||||
|
|
||||||
import torch
|
|
||||||
from torch.utils.data import DataLoader
|
|
||||||
from torchvision import datasets
|
|
||||||
from torch.autograd import Variable
|
|
||||||
|
|
||||||
import matplotlib.pyplot as plt
|
|
||||||
import matplotlib.patches as patches
|
|
||||||
from matplotlib.ticker import NullLocator
|
|
||||||
|
|
||||||
import onnx
|
|
||||||
import onnxruntime
|
import onnxruntime
|
||||||
|
from darknet_api import process_img, get_boxes, draw_box
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
parser = argparse.ArgumentParser()
|
|
||||||
parser.add_argument("--image_folder", type=str, default="data/samples", help="path to dataset")
|
|
||||||
parser.add_argument("--model_def", type=str, default="config/yolov3.cfg", help="path to model definition file")
|
|
||||||
parser.add_argument("--weights_path", type=str, default="weights/yolov3.weights", help="path to weights file")
|
|
||||||
parser.add_argument("--class_path", type=str, default="data/coco.names", help="path to class label file")
|
|
||||||
parser.add_argument("--conf_thres", type=float, default=0.8, help="object confidence threshold")
|
|
||||||
parser.add_argument("--nms_thres", type=float, default=0.4, help="iou thresshold for non-maximum suppression")
|
|
||||||
parser.add_argument("--batch_size", type=int, default=1, help="size of the batches")
|
|
||||||
parser.add_argument("--n_cpu", type=int, default=0, help="number of cpu threads to use during batch generation")
|
|
||||||
parser.add_argument("--img_size", type=int, default=416, help="size of each image dimension")
|
|
||||||
parser.add_argument("--checkpoint_model", type=str, help="path to checkpoint model")
|
|
||||||
opt = parser.parse_args()
|
|
||||||
print(opt)
|
|
||||||
|
|
||||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
# 定义日志格式
|
||||||
|
def log_set():
|
||||||
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||||
|
|
||||||
os.makedirs("output", exist_ok=True)
|
|
||||||
|
|
||||||
model = onnxruntime.InferenceSession("output/yolov3.onnx")
|
# 加载onnx模型
|
||||||
onnx.checker.check_model(model)
|
def load_model(onnx_model):
|
||||||
|
sess = onnxruntime.InferenceSession(onnx_model)
|
||||||
|
in_name = [input.name for input in sess.get_inputs()][0]
|
||||||
|
out_name = [output.name for output in sess.get_outputs()]
|
||||||
|
logging.info("输入的name:{}, 输出的name:{}".format(in_name, out_name))
|
||||||
|
|
||||||
dataloader = DataLoader(
|
return sess, in_name, out_name
|
||||||
ImageFolder(opt.image_folder,img_size=opt.img_size),
|
|
||||||
batch_size = opt.batch_size,
|
|
||||||
shuffle=False,
|
|
||||||
num_workers=opt.n_cpu,
|
|
||||||
)
|
|
||||||
|
|
||||||
classes = load_classes(opt.class_path)
|
|
||||||
|
|
||||||
Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
|
if __name__ == '__main__':
|
||||||
|
log_set()
|
||||||
|
input_shape = (416 , 416)
|
||||||
|
|
||||||
imgs = []
|
# anchors
|
||||||
img_detections = []
|
anchors_yolo = [[(116, 90), (156, 198), (373, 326)], [(30, 61), (62, 45), (59, 119)],
|
||||||
|
[(10, 13), (16, 30), (33, 23)]]
|
||||||
|
anchors_yolo_tiny = [[(81, 82), (135, 169), (344, 319)], [(10, 14), (23, 27), (37, 58)]]
|
||||||
|
session, inname, outname = load_model(onnx_model='output/yolov3-416.onnx')
|
||||||
|
logging.info("开始Inference....")
|
||||||
|
# 照片的批量inference
|
||||||
|
img_files_path = 'data/samples'
|
||||||
|
imgs = os.listdir(img_files_path)
|
||||||
|
|
||||||
print("\nPerforming object detection:")
|
logging.debug(imgs)
|
||||||
prev_time = time.time()
|
for img_name in imgs:
|
||||||
for batch_i,(img_paths,input_imgs) in enumerate(dataloader):
|
img_full_path = os.path.join(img_files_path, img_name)
|
||||||
input_imgs = Variable(input_imgs.type(Tensor))
|
logging.debug(img_full_path)
|
||||||
|
img, img_shape, testdata = process_img(img_path=img_full_path,
|
||||||
|
input_shape=input_shape)
|
||||||
|
s = time.time()
|
||||||
|
prediction = session.run(outname, {inname: testdata})
|
||||||
|
|
||||||
yolo_inputs = {'input': input_imgs}
|
# logging.info("推理照片 %s 耗时：% .2fms" % (img_name, ((time.time() - s)*1000)))
|
||||||
yolo_output = model.run(['output'], yolo_inputs)[0]
|
boxes = get_boxes(prediction=prediction,
|
||||||
|
anchors=anchors_yolo,
|
||||||
detections = non_max_suppression(yolo_output, opt.conf_thres, opt.nms_thres)
|
img_shape=input_shape)
|
||||||
|
draw_box(boxes=boxes,
|
||||||
# Log progress
|
img=img,
|
||||||
current_time = time.time()
|
img_shape=img_shape)
|
||||||
inference_time = datetime.timedelta(seconds=current_time - prev_time)
|
logging.info("推理照片 %s 耗时：% .2fms" % (img_name, ((time.time() - s)*1000)))
|
||||||
prev_time = current_time
|
|
||||||
print("\t+ Batch %d, Inference Time: %s" % (batch_i, inference_time))
|
|
||||||
|
|
||||||
# Save image and detections
|
|
||||||
imgs.extend(img_paths)
|
|
||||||
img_detections.extend(detections)
|
|
||||||
|
|
||||||
# Bounding-box colors
|
|
||||||
cmap = plt.get_cmap("tab20b")
|
|
||||||
colors = [cmap(i) for i in np.linspace(0, 1, 20)]
|
|
||||||
|
|
||||||
print("\nSaving images:")
|
|
||||||
# Iterate through images and save plot of detections
|
|
||||||
for img_i, (path, detections) in enumerate(zip(imgs, img_detections)):
|
|
||||||
|
|
||||||
print("(%d) Image: '%s'" % (img_i, path))
|
|
||||||
|
|
||||||
# Create plot
|
|
||||||
img = np.array(Image.open(path))
|
|
||||||
plt.figure()
|
|
||||||
fig, ax = plt.subplots(1)
|
|
||||||
ax.imshow(img)
|
|
||||||
|
|
||||||
# Draw bounding boxes and labels of detections
|
|
||||||
if detections is not None:
|
|
||||||
# Rescale boxes to original image
|
|
||||||
detections = rescale_boxes(detections, opt.img_size, img.shape[:2])
|
|
||||||
unique_labels = detections[:, -1].cpu().unique()
|
|
||||||
n_cls_preds = len(unique_labels)
|
|
||||||
bbox_colors = random.sample(colors, n_cls_preds)
|
|
||||||
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
|
|
||||||
print("\t+ Label: %s, Conf: %.5f" % (classes[int(cls_pred)], cls_conf.item()))
|
|
||||||
|
|
||||||
box_w = x2 - x1
|
|
||||||
box_h = y2 - y1
|
|
||||||
|
|
||||||
color = bbox_colors[int(np.where(unique_labels == int(cls_pred))[0])]
|
|
||||||
# Create a Rectangle patch
|
|
||||||
bbox = patches.Rectangle((x1, y1), box_w, box_h, linewidth=2, edgecolor=color, facecolor="none")
|
|
||||||
# Add the bbox to the plot
|
|
||||||
ax.add_patch(bbox)
|
|
||||||
# Add label
|
|
||||||
plt.text(
|
|
||||||
x1,
|
|
||||||
y1,
|
|
||||||
s=classes[int(cls_pred)],
|
|
||||||
color="white",
|
|
||||||
verticalalignment="top",
|
|
||||||
bbox={"color": color, "pad": 0},
|
|
||||||
)
|
|
||||||
|
|
||||||
# Save generated image with detections
|
|
||||||
plt.axis("off")
|
|
||||||
plt.gca().xaxis.set_major_locator(NullLocator())
|
|
||||||
plt.gca().yaxis.set_major_locator(NullLocator())
|
|
||||||
filename = path.split("/")[-1].split(".")[0]
|
|
||||||
plt.savefig(f"output/{filename}.png", bbox_inches="tight", pad_inches=0.0)
|
|
||||||
plt.close()
|
|
||||||
|
|
|
||||||
|
|
@ -448,7 +448,8 @@ class GraphBuilderONNX(object):
|
||||||
if verbose:
|
if verbose:
|
||||||
print(helper.printable_graph(self.graph_def))
|
print(helper.printable_graph(self.graph_def))
|
||||||
model_def = helper.make_model(self.graph_def,
|
model_def = helper.make_model(self.graph_def,
|
||||||
producer_name='NVIDIA TensorRT sample')
|
producer_name='NVIDIA TensorRT sample',
|
||||||
|
opset_imports=[helper.make_opsetid(domain="", version=17)])
|
||||||
return model_def
|
return model_def
|
||||||
|
|
||||||
def _make_onnx_node(self, layer_name, layer_dict):
|
def _make_onnx_node(self, layer_name, layer_dict):
|
||||||
|
|
@ -757,13 +758,15 @@ class GraphBuilderONNX(object):
|
||||||
assert channels > 0
|
assert channels > 0
|
||||||
upsample_params = UpsampleParams(layer_name, scales)
|
upsample_params = UpsampleParams(layer_name, scales)
|
||||||
scales_name = upsample_params.generate_param_name()
|
scales_name = upsample_params.generate_param_name()
|
||||||
|
|
||||||
# For ONNX opset >= 9, the Upsample node takes the scales array
|
# For ONNX opset >= 9, the Upsample node takes the scales array
|
||||||
# as an input.
|
# as an input.
|
||||||
|
inputs.append("")
|
||||||
inputs.append(scales_name)
|
inputs.append(scales_name)
|
||||||
|
|
||||||
upsample_node = helper.make_node(
|
upsample_node = helper.make_node(
|
||||||
'Resize',
|
"Resize",
|
||||||
mode='nearest',
|
mode="nearest",
|
||||||
inputs=inputs,
|
inputs=inputs,
|
||||||
outputs=[layer_name],
|
outputs=[layer_name],
|
||||||
name=layer_name,
|
name=layer_name,
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue