vehicle-demo/yolov5onnxtest/testonnxvideo.py

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()