最近用在用YOLO11算法做检测的时候,灵光乍现,发现某一个label有更好的命名方式,不改吧又觉得心里难受

我以为detect.py会调用数据集中的类别文件,于是尝试着先修改classes.txt文件中的标签名称,重新检测,发现没有变化,查看了一下代码,并没有重新读取修改后的classes.txt

然后打开训练好的权重文件,全是乱码,偶尔有几个能够看懂的英文,使用搜索查找了一下我想要修改的标签名,然后就改掉试试。

结果好家伙,检测直接报错,看来虽然标签在权重里面不是二进制代码,还是不能直接修改

然后又开始寻找新的解决办法,最后还真被我找到了

先在训练好的权重文件所在的目录下创建一个python脚本文件,然后输入以下代码

from ultralytics import YOLO

# 加载训练好的模型
model = YOLO("best.pt")

# 修改标签名(注意:顺序必须和原始一致)
model.model.names = ['label_1', 'label_2',..., 'label_n']

# 保存为新的权重文件
model.save("best_renamed.pt")

这样就能生成一个新的权重文件,在detect.py中使用新生成的权重文件来检测,如下所示

model = YOLO(model='runs/train/exp/weights/best_renamed.pt')

就能够使用新的标签来进行检测,就不用重新训练权重了,从而大大缩短了等待模型训练的大量摸鱼时间(bushi

这个脚本我只在YOLO11的权重上面尝试过,YOLOv8应该也都能使用,希望能够帮助到大家

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