openvino +yolov5 自己训练模型并测试

最新推荐文章于 2023-09-18 17:00:56 发布

CHHC1880

最新推荐文章于 2023-09-18 17:00:56 发布

阅读量995

点赞数 1

分类专栏：人工智能文章标签： openvino YOLO python

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

版权

人工智能专栏收录该内容

4 篇文章 0 订阅

订阅专栏

该文详细介绍了如何将Yolov5模型转换为ONNX格式，然后利用Intel的OpenVINO工具套件进行模型优化和部署。通过设置环境变量、安装必要组件、执行模型转换脚本，以及训练和测试自定义数据集，最终实现了一个基于OpenVINO的Yolov5目标检测应用。

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

https://github.com/ultralytics/yolov5/releases/tag/v6.1

conda create -n openvino_yolov6 python=3.7 -y
conda activate openvino_yolov6
pip install -r requirements.txt

export.py
export_onnx(model, im, file, 10, train, False, simplify) # opset 12 10

python export.py --weights yolov5n.pt --img 640 --batch 1

* 设置一个临时的环境变量
cd C:\Program Files (x86)\Intel\openvino_2021.4.752\bin
setupvars.bat

* 安装一些常用的模型优化器
cd C:\Program Files (x86)\Intel\openvino_2021.4.752\deployment_tools\model_optimizer\install_prerequisites
install_prerequisites.bat

* 模型转换
cd C:\Program Files (x86)\Intel\openvino_2021.4.752\deployment_tools\model_optimizer
python mo_onnx.py --input_model H:\tt\yolov5-6.1\yolov5n.onnx --output_dir E:\yolo5n_IR\

测试
python detect.py --weights yolov5n.pt --source data/images/bus.jpg

训练 yolov5 - CHHC - 博客园 (cnblogs.com)

python train.py --weights yolov5s.pt --data data/catdog.yaml --workers 1 --batch-size 8

训练测试
python detect.py --weights catdog.pt --source data/images/1.png
conda activate openvino_yolov6
python export.py --weights catdog.pt --img 640 --batch 1
python mo_onnx.py --input_model H:\tt\yolov5-6.1\catdog.onnx --output_dir E:\IR\catdog

app.py

#!/usr/local/bin/python3
# encodin: utf-8
import cv2
import time
import os

from OpenVinoYoloV5Detector import OpenVinoYoloV5Detector

classes = []

classes_base = ['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

classes_catdog = ['cat', 'dog']  # class names

def get_milsecond():
    t = time.time()
    return (int(round(t * 1000)))


if __name__ == '__main__':
    url = '1.png'
    box_color = (0, 255, 0)

    # yolov5
    conf = {
        # "weight_file": "weights/yolov5n_openvino_model/yolov5n.xml",
        # "weight_file": "weights/yolov5s_openvino_model/yolov5s.xml",
        "weight_file": "weights/catdog_openvino_model/catdog.xml",
        "device": "CPU"
    }
    classes = classes_catdog
    detector = OpenVinoYoloV5Detector(IN_conf=conf)

    # ssd
    # conf = {
    #     "model_xml": "./weights/ssdlite_mobilenet_v2/FP16/ssdlite_mobilenet_v2.xml",
    #     "model_bin": "./weights/ssdlite_mobilenet_v2/FP16/ssdlite_mobilenet_v2.bin",
    #     "device": "CPU"
    # }
    # rtscap.detector = OpenVinoSSDDetector(IN_conf=conf)

    cap = cv2.VideoCapture(url)

    while True:
        ret, frame = cap.read()
        if not ret or frame is None:
            break

        starttime = get_milsecond()
        detect_num, detect_data = detector.detect(frame)
        if len(detect_data):
            #cv2.imwrite('1.png', frame, [int(cv2.IMWRITE_JPEG_QUALITY), 95])

            for m in detect_data:
                classid= int(m.get('class_id'))
                classname = classes[classid]
                score = m.get('score')
                location = m.get('location')
                box_l, box_t = int(location.get('x1')), int(location.get('y1'))
                box_r, box_b = int(location.get('x2')), int(location.get('y2'))
                frame = cv2.rectangle(frame, (box_l, box_t), (box_r, box_b), box_color, 2)
                frame = cv2.putText(frame, classname + " " + str(score), (box_l, box_t + 15), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 255), 2)

        endtime = get_milsecond();
        print('检测耗时:' + str(endtime - starttime))

        cv2.imshow('openvino detection', frame)
        if cv2.waitKey(5) & 0xFF == ord('q'):
            break

    os.system('pause'