YOLOv5-6.x输出中文标签修改标签框的位置和大小

Posted 2022-05-16 嗜睡的篠龙

文章目录

• • 前言
  • 输出中文标签
  • 修改标签框的位置和大小

前言

1. 使用英文标签训练好的权重文件，检测输出对应的中文标签；
2. 当文本太长导致其右侧超出标签框甚至图片时，可以将文本的位置由原来的标签框左上角，改为标签框正上方。
    

输出中文标签

• 将权重文件中类别的英文标签，直接替换为对应中文标签
• 可以直接把下面这个函数，放在detect.py文件中（例如，可以直接放在def run(...)函数上面）

# 将英文标签训练好的权重文件，转换为中文标签
def en_to_ch(chinese_names, weights):
    # 读取权重文件
    weights_dict = torch.load(weights)
    # 将原来的英文标签，替换为中文标签
    weights_dict['model'].names = chinese_names
    # 最后保存到原文件中
    torch.save(weights_dict, weights)

• chinese_names：就是对应的中文标签，例如：['人类','领带']，注意这里的顺序要和原来的英文标签一致；
  • 如果只希望修改一部分为中文，另一份仍然为英文，则只需要将要修改的部分改为中文即可，例如：['person','领带']
• weights：为要修改的权重文件

效果如下：

 

修改标签框的位置和大小

• 根据需要（便于观察），如果文本太长导致其右侧超出标签框甚至图片时，可以将文本的位置改为在标签框正上方，未超出时就不用改
• 可以直接用下面的代码替换原来的box_label函数（直接copy即可）

def box_label(self, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
    # Add one xyxy box to image with label
    if self.pil or not is_ascii(label):
        self.draw.rectangle(box, width=self.lw, outline=color)  # box
        if label:
            w, h = self.font.getsize(label)  # text width, height
            outside = box[1] - h >= 0  # label fits outside box
            self.draw.rectangle((box[0],
                                 box[1] - h if outside else box[1],
                                 box[0] + w + 1,
                                 box[1] + 1 if outside else box[1] + h + 1), fill=color)
            # self.draw.text((box[0], box[1]), label, fill=txt_color, font=self.font, anchor='ls')  # for PIL>8.0
            self.draw.text((box[0], box[1] - h if outside else box[1]), label, fill=txt_color, font=self.font)
    else:  # cv2
        p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
        cv2.rectangle(self.im, p1, p2, color, thickness=self.lw, lineType=cv2.LINE_AA)
        if label:
            tf = max(self.lw - 1, 1)  # font thickness
            w, h = cv2.getTextSize(label, 0, fontScale=self.lw / 3, thickness=tf)[0]  # text width, height
            outside = p1[1] - h - 3 >= 0  # label fits outside box
            p2 = p1[0] + w, p1[1] - h - 5 if outside else p1[1] + h + 3
            # 当文本超出标签框时，改为标签框正上方
            # 未超出，仍是标签框左上角
            move_size = int((w - (box[2] - box[0])) / 2) if (box[2] - box[0]) < w else 0
            # p1: 左下角(x1, y1)  p2: 右上角(x2, y2)
            cv2.rectangle(self.im,
                          (p1[0] - move_size, p1[1]), (p2[0] - move_size, p2[1]),
                          color, -1, cv2.LINE_AA)  # filled
            # p1[0], p1[1] - 2 文本左下角坐标
            cv2.putText(self.im, label, (p1[0] - move_size, p1[1] - 3 if outside else p1[1] + h + 2), 0,
                        fontScale=self.lw / 3,
                        color=txt_color, thickness=tf, lineType=cv2.LINE_AA)

效果如下：

创作打卡挑战赛 赢取流量/现金/CSDN周边激励大奖