YOLOv8训练自定义数据集后，模型预测结果没有显示正确的标签，而是claa_0、class_1

weixin_74350204

已于 2024-08-25 00:22:48 修改

阅读量566

点赞数 5

文章标签： YOLO

于 2024-08-25 00:20:27 首次发布

本文链接：https://blog.csdn.net/weixin_74350204/article/details/141507354

版权

前言

根据YOLOv8的相关教程训练自定义数据集后，却发现图片的预测结果并没有显示正确的标签，而是class_0、class_1.....

对此，可以通过对YOLOv8模型预测后图片的绘制代码进行修改，以显示正确的标签类别。

假设自定义数据集的data.yaml如下：

nc: 3
name:
    0: "cat"
    1: "dog"
    2: "others"

根据标签类别的一一对应关系，自定义一个标签映射字典：

classes = {"class_0":"cat",
           "class_1":"dog",
           "class_2":"others"}

在YOLOv8官方源码ultralytics文件夹下，进入yolo目录下，打开plotting.py文件。

定位到def box_label函数：

参数label为模型预测的标签字符串，为了利用该函数的已知参数，以减少修改代码量，将参数label修改成label1：

再对自定义的标签字典进行键值对的查找、拼接，代码如下：

classes = {"class_0":"cat",
           "class_1":"dog",
           "class_2":"others"}
label = classes[label1[0:7]] + label1[7:]

对于YOLOv8训练自定义数据集的模型，如果出现无法正确显示预测标签信息，可以通过预测后的图片绘制进行修正。

关注