Yolo v5 Convert to ONNX

yolov5 Weight로 ONNX Convert

1. 아래 스크립트를 shell에 입력하면 .pt파일이 있던 경로에 같은 이름의 .onnx 파일이 생성된다.

python export.py --weights <.pt 파일 경로> --include torchscript onnx

2. detect.py를 사용할 때 참고할 argument들

weights=ROOT / 'yolov5s.pt',  # model path or triton URL
        source=ROOT / 'data/images',  # file/dir/URL/glob/screen/0(webcam)
        data=ROOT / 'data/coco128.yaml',  # dataset.yaml path
        imgsz=(640, 640),  # inference size (height, width)
        conf_thres=0.25,  # confidence threshold
        iou_thres=0.45,  # NMS IOU threshold
        max_det=1000,  # maximum detections per image
        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
        view_img=False,  # show results
        save_txt=False,  # save results to *.txt
        save_conf=False,  # save confidences in --save-txt labels
        save_crop=False,  # save cropped prediction boxes
        nosave=False,  # do not save images/videos
        classes=None,  # filter by class: --class 0, or --class 0 2 3
        agnostic_nms=False,  # class-agnostic NMS
        augment=False,  # augmented inference
        visualize=False,  # visualize features
        update=False,  # update all models
        project=ROOT / 'runs/detect',  # save results to project/name
        name='exp',  # save results to project/name
        exist_ok=False,  # existing project/name ok, do not increment
        line_thickness=3,  # bounding box thickness (pixels)
        hide_labels=False,  # hide labels
        hide_conf=False,  # hide confidences
        half=False,  # use FP16 half-precision inference
        dnn=False,  # use OpenCV DNN for ONNX inference
        vid_stride=1,  # video frame-rate stride

3. 생성된 .onnx 파일은 아래 스크립트를 통해 inference 할 수 있다.

python detect.py \
--weights yolo_best.onnx \
--source ./sample_video \
--data data/koreafood.yaml \
--img 640 \
--conf 0.25 \
--save-txt \
--save-conf \
--hide-conf \
--vid-stride 30 \
--save-crop # 현재 물체 검출이 되어 crop될 경우 크롭된 부분만 가져오지 않고 전체를 가져오도록 코드를 수정

참고 사이트

ultralytics yolov5 사용법 완벽 정리

https://github.com/ultralytics/yolov5/issues/5832