Explorar el Código

单独yolov5检测测试demo,需要的改进点:降低检测频率等方式降低CPU使用。现有结果和情况在environment.txt文件里

Curious hace 6 meses
padre
commit
db74e287ce
Se han modificado 2 ficheros con 290 adiciones y 0 borrados
  1. 25 0
      yolov5onnxtest/enviroment.txt
  2. 265 0
      yolov5onnxtest/testonnxvideo.py

+ 25 - 0
yolov5onnxtest/enviroment.txt

@@ -0,0 +1,25 @@
+python==3.8
+onnx==1.11
+
+CPU时常:130-140ms。GPU待解决
+CPU使用情况:总的为600%
+onnxruntime-gpu版本待定均有报错:
+1、1.11及1.10版本可以使用CPU但无法使用cuda,报错信息为:
+2024-09-04 15:33:24.369159907 [E:onnxruntime:Default, provider_bridge_ort.cc:1022 Get] Failed to load library libonnxruntime_providers_cuda.so with error: libcublas.so.10: cannot open shared object file: No such file or directory
+2024-09-04 15:33:24.369221636 [W:onnxruntime:Default, onnxruntime_pybind_state.cc:552 CreateExecutionProviderInstance] Failed to create CUDAExecutionProvider. Please reference https://onnxruntime.ai/docs/reference/execution-providers/CUDA-ExecutionProvider.html#requirements to ensure all dependencies are met.
+
+2、1.12.1及1.17(测试了两个版本)
+不再有libcublas.so.10: cannot open shared object file: No such file or directory的错误,但是有新的报错:
+Model correct
+Using CUDA for inference.
+2024-09-04 16:15:58.913593544 [E:onnxruntime:, sequential_executor.cc:368 Execute] Non-zero status code returned while running Conv node. Name:'/model.1/conv/Conv' Status Message: :0: cudaFuncSetAttribute(kernel_entry, cudaFuncAttributeMaxDynamicSharedMemorySize, integer_cast<int32_t>(launch_configs[0].smemSizeInBytes)): invalid device function
+Traceback (most recent call last):
+  File "testonnxvideo.py", line 263, in <module>
+    main()
+  File "testonnxvideo.py", line 242, in main
+    output, org_img = model.inference(frame)
+  File "testonnxvideo.py", line 86, in inference
+    pred = self.onnx_session.run(None, input_feed)[0]
+  File "/usr/local/lib/python3.8/dist-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 200, in run
+    return self._sess.run(output_names, input_feed, run_options)
+onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Conv node. Name:'/model.1/conv/Conv' Status Message: :0: cudaFuncSetAttribute(kernel_entry, cudaFuncAttributeMaxDynamicSharedMemorySize, integer_cast<int32_t>(launch_configs[0].smemSizeInBytes)): invalid device function

+ 265 - 0
yolov5onnxtest/testonnxvideo.py

