YOLOX 训练自己的VOC数据集

【YOLOX训练部署】YOLOX训练自己的VOC数据集_乐亦亦乐的博客-CSDN博客YOLOX 环境安装与训练自己标注的VOC数据集；https://blog.csdn.net/qq_41251963/article/details/122262738

训练好的模型：模型所在路径YOLOX/YOLOX_outputs/yolox_voc_s 

将自己的模型转化成ONNX

参考redeme：YOLOX/demo/ONNXRuntime at main · Megvii-BaseDetection/YOLOX · GitHub

pip install onnx

pip install onnxsim-simplifier

运行以下命令进行模型转化：

python3 tools/export_onnx.py --output-name your_yolox.onnx -f exps/your_dir/your_yolox.py -c your_yolox.pth

python tools/export_onnx.py --output-name my_yolox_s.onnx -f exps/example/yolox_voc/yolox_voc_s.py -c YOLOX_outputs/yolox_voc_s/best_ckpt.pth

模型转化成功。

ONNXRuntime 推理

cd demo/ONNXRuntime/

修改demo/ONNXRuntime/onnx_inference.py 文件：

from yolox.data.datasets import VOC_CLASSES

class_names=VOC_CLASSES 

运行：

python3 onnx_inference.py -m <ONNX_MODEL_PATH> -i <IMAGE_PATH> -o <OUTPUT_DIR> -s 0.3 --input_shape 640,640

Notes:

• -m: your converted onnx model
• -i: input_image
• -s: score threshold for visualization.
• --input_shape: should be consistent with the shape you used for onnx convertion.

python onnx_inference.py -m /media/liqiang/新加卷/YOLOX/my_yolox_s.onnx -i /media/liqiang/新加卷/YOLOX/assets/8_169.jpg -o /media/liqiang/新加卷/YOLOX -s 0.3 --input_shape 640,640

 成功运行！

视频推理代码：onnx_inference_video.py

'''
Descripttion: 
version: 
Author: LiQiang
Date: 2022-01-01 09:39:19
LastEditTime: 2022-01-01 10:12:11
'''
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.

import argparse
import os

import cv2
import numpy as np

import onnxruntime

from yolox.data.data_augment import preproc as preprocess
# from yolox.data.datasets import COCO_CLASSES
from yolox.data.datasets import VOC_CLASSES
from yolox.utils import mkdir, multiclass_nms, demo_postprocess, vis


def make_parser():
    parser = argparse.ArgumentParser("onnxruntime inference sample")
    parser.add_argument(
        "-m",
        "--model",
        type=str,
        default="yolox.onnx",
        help="Input your onnx model.",
    )
    parser.add_argument(
        "-i",
        "--video_path",
        type=str,
        # default='test_image.png',
        help="Path to your input image.",
    )
    parser.add_argument(
        "-o",
        "--output_dir",
        type=str,
        default='demo_output',
        help="Path to your output directory.",
    )
    parser.add_argument(
        "-s",
        "--score_thr",
        type=float,
        default=0.3,
        help="Score threshould to filter the result.",
    )
    parser.add_argument(
        "--input_shape",
        type=str,
        default="640,640",
        help="Specify an input shape for inference.",
    )
    parser.add_argument(
        "--with_p6",
        action="store_true",
        help="Whether your model uses p6 in FPN/PAN.",
    )
    return parser


if __name__ == '__main__':
    args = make_parser().parse_args()
    input_shape = tuple(map(int, args.input_shape.split(',')))
    # origin_img = cv2.imread(args.image_path)
    session = onnxruntime.InferenceSession(args.model)
    cap = cv2.VideoCapture(args.video_path)
    while True:
        ret, origin_img = cap.read()

        img, ratio = preprocess(origin_img, input_shape)

        ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
        output = session.run(None, ort_inputs)
        predictions = demo_postprocess(output[0], input_shape, p6=args.with_p6)[0]

        boxes = predictions[:, :4]
        scores = predictions[:, 4:5] * predictions[:, 5:]

        boxes_xyxy = np.ones_like(boxes)
        boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
        boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
        boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
        boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
        boxes_xyxy /= ratio
        dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
        if dets is not None:
            final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
            origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,
                             conf=args.score_thr, class_names=VOC_CLASSES)
        cv2.imshow('result', origin_img)
        c = cv2.waitKey(1)
        if c == 27:
            break
        # mkdir(args.output_dir)
        # output_path = os.path.join(args.output_dir, args.image_path.split("/")[-1])
        # cv2.imwrite(output_path, origin_img)

运行：

python onnx_inference_video.py -m /media/liqiang/新加卷/YOLOX/my_yolox_s.onnx -i ./4.mp4 -o /media/liqiang/新加卷/YOLOX -s 0.3 --input_shape 640,640

版权声明：本文为CSDN博主「乐亦亦乐」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_41251963/article/details/122265224