yolov5 detection box displays Chinese labels

yolov5 Detection Box Displaying Chinese Labels#

Note: This annotation corresponds to version 5.0 of yolov5!!!#

2. Modify yolov5 Code to Support Chinese Labels

Introduction#

Many people find that when training the yolov5 object detection model, the labels can only display in English. How can we train a model that can detect objects and display Chinese labels? Let's do it step by step.

1. Dataset with Chinese Labels Dataset #

When collecting datasets, many publicly available datasets have labels in English and are in VOC format, with files in XML format. The XML format is particularly intuitive, as it clearly shows the image size, annotation position, and annotation categories. Although the yolov5 training uses YOLO format datasets, it is advisable to annotate datasets in VOC format for convenience and clarity, and there is code available for one-click conversion when training the YOLOv5 model. Code blog: Object Detection --- Dataset Format Conversion and Training/Validation Set Division.

If you are creating a dataset with Chinese labels, you can use the labelimg tool to annotate the dataset and simply change the labels to Chinese. For a detailed tutorial, refer to the blog: Object Detection --- Using labelimg to Create Your Own Deep Learning Object Detection Dataset.

If you have an English label VOC format dataset, you can use the following code to convert the English labels in the dataset to Chinese labels.

# 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)         # Get a list of all file names in the current working directory
for fileName in fileList:               # Get files from the file list
    if fileName.endswith(".xml"):       # Only look at xml files
        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: # First level parsing
            if child.tag == "object":   # Find object tag
                print(child.tag)
                for sub in child:
                    if sub.tag == "name":
                        print("Label name:", sub.tag, "; Text content:", 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)

This code can also convert a VOC format dataset with Chinese labels to a dataset with English labels. As shown in the figure below, you can simply swap the positions of the Chinese and English labels.

For a VOC format dataset with Chinese labels, you can use the code in the blog Object Detection --- Dataset Format Conversion and Training/Validation Set Division to convert the VOC format dataset to YOLO format and divide it into training and validation sets.

Thus, the dataset with Chinese labels is ready.

2. Modify yolov5 Code to Support Chinese Labels#

It is important to note that this blog uses version 5.0 of yolov5.

(1) Modify the train.py file, change line 63 in the py file as follows:

with open(opt.data, encoding='UTF-8') as f:

(2) Modify the test.py file, change line 73 in the py file as follows:

with open(opt.data, encoding='UTF-8') as f:

(3) Modify the utils/general.py file, import the following package in this code:

from PIL import Image, ImageDraw, ImageFont

(4) Modify the utils/plots.py file, change line 64 in the py file, modify the plot_one_box function, change the code after if label to:

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) Modify the utils/plots.py file, change line 144 in the py file, modify the plot_images function:

mosaic = plot_one_box(box, mosaic, label=label, color=color, line_thickness=tl)

(6) Modify the detect.py file, change line 178 in the py file as follows:

im0 = plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)

Note that this im0 is very important. If your image has already been framed but the label does not display text, it is likely that this im0 is missing. This im0 must correspond to the im0 in the later plot_one_box(xyxy, im0.

Thus, the code modifications to support Chinese labels are complete.

If you need to refer to the model training and inference testing work later, you can check the blog: Object Detection --- Teaching You to Train Your Own Object Detection Model Using yolov5.

It is particularly important to note that when modifying the yaml file in the data directory, the labels must correspond to the annotated Chinese categories.

Reference Links:

YoloV5 Implementation of Chinese Label Object Detection - Zhihu (zhihu.com)

YoloV5 Implementation of Chinese Label Object Detection#

1. Copy the yolov5 Project to Local#

https://github.com/ultralytics/yolov5

2. Check Fonts#

By default, it calls fonts from C:\Windows\Fonts, find a suitable one, for example, I used simhei.ttf this time.

3. `train.py File`#

with open(opt.data) as f:

Change to

with open(opt.data, encoding='UTF-8') as f:

4. test.py File#

with open(data) as f:

Change to

with open(data, encoding='UTF-8') as f:

5. utils/general.py File#

Import the package

from PIL import Image, ImageDraw, ImageFont

6. utils/plot.py File#

Modify the plot_one_box function, change the code after if label to

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)

7. utils/plot.py File#

In the plot_images function

plot_one_box(box, mosaic, label=label, color=color, line_thickness=tl)

Change to

mosaic = plot_one_box(box, mosaic, label=label, color=color, line_thickness=tl)

8. detect.py File#

plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)

Change to

im0 = plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)