yolov5檢測框顯示中文標籤

yolov5 檢測框顯示中文標籤#

注意：本註解對應 yolov5 的 5.0 版本！！！#

前言#

很多人在訓練yolov5目標檢測的時候，標籤只能顯示英文的。怎麼樣才可以訓練一個可以檢測物體並且顯示中文標籤的模型呢。下面我們來一步一步的做。

1、有中文標籤的數據集 #

首先在收集數據集的時候，打部分公開的數據集，標籤都是英文的，格式都是 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 數據集，並劃分為訓練集和驗證集。

至此中文標籤的數據集就準備好了

2、yolov5 代碼修改為支持中文標籤#

特別要說明一下，該博客所用的是 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 文件中的 144 行，修改 plot_images 函數：

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

（6）detect.py 文件修改，在 py 文件中的 178 行，代碼修改如下：

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

注意，這個 im0 非常重要，如果你的圖片已經框出來了，但是標籤上不顯示字，那麼很有可能就是這個im0= 缺失，這個im0與後面plot_one_box(xyxy, im0裡的 im0 要對應

至此為了支持中文標籤的代碼就修改完了。

如果後面的模型訓練和推理測試工作可以參考博客：目標檢測 --- 教你利用 yolov5 訓練自己的目標檢測模型。

這裡特別要說明一下，就是修改 data 目錄下的 yaml 文件的時候，那個標籤也一定要對應標註的中文類別。

參考鏈接：（下文）

YoloV5 實現中文標籤目標檢測 - 知乎 (zhihu.com)

YoloV5 實現中文標籤目標檢測#

1、複製 yolov5 工程到本地#

https://github.com/ultralytics/yolov5

2、查看字體#

默認調用C:\Windows\Fonts 下的字體，找個合適的，比如我這次使用simhei.ttf 。

3、`train.py文件`#

with open(opt.data) as f:

改為

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

4、test.py 文件#

with open(data) as f:

改為

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

5、utils/general.py 文件#

導包

from PIL import Image, ImageDraw, ImageFont

6、utils/plot.py 文件#

修改plot_one_box 函數，if label之後的代碼改為

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 文件#

plot_images函數中

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

改為

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

8、detect.py 文件#

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

改為

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