1-python onnexruntime 部署Yolov5深度学习模型

1-效果图

**

2-实现代码

**

import onnxruntime
from torchvision import transforms
import cv2 as cv
import time
import numpy as np


img_transform = transforms.Compose([transforms.ToTensor(),
                                    transforms.Resize((640, 640))
                                    ])
# pip install onnxruntime-gpu==1.7 -i https://pypi.tuna.tsinghua.edu.cn/simple
def load_classes():
    with open("classes.txt", "r") as f:
        class_list = [cname.strip() for cname in f.readlines()]
    return class_list


def format_yolov5(frame):
    row, col, _ = frame.shape
    _max = max(col, row)
    result = np.zeros((_max, _max, 3), np.uint8)
    result[0:row, 0:col] = frame
    result = cv.cvtColor(result, cv.COLOR_BGR2RGB)
    return result

def wrap_detection(input_image, output_data):
    class_ids = []
    confidences = []
    boxes = []
    # print(output_data.shape)
    rows = output_data.shape[0]

    image_width, image_height, _ = input_image.shape

    x_factor = image_width / 640.0
    y_factor = image_height / 640.0

    for r in range(rows):
        row = output_data[r]
        confidence = row[4]
        if confidence >= 0.4:

            classes_scores = row[5:]
            _, _, _, max_indx = cv.minMaxLoc(classes_scores)
            class_id = max_indx[1]
            if (classes_scores[class_id] > .25):
                confidences.append(confidence)

                class_ids.append(class_id)

                x, y, w, h = row[0].item(), row[1].item(), row[2].item(), row[3].item()
                left = int((x - 0.5 * w) * x_factor)
                top = int((y - 0.5 * h) * y_factor)
                width = int(w * x_factor)
                height = int(h * y_factor)
                box = np.array([left, top, width, height])
                boxes.append(box)

    indexes = cv.dnn.NMSBoxes(boxes, confidences, 0.25, 0.45)

    result_class_ids = []
    result_confidences = []
    result_boxes = []

    for i in indexes:
        result_confidences.append(confidences[i])
        result_class_ids.append(class_ids[i])
        result_boxes.append(boxes[i])

    return result_class_ids, result_confidences, result_boxes

def gpu_ort_demo():
    device_name = onnxruntime.get_device()
    print(device_name)
    class_list = load_classes()
    session = onnxruntime.InferenceSession("yolov5s.onnx", providers=['CUDAExecutionProvider'])
    # capture = cv.VideoCapture("D:/images/video/sample.mp4")
    capture = cv.VideoCapture("testDnn.mp4")
    colors = [(255, 255, 0), (0, 255, 0), (0, 255, 255), (255, 0, 0)]
    while True:
        start = time.time()
        _, frame = capture.read()
        if frame is None:
            print("End of stream")
            break
        image = format_yolov5(frame)
        x_input = img_transform(image).view(1, 3, 640, 640)
        ort_inputs = {session.get_inputs()[0].name: x_input.numpy()}
        ort_outs = session.run(None, ort_inputs)
        out_prob = ort_outs[0]
        # print(out_prob.shape)
        class_ids, confidences, boxes = wrap_detection(image, out_prob[0])
        for (classid, confidence, box) in zip(class_ids, confidences, boxes):
            color = colors[int(classid) % len(colors)]
            cv.rectangle(frame, box, color, 2)
            cv.rectangle(frame, (box[0], box[1] - 20), (box[0] + box[2], box[1]), color, -1)
            cv.putText(frame, class_list[classid], (box[0], box[1] - 10), cv.FONT_HERSHEY_SIMPLEX, .5, (0, 0, 0))

        end = time.time()
        inf_end = end - start
        fps = 1 / inf_end
        fps_label = "FPS: %.2f" % fps
        cv.putText(frame, fps_label, (10, 25), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        cv.imshow("YOLOv5 + ONNXRUNTIME by gloomyfish", frame)
        cc = cv.waitKey(1)
        if cc == 27:
            break
    cv.waitKey(0)
    cv.destroyAllWindows()


if __name__ == "__main__":
    gpu_ort_demo()