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

Posted 2023-03-11 Arno1988

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'