@@ -0,0 +1,265 @@
+import os
+import sys
+import onnx
+import onnxruntime as ort
+import cv2
+import numpy as np
+import time  # 导入时间模块
+
+# CLASSES = ['person']
+CLASSES = ['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']  # coco80类别
+
+
+class Yolov5ONNX(object):
+    def __init__(self, onnx_path):
+        onnx_model = onnx.load(onnx_path)
+        try:
+            onnx.checker.check_model(onnx_model)
+        except Exception:
+            print("Model incorrect")
+        else:
+            print("Model correct")
+
+        # 设置 ONNX Runtime 会尝试优先使用 CUDAExecutionProvider,如果不可用则使用 CPUExecutionProvider
+        self.onnx_session = ort.InferenceSession(onnx_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+
+        # 检查是否使用了 CUDA
+        providers = self.onnx_session.get_providers()
+        if 'CUDAExecutionProvider' in providers:
+            print("Using CUDA for inference.")
+        else:
+            print("CUDA is not available, using CPU for inference.")
+
+        self.input_name = self.get_input_name()
+        self.output_name = self.get_output_name()
+        self.input_size = (640, 640)
+
+    def get_input_name(self):
+        input_name = []
+        for node in self.onnx_session.get_inputs():
+            input_name.append(node.name)
+        return input_name
+
+    def get_output_name(self):
+        output_name = []
+        for node in self.onnx_session.get_outputs():
+            output_name.append(node.name)
+        return output_name
+
+    def get_input_feed(self, image_numpy):
+        input_feed = {}
+        for name in self.input_name:
+            input_feed[name] = image_numpy
+        return input_feed
+
+    def inference(self, img):
+        h, w, _ = img.shape
+        new_w, new_h = self.input_size
+
+        scale = min(new_w / w, new_h / h)
+        new_w = int(w * scale)
+        new_h = int(h * scale)
+
+        img_resized = cv2.resize(img, (new_w, new_h))
+
+        padded_img = np.zeros((self.input_size[1], self.input_size[0], 3), dtype=np.uint8)
+        pad_x = (self.input_size[0] - new_w) // 2
+        pad_y = (self.input_size[1] - new_h) // 2
+        padded_img[pad_y:pad_y + new_h, pad_x:pad_x + new_w] = img_resized
+
+        img_rgb = cv2.cvtColor(padded_img, cv2.COLOR_BGR2RGB).transpose(2, 0, 1)
+        img_rgb = img_rgb.astype(np.float32)
+        img_rgb /= 255.0
+        img_rgb = np.expand_dims(img_rgb, axis=0)
+
+        input_feed = self.get_input_feed(img_rgb)
+
+        start_time = time.time()  # 记录开始时间
+        pred = self.onnx_session.run(None, input_feed)[0]
+        end_time = time.time()  # 记录结束时间
+
+        inference_time = end_time - start_time  # 计算推理时间
+        print(f"Inference time: {inference_time:.4f} seconds")  # 输出推理时间
+
+        return pred, padded_img
+
+
+def nms(dets, thresh):
+    x1 = dets[:, 0]
+    y1 = dets[:, 1]
+    x2 = dets[:, 2]
+    y2 = dets[:, 3]
+
+    areas = (y2 - y1 + 1) * (x2 - x1 + 1)
+    scores = dets[:, 4]
+
+    keep = []
+    index = scores.argsort()[::-1]
+
+    while index.size > 0:
+        i = index[0]
+        keep.append(i)
+
+        x11 = np.maximum(x1[i], x1[index[1:]])
+        y11 = np.maximum(y1[i], y1[index[1:]])
+        x22 = np.minimum(x2[i], x2[index[1:]])
+        y22 = np.minimum(y2[i], y2[index[1:]])
+
+        w = np.maximum(0, x22 - x11 + 1)
+        h = np.maximum(0, y22 - y11 + 1)
+
+        overlaps = w * h
+        ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
+
+        idx = np.where(ious <= thresh)[0]
+        index = index[idx + 1]
+    return keep
+
+
+def xywh2xyxy(x):
+    y = np.copy(x)
+    y[:, 0] = x[:, 0] - x[:, 2] / 2
+    y[:, 1] = x[:, 1] - x[:, 3] / 2
+    y[:, 2] = x[:, 0] + x[:, 2] / 2
+    y[:, 3] = x[:, 1] + x[:, 3] / 2
+
+    return y
+
+
+def filter_box(org_box, conf_thres, iou_thres):
+    org_box = np.squeeze(org_box)
+    conf = org_box[..., 4] > conf_thres
+    box = org_box[conf == True]
+
+    if box.size == 0:
+        return np.array([])  # 确保返回一个空的numpy数组
+
+    cls_cinf = box[..., 5:]
+    cls = [int(np.argmax(cls_cinf[i])) for i in range(len(cls_cinf))]
+
+    # 只保留类别为“person”的框
+    person_boxes = [box[i] for i in range(len(cls)) if cls[i] == 0]  # 类别“person”对应的索引是0
+
+    if len(person_boxes) == 0:
+        return np.array([])  # 如果没有“person”框,返回一个空的numpy数组
+
+    person_boxes = np.array(person_boxes)
+    person_boxes = xywh2xyxy(person_boxes)
+
+    person_out_box = nms(person_boxes, iou_thres)
+    output = [person_boxes[k] for k in person_out_box]
+
+    return np.array(output)
+
+
+def draw(image, box_data):
+    if box_data.size == 0:
+        return image
+
+    boxes = box_data[..., :4].astype(np.int32)
+    scores = box_data[..., 4]
+    classes = box_data[..., 5].astype(np.int32)
+    for box, score, cl in zip(boxes, scores, classes):
+        top, left, right, bottom = box
+
+        cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)
+        cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score),
+                    (top, left),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    0.6, (0, 0, 255), 2)
+    return image
+
+
+# def main():
+#     onnx_path = 'D:/Thework/yolov5-5.0/yolov5-5.0/runs/train/exp/weights/yolov5s.onnx'
+#     model = Yolov5ONNX(onnx_path)
+#
+#     cap = cv2.VideoCapture(0)  # 打开摄像头
+#     if not cap.isOpened():
+#         print("无法打开摄像头")
+#         sys.exit(0)
+#
+#     while True:
+#         ret, frame = cap.read()
+#         if not ret:
+#             print("无法读取摄像头图像")
+#             break
+#
+#         start_time = time.time()  # 记录开始时间
+#         output, org_img = model.inference(frame)
+#         outbox = filter_box(output, 0.5, 0.5)
+#         org_img = draw(org_img, outbox)
+#         end_time = time.time()  # 记录结束时间
+#
+#         inference_time = end_time - start_time  # 计算推理时间
+#         print(f"Total processing time: {inference_time:.4f} seconds")  # 输出总处理时间
+#
+#         cv2.imshow('result', org_img)
+#
+#         if cv2.waitKey(1) & 0xFF == ord('q'):
+#             break
+#
+#     cap.release()
+#     cv2.destroyAllWindows()
+#
+#
+# if __name__ == "__main__":
+#     main()
+def main():
+    onnx_path = 'yolov5s.onnx'
+    model = Yolov5ONNX(onnx_path)
+
+    # cap = cv2.VideoCapture(0)  # 打开摄像头
+    cap = cv2.VideoCapture(7, cv2.CAP_V4L2)
+    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
+    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
+    if not cap.isOpened():
+        print("无法打开摄像头")
+        sys.exit(0)
+
+    last_save_time = time.time()  # 记录最后一次保存的时间
+    frame_count = 0  # 用于命名保存的图像
+
+    # 创建保存图片的文件夹
+    output_dir = "saved_images"
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            print("无法读取摄像头图像")
+            break
+
+        output, org_img = model.inference(frame)
+        outbox = filter_box(output, 0.5, 0.5)
+        org_img = draw(org_img, outbox)
+
+        # 如果距离上次保存超过2秒,则保存图像
+        if time.time() - last_save_time >= 2:
+            frame_count += 1
+            image_path = os.path.join(output_dir, f'result_{frame_count}.jpg')
+            cv2.imwrite(image_path, org_img)  # 保存图像
+            print(f"Image saved: {image_path}")  # 输出保存路径
+            last_save_time = time.time()  # 更新最后一次保存的时间
+
+        # 如果保存了5张图片,退出循环
+        if frame_count >= 5:
+            print("保存了5张图片,程序退出")
+            break
+
+    cap.release()
+    cv2.destroyAllWindows()
+
+if __name__ == "__main__":
+    main()
+
+