成功运行onnx推理模型；
parent
dfd9569120
commit
29c2033053
|
|
@ -3,7 +3,7 @@
|
|||
#batch=1
|
||||
#subdivisions=1
|
||||
# Training
|
||||
batch=16
|
||||
batch=1
|
||||
subdivisions=1
|
||||
width=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
|
||||
|
||||
from models import *
|
||||
from utils.utils import *
|
||||
from utils.datasets import *
|
||||
# coding: utf-8
|
||||
# author: hxy
|
||||
# 2019-12-10
|
||||
|
||||
"""
|
||||
照片的inference；
|
||||
默认推理过程在CPU上；
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import datetime
|
||||
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 logging
|
||||
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.checker.check_model(model)
|
||||
# 加载onnx模型
|
||||
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(
|
||||
ImageFolder(opt.image_folder,img_size=opt.img_size),
|
||||
batch_size = opt.batch_size,
|
||||
shuffle=False,
|
||||
num_workers=opt.n_cpu,
|
||||
)
|
||||
return sess, in_name, out_name
|
||||
|
||||
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 = []
|
||||
img_detections = []
|
||||
# anchors
|
||||
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:")
|
||||
prev_time = time.time()
|
||||
for batch_i,(img_paths,input_imgs) in enumerate(dataloader):
|
||||
input_imgs = Variable(input_imgs.type(Tensor))
|
||||
logging.debug(imgs)
|
||||
for img_name in imgs:
|
||||
img_full_path = os.path.join(img_files_path, img_name)
|
||||
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}
|
||||
yolo_output = model.run(['output'], yolo_inputs)[0]
|
||||
|
||||
detections = non_max_suppression(yolo_output, opt.conf_thres, opt.nms_thres)
|
||||
|
||||
# Log progress
|
||||
current_time = time.time()
|
||||
inference_time = datetime.timedelta(seconds=current_time - prev_time)
|
||||
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()
|
||||
# logging.info("推理照片 %s 耗时：% .2fms" % (img_name, ((time.time() - s)*1000)))
|
||||
boxes = get_boxes(prediction=prediction,
|
||||
anchors=anchors_yolo,
|
||||
img_shape=input_shape)
|
||||
draw_box(boxes=boxes,
|
||||
img=img,
|
||||
img_shape=img_shape)
|
||||
logging.info("推理照片 %s 耗时：% .2fms" % (img_name, ((time.time() - s)*1000)))
|
||||
|
|
|
|||
|
|
@ -448,7 +448,8 @@ class GraphBuilderONNX(object):
|
|||
if verbose:
|
||||
print(helper.printable_graph(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
|
||||
|
||||
def _make_onnx_node(self, layer_name, layer_dict):
|
||||
|
|
@ -757,13 +758,15 @@ class GraphBuilderONNX(object):
|
|||
assert channels > 0
|
||||
upsample_params = UpsampleParams(layer_name, scales)
|
||||
scales_name = upsample_params.generate_param_name()
|
||||
|
||||
# For ONNX opset >= 9, the Upsample node takes the scales array
|
||||
# as an input.
|
||||
inputs.append("")
|
||||
inputs.append(scales_name)
|
||||
|
||||
upsample_node = helper.make_node(
|
||||
'Resize',
|
||||
mode='nearest',
|
||||
"Resize",
|
||||
mode="nearest",
|
||||
inputs=inputs,
|
||||
outputs=[layer_name],
|
||||
name=layer_name,
|
||||
|
|
|
|||
Loading…
Reference in New Issue