From f709af43004be4e3c280d3fdb86eb7f659bd1e13 Mon Sep 17 00:00:00 2001
From: 12345qiupeng <qpeng0504@163.com>
Date: Mon, 6 Mar 2023 11:41:08 +0800
Subject: [PATCH] =?UTF-8?q?=E6=88=90=E5=8A=9F=E5=9C=A8NX=E4=B8=8A=E8=BF=90?=
 =?UTF-8?q?=E8=A1=8C=E6=8E=A8=E7=90=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 darknet_api.py | 25 +++++++++++++++++++++++++
 detect_onnx.py | 11 ++++++++++-
 yolo2onnx.py   |  2 +-
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/darknet_api.py b/darknet_api.py
index c4bf5c7..d0eb1de 100644
--- a/darknet_api.py
+++ b/darknet_api.py
@@ -153,6 +153,31 @@ def draw_box(boxes, img, img_shape):
     cv2.waitKey(10)
     cv2.destroyAllWindows()
 
+def draw_box_save(boxes,img,img_shape,img_path):
+    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.imwrite(img_path,img)
 
 # 获取预测框
 def get_boxes(prediction, anchors, img_shape, confidence_threshold=0.25, nms_threshold=0.6):
diff --git a/detect_onnx.py b/detect_onnx.py
index b2c0f30..5b7d72d 100644
--- a/detect_onnx.py
+++ b/detect_onnx.py
@@ -10,7 +10,7 @@ import os
 import time
 import logging
 import onnxruntime
-from darknet_api import process_img, get_boxes, draw_box
+from darknet_api import process_img, get_boxes, draw_box,draw_box_save
 
 
 # 定义日志格式
@@ -41,10 +41,13 @@ if __name__ == '__main__':
     # 照片的批量inference
     img_files_path = 'data/samples'
     imgs = os.listdir(img_files_path)
+    img_output_path = 'output/'
 
     logging.debug(imgs)
     for img_name in imgs:
         img_full_path = os.path.join(img_files_path, img_name)
+        img_save_path = os.path.join(img_output_path,img_name)
+
         logging.debug(img_full_path)
         img, img_shape, testdata = process_img(img_path=img_full_path,
                                                input_shape=input_shape)
@@ -58,4 +61,10 @@ if __name__ == '__main__':
         draw_box(boxes=boxes,
                  img=img,
                  img_shape=img_shape)
+        draw_box_save(
+            boxes=boxes,
+            img=img,
+            img_shape=img_shape,
+            img_path=img_save_path
+        )
         logging.info("推理照片 %s 耗时：% .2fms" % (img_name, ((time.time() - s)*1000)))
diff --git a/yolo2onnx.py b/yolo2onnx.py
index 6341011..2843dfe 100644
--- a/yolo2onnx.py
+++ b/yolo2onnx.py
@@ -449,7 +449,7 @@ class GraphBuilderONNX(object):
             print(helper.printable_graph(self.graph_def))
         model_def = helper.make_model(self.graph_def,
                                       producer_name='NVIDIA TensorRT sample',
-                                      opset_imports=[helper.make_opsetid(domain="", version=17)])
+                                      opset_imports=[helper.make_opsetid(domain="", version=15)])
         return model_def
 
     def _make_onnx_node(self, layer_name, layer_dict):