瀏覽代碼

增加更多 yolo 资源

zhzhenqin 5 月之前
父節點
當前提交
857b4c41ab
共有 4 個文件被更改,包括 55 次插入1 次删除
  1. 2 1
      docs/YOLO理论研究.md
  2. 48 0
      onnx_labels.py
  3. 3 0
      python-export-onnx.sh
  4. 2 0
      requirements.txt

+ 2 - 1
docs/YOLO理论研究.md

@@ -1,2 +1,3 @@
 https://blog.csdn.net/weixin_45303602/article/details/129175854
-https://blog.csdn.net/qq_32892383/article/details/136505299
+https://blog.csdn.net/qq_32892383/article/details/136505299
+https://blog.csdn.net/weixin_45277161/article/details/136566919

+ 48 - 0
onnx_labels.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python
+# -*- coding:utf-8 -*-
+
+# Created by zhenqin.
+# User: zhenqin
+# Date: 2024-04-25
+# Time: 下午1:58
+# Verdor: yiidata.com
+#
+
+# 针对有网络模型,但还没有训练保存 .pth 文件的情况
+import onnxruntime
+import netron
+import onnx
+import torch.onnx
+
+def printLabels():
+    # Load Label from Model's meta data
+    session = onnxruntime.InferenceSession("./yolov8s.onnx", providers=['CPUExecutionProvider'])
+    meta = session.get_modelmeta().custom_metadata_map  # metadata
+    label_names = eval(meta['names'])
+    print(label_names)
+
+def addLabels():
+    f = "./yolov7.onnx"
+    new_f = "./yolov7-425-labels.onnx"
+    model_onnx = onnx.load(f)  # load onnx model
+    onnx.checker.check_model(model_onnx)  # check onnx model
+    labels = ["person", "bicycle", "car", "motorcycle", "airplane", "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", "couch",
+                "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
+                "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
+                "hair drier", "toothbrush"]
+    index = 1
+    meta = model_onnx.metadata_props.add()
+    meta.key, meta.value = "labels", ",".join(labels)
+    print(meta.key + ":" + meta.value)
+    onnx.save(model_onnx, new_f)
+
+if __name__ == '__main__':
+    # printLabels()
+    # netron.start("./yolov8s.onnx")  # 输出网络结构
+    netron.start("./yolov7-425-labels.onnx")
+    # addLabels()

+ 3 - 0
python-export-onnx.sh

@@ -0,0 +1,3 @@
+python export-det.py --weights water_strean_model.pt \
+--iou-thres 0.65 --conf-thres 0.25 --topk 100 --opset 11 \
+--sim --input-shape 1 3 640 640 --device cpu

+ 2 - 0
requirements.txt

@@ -0,0 +1,2 @@
+onnxsim
+tensorrt==10.2.0