用python写openvino yolov5目标检测代码

最新推荐文章于 2024-10-09 11:24:45 发布

CHHC1880

最新推荐文章于 2024-10-09 11:24:45 发布

阅读量452

点赞数

文章标签： python openvino YOLO

本文链接：https://blog.csdn.net/arno1988/article/details/129667971

版权

这段代码展示了如何使用OpenVINO库在Python中实现YOLOv5目标检测器。它加载模型权重，处理输入图像，并通过OpenVINO进行推理，以检测图像中的物体。检测到的物体包括各种类别，如人、车辆、动物等。程序还包含了后处理步骤，如NMS（非极大值抑制），以去除重复的检测框。

摘要由CSDN通过智能技术生成

检测代码

import cv2
import numpy as np
import time

#from openvino.runtime import Core  # the version of openvino >= 2022.1 # openvino 2022.1.0 has requirement numpy<1.20,>=1.16.6
from openvino.inference_engine import IECore # the version of openvino <= 2021.4.2

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']  # class names

class OpenVinoYoloV5Detector():

    def __init__(self,IN_conf):
        # ie = Core()  # Initialize Core version>=2022.1
        # self.Net = ie.compile_model(model=IN_conf.get("weight_file"),device_name=IN_conf.get("device"))

        ie = IECore()  # Initialize IECore  openvino <= 2021.4.2
        self.Net = ie.load_network(network=IN_conf.get("weight_file"), device_name=IN_conf.get("device"))

        self.INPUT_HEIGHT = 640
        self.INPUT_WIDTH = 640

    # YOLOv5目标检测
    def detect(self, image):

        blob = cv2.dnn.blobFromImage(image, 1 / 255.0, (self.INPUT_WIDTH, self.INPUT_HEIGHT), swapRB=True, crop=False)

        # openvino >= 2022.1
        # results = net([blob])[next(iter(net.outputs))]
        # results = self.Net([blob])[self.Net.output(0)]

        # openvino <= 2021.4.2
        results = self.Net.infer(inputs={"images": blob})
        results = results["output"]

        self.process_results(image, results)

    # YOLOv5的后处理函数，解析模型的输出
    def process_results(self,image, results,thresh=0.25):

        h, w, _ = image.shape

        class_ids = []
        boxes = []
        scores = []

        results = results[0]
        rows = results.shape[0]

        y_factor = h / self.INPUT_HEIGHT
        x_factor = w / self.INPUT_WIDTH

        for r in range(rows):
            row = results[r]
            score = row[4]

            if score >= 0.4 :

                classes_scores = row[5:]

                _, _, _, max_indexes = cv2.minMaxLoc(classes_scores)
                class_id = max_indexes[1]

                if classes_scores[class_id] > 0.25:
                    x, y, w, h = row[0].item(), row[1].item(), row[2].item(), row[3].item()

                    x1 = int((x - 0.5 * w) * x_factor)
                    y1 = int((y - 0.5 * h) * y_factor)
                    x2 = x1 + int(w * x_factor)
                    y2 = y1 + int(h * y_factor)
                    boxes.append((x1,y1, x2, y2))
                    class_ids.append(class_id)
                    scores.append(score)

        # default score_threshold=0.25, nms_threshold=0.45
        indices = cv2.dnn.NMSBoxes(bboxes=boxes, scores=scores, score_threshold=thresh, nms_threshold=0.45)

        for index in indices:
            x1, y1, x2,y2 = boxes[index][0], boxes[index][1], boxes[index][2], boxes[index][3]
            cv2.rectangle(frame, (x1,y1), (x2,y2), (0,255,0), 2)
            cv2.putText(frame, str(classes[class_ids[index]]), (x1, y1 + 20), cv2.FONT_HERSHEY_SIMPLEX, .5,(255,255,255))

if __name__ == '__main__':
    IN_conf = {
        "weight_file": "weights/yolov5n_openvino_model/yolov5n.xml",
        "device": "CPU"#"GPU"
    }
    detector = OpenVinoYoloV5Detector(IN_conf=IN_conf)

    url = 'bus.jpg'

    cap = cv2.VideoCapture(url)

    while True:
        r, frame = cap.read()
        if r:
            t1 = time.time()
            detector.detect(frame)
            cv2.imshow('OpenVinoYoloV5Detector.py', frame)
            t2 = time.time()
        else:
            print("读取%s结束" % str(url))
            break

    cv2.waitKey(0)
    cap.release()
    cv2.destroyAllWindows()