目录
前言
1、有中文标签的数据集
2、yolov5代码修改为支持中文标签
很多人在训练yolov5目标检测的时候,标签只能显示英文的。怎么样才可以训练一个可以检测物体并且显示中文标签的模型呢。下面我们来一步一步的做。
首先在收集数据集的时候,打部分公开的数据集,标签都是英文的,格式都是voc格式,文件是xml文件,当然xml文件格式的标签是特别直观的,直观的知道标签中的图片大小,标注位置,还有标注的类别。虽然训练yolov5的时候用的是yolo格式的数据集,但是大家在标注数据集的时候尽量标注为voc格式,这样方便又直观,然后训练YOLOv5模型的时候有代码可以一键转换。代码博客:目标检测---数据集格式转化及训练集和验证集划分。
自己制作中文标签的数据集的话,利用labelimg这个工具标注数据集的时候,将标签改成中文就好了,具体的教程博客:目标检测---利用labelimg制作自己的深度学习目标检测数据集。
如果这里有一份英文标签的voc格式的数据集,可以利用如下的代码将数据集中的英文标签转为中文标签。
# encoding:utf-8 import os import xml.etree.ElementTree as ET count = 0 list_xml = [] dict = {"ball": "足球", "messi": "梅西", } openPath = "VOCdevkit\VOC2007\Annotations" savePath = "VOCdevkit\VOC2007\Annotations1" fileList = os.listdir(openPath) # 得到进程当前工作目录中的所有文件名称列表 for fileName in fileList: # 获取文件列表中的文件 if fileName.endswith(".xml"): # 只看xml文件 print("filename=:", fileName) tree = ET.parse(os.path.join(openPath, fileName)) root = tree.getroot() print("root-tag=:", root.tag) # ',root-attrib:', root.attrib, ',root-text:', root.text) for child in root: # 第一层解析 if child.tag == "object": # 找到object标签 print(child.tag) for sub in child: if sub.tag == "name": print("标签名字:", sub.tag, ";文本内容:", sub.text) if sub.text not in list_xml: list_xml.append(sub.text) if sub.text in list(dict.keys()): sub.text = dict[sub.text] print(sub.text) count = count + 1 tree.write(os.path.join(savePath, fileName), encoding='utf-8') print("=" * 20) print(count) for i in list_xml: print(i)
这个代码也可以将中文标签的voc格式的数据集转化为英文标签的数据集。如下图所示,将中文和英文的位置调换就可以了。
中文标签的voc格式的数据集,可以利用目标检测---数据集格式转化及训练集和验证集划分这篇博客中的代码将voc格式的数据集转化为yolo数据集,并划分为训练集和验证集。
至此中文标签的数据集就准备好了
特别要说明一下,该博客所用的是yolov5的5.0版本。
(1)train.py文件修改,在py文件中的63行,代码修改如下:
with open(opt.data, encoding='UTF-8') as f:
(2)test.py文件文件修改,在py文件中的73行,代码修改如下:
with open(opt.data, encoding='UTF-8') as f:
(3)utils/general.py文件文件修改,在这个代码中导入如下的包:
from PIL import Image, ImageDraw, ImageFont
(4)utils/plots.py文件修改,在py文件中的64行,修改 plot_one_box 函数,if label之后的代码改为:
tf = max(tl - 1, 1) # font thickness t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] font_size = t_size[1] font = ImageFont.truetype('msyh.ttc', font_size) t_size = font.getsize(label) c2 = c1[0] + t_size[0], c1[1] - t_size[1] cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled img_PIL = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) draw = ImageDraw.Draw(img_PIL) draw.text((c1[0], c2[1] - 2), label, fill=(255, 255, 255), font=font) return cv2.cvtColor(np.array(img_PIL), cv2.COLOR_RGB2BGR)
(5)utils/plots.py文件修改,在py文件中的184行,修改plot_images 函数:
mosaic = plot_one_box(box, mosaic, label=label, color=color, line_thickness=tl)
(6)detect.py文件修改,在py文件中的113行,代码修改如下:
im0 = plot_one_box(xyxy, im0, label=label, color=colors(c, True), line_thickness=3)
至此为了支持中文标签的代码就修改完了。
如果后面的模型训练和推理测试工作可以参考博客:目标检测---教你利用yolov5训练自己的目标检测模型。
这里特别要说明一下,就是修改data目录下的yaml文件的时候,那个标签也一定要对应标注的中文类别。