From 29c203305310bd3b250b8f747f154c39d09a38f5 Mon Sep 17 00:00:00 2001 From: 12345qiupeng Date: Sat, 4 Mar 2023 19:47:42 +0800 Subject: [PATCH] =?UTF-8?q?=E6=88=90=E5=8A=9F=E8=BF=90=E8=A1=8Connx?= =?UTF-8?q?=E6=8E=A8=E7=90=86=E6=A8=A1=E5=9E=8B=EF=BC=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- config/yolov3-416.cfg | 2 +- darknet_api.py | 165 +++++++++++++++++++++++++++++++++++++++++ detect_onnx.py | 166 ++++++++++++------------------------------ yolo2onnx.py | 9 ++- 4 files changed, 219 insertions(+), 123 deletions(-) create mode 100644 darknet_api.py diff --git a/config/yolov3-416.cfg b/config/yolov3-416.cfg index 946e015..1725128 100644 --- a/config/yolov3-416.cfg +++ b/config/yolov3-416.cfg @@ -3,7 +3,7 @@ #batch=1 #subdivisions=1 # Training -batch=16 +batch=1 subdivisions=1 width=416 height=416 diff --git a/darknet_api.py b/darknet_api.py new file mode 100644 index 0000000..c4bf5c7 --- /dev/null +++ b/darknet_api.py @@ -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 diff --git a/detect_onnx.py b/detect_onnx.py index 649571a..b2c0f30 100644 --- a/detect_onnx.py +++ b/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() \ No newline at end of file + # 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))) diff --git a/yolo2onnx.py b/yolo2onnx.py index 1fc7a5b..6341011 100644 --- a/yolo2onnx.py +++ b/yolo2onnx.py @@ -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,