OpenVino快速落地部署教程

SkyXZ發表於2024-08-27

OpenVino快速落地部署教程

Openvino是由Intel開發的專門用於最佳化和部署人工智慧推理的半開源的工具包，主要用於對深度推理做最佳化。本教程適用於Yolov5-7.0，直接跑Yolov5為6FPS，使用OpenVino後為30FPS，未來將會出一系列其他模型（Paddle等）的OpenVino部署教程，測試平臺——Intel Nuc 11代i5處理器

一、安裝OpenVino

進入OpenVino官網

https://docs.openvino.ai/2024/get-started/install-openvino.html

選擇自己喜歡的下載方式，本教程採用OpenVino-2022.3.1版本

二、模型轉換

  1. 透過Yolov5自帶的export.py檔案將.pt轉為.onnx格式

    python3 export.py --weights xxxx/xxxxx.pt --include onnx --batch_size 1 --opset 10

PS:如果出現轉換失敗的提示，如：opset 10不支援或是onnx版本問題請重新搭建yolov5環境，按照requirements.txt裡庫的最低版本進行安裝

  2. 使用OpenVino工具鏈將.onnx轉為xml、bin模型

    mo --input_model xxx/xxx.onnx

PS：如果openvino環境安裝成功將可以在yolov5的環境中直接使用mo命令

PS：轉換完成後請一定用模型視覺化工具檢視轉換是否正確

三、採用以下程式碼快速部署

import openvino.runtime as ov
import cv2
import numpy as np
import openvino.preprocess as op

class ObjectDetector:
    def __init__(self, model_xml, model_bin, labels, device="CPU"):
        self.core = ov.Core()
        self.model = self.core.read_model(model_xml, model_bin)
        self.labels = labels
        self.preprocess_model()
        self.compiled_model = self.core.compile_model(self.model, device)
        self.infer_request = self.compiled_model.create_infer_request()

    def preprocess_model(self):
        premodel = op.PrePostProcessor(self.model)
        premodel.input().tensor().set_element_type(ov.Type.u8).set_layout(ov.Layout("NHWC")).set_color_format(op.ColorFormat.BGR)
        premodel.input().preprocess().convert_element_type(ov.Type.f32).convert_color(op.ColorFormat.RGB).scale([255., 255., 255.])
        premodel.input().model().set_layout(ov.Layout("NCHW"))
        premodel.output(0).tensor().set_element_type(ov.Type.f32)
        self.model = premodel.build()

    def infer(self, img):
        detections = []
        img_re, dw, dh = self.resizeimg(img, (640, 640))
        input_tensor = np.expand_dims(img_re, 0)
        self.infer_request.infer({0: input_tensor})
        output = self.infer_request.get_output_tensor(0)
        detections = self.process_output(output.data[0])
        return detections

    def process_output(self, detections):
        boxes = []
        class_ids = []
        confidences = []
        for prediction in detections:
            confidence = prediction[4].item()
            if confidence >= 0.6:
                classes_scores = prediction[5:]
                _, _, _, max_indx = cv2.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 = prediction[0].item(), prediction[1].item(), prediction[2].item(), prediction[3].item()
                    xmin = x - (w / 2)
                    ymin = y - (h / 2)
                    box = np.array([xmin, ymin, w, h])
                    boxes.append(box)
        indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.5)
        detections = []
        for i in indexes:
            j = i.item()
            detections.append({"class_index": class_ids[j], "confidence": confidences[j], "box": boxes[j]})
        return detections

    def resizeimg(self, image, new_shape):
        old_size = image.shape[:2]
        ratio = float(new_shape[-1] / max(old_size))
        new_size = tuple([int(x * ratio) for x in old_size])
        image = cv2.resize(image, (new_size[1], new_size[0]))
        delta_w = new_shape[1] - new_size[1]
        delta_h = new_shape[0] - new_size[0]
        color = [100, 100, 100]
        new_im = cv2.copyMakeBorder(image, 0, delta_h, 0, delta_w, cv2.BORDER_CONSTANT, value=color)
        return new_im, delta_w, delta_h


if __name__ == "__main__":
    # Example usage:
    labels = [
        "right",
        "warning",
        "left",
        "people",
        "10",
        "pullover",
        "10off",
        "green",
        "red"
    ]
    detector = ObjectDetector("/home/nuc/MyCar/yolov5-7.0/best.xml", "/home/nuc/MyCar/yolov5-7.0/best.bin", labels, "CPU")
    cap = cv2.VideoCapture(0)
    while True:
        ret, frame = cap.read()
        detections = detector.infer(frame)
        for detection in detections:
            classId = detection["class_index"]
            confidence = detection["confidence"]
            label = labels[classId]
            box = detection["box"]
            area = box[2] * box[3]
            print(f"Detected object: {label}, Confidence: {confidence}, Area: {area}")
    cap.release()

相關文章