YOLOX 重写推理携带标签信息

最新推荐文章于 2026-07-13 15:50:59 发布

原创最新推荐文章于 2026-07-13 15:50:59 发布 · 253 阅读

0 ·

本内容遵循CC 4.0 BY-SA版权协议

收录于

码农的日常

本文介绍了一种YOLOX模型推理代码的重写方法，该方法不仅可以保存推理后的图片，还可以在图片上显示推理框及对应的标签信息。通过预处理、模型推理和后处理等步骤，实现了对输入图像的目标检测。

YOLOX 重写推理携带标签信息

YOLOX模型推理重写携带推理label

重新编写官方推理代码，可以保存推理图片，同时携带推理框和推理label

引入

import cv2
import torch
from yolox.data.datasets.dt_classes import VOC_CLASSES as COCO_CLASSES
from yolox.data.data_augment import preproc
from yolox.utils import postprocess
from exps.example.yolox_voc.yolox_voc_YZmulti_s_resume import Exp

推理代码

class Detector:
    def __init__(self, model_path=" ", rank=0):#自定义模型路径
        exp = Exp()
        self.model = exp.get_model()
        self.class_names = COCO_CLASSES
        torch.cuda.set_device(rank)
        loc = rank if rank == "cpu" else "cuda:{0}".format(rank)
        self.device = torch.device(loc)
        ckpt = torch.load(model_path, map_location=loc)
        # load the model state dict
        self.model.load_state_dict(ckpt["model"])
        self.num_classes = exp.num_classes
        self.confthre = 0.75
        self.nmsthre = 0.75
        self.test_size = exp.test_size

        # cuda and eval
        self.model = self.model.to(self.device)
        self.model.eval()

    @torch.no_grad()
    def detect(self, input_img, flag=0):
        img, ratio = preproc(input_img, self.test_size)
        img = torch.from_numpy(img).unsqueeze(0).to(self.device)
        outputs = self.model(img)
        outputs = postprocess(outputs, self.num_classes, self.confthre, self.nmsthre)
        if outputs[0] is not None:
            outputs_list = outputs[0].cpu().numpy().tolist()
            for d in outputs_list:
                boxes = [int(d[i] / ratio) for i in range(4)]
                label = '{}'.format(self.class_names[int(d[-1])])
                Detector.plot_one_box(boxes, input_img, color=(0, 0, 255), label=label)
                cv2.imwrite('./new.jpg',input_img)
        return input_img

    @staticmethod
    def plot_one_box(x, img, color, label=None, line_thickness=None):
        tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1  # line/font thickness
        c1, c2 = (x[0], x[1]), (x[2], x[3])
        cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
        if label:
            tf = max(tl - 1, 1)  # font thickness
            t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
            c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
            cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled
            cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)