yoloV5 Ground Truth box框框的绘制

Posted 2023-03-04 屋顶上的蓝胖子

GT标签格式就是 yolo label的形式，也就是.txt格式，而且命名除了后缀与图片不同其他是一样的，具体的排放路径如下，其中LGT是主目录：

图片就全部放在images中，标签就全部放在labels中，最终的输出图片默认放在output中

有一点需要说明，就是框的颜色问题，根据yolov5的plots.py文件，颜色的顺序是这样的，具体可以查看源文件：

hex = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
               '2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')

这些颜色出现的顺序是与你训练时使用的yaml文件有关的，比如，如果yaml文件中标签的顺序是这样：

# Classes
nc: 4  # number of classes
names: ['D10', 'D00', 'D20', 'D40']  # class names

这就意味着D10标签的颜色是 FF3838，以此类推。这里颜色的表示是HEX格式，利用opencv画图传入的颜色格式是BGR！注意是BGR不是RGB！！！
几乎所有的颜色转换网站默认的颜色格式是rgb，需要自己换一下，切记！

# -- coding: utf-8 --
import os
import shutil
from pathlib import Path
import numpy as np
import cv2
from tqdm import tqdm

# 修改输入图片文件夹
img_folder = "LGT/images"
img_list = os.listdir(img_folder)
img_list.sort()
# 修改输入标签文件夹
label_folder = "LGT/labesl"
label_list = os.listdir(label_folder)
label_list.sort()
# 输出图片文件夹位置
output_folder = 'LGT/output'

labels = ['D10', 'D00', 'D20', 'D40']  # 这里修改为自己的类别

# 色盘，可根据类别添加新颜色，注意第一个别动，这是字体的颜色也就是白色，往后面推
colormap = [(255, 255, 255), (56, 56, 255), (151, 157, 255), (31, 112, 255), (29, 178, 255)]   # 不是RGB,是BGR


# 坐标转换
def xywh2xyxy(x, w1, h1, img):
    label, x, y, w, h = x

    label = int(label)
    label_ind = label

    # 边界框反归一化
    x_t = x * w1
    y_t = y * h1
    w_t = w * w1
    h_t = h * h1

    # 计算坐标
    top_left_x = x_t - w_t / 2
    top_left_y = y_t - h_t / 2
    bottom_right_x = x_t + w_t / 2
    bottom_right_y = y_t + h_t / 2

    p1, p2 = (int(top_left_x), int(top_left_y)), (int(bottom_right_x), int(bottom_right_y))
    # 绘制矩形框
    cv2.rectangle(img, p1, p2, colormap[label_ind+1], thickness=2, lineType=cv2.LINE_AA)
    label = labels[label_ind]
    if label:
        w, h = cv2.getTextSize(label, 0, fontScale=2 / 3, thickness=2)[0]  # text width, height
        outside = p1[1] - h - 3 >= 0  # label fits outside box
        p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
        # 绘制矩形框填充
        cv2.rectangle(img, p1, p2, colormap[label_ind+1], -1, cv2.LINE_AA)
        # 绘制标签
        cv2.putText(img, label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), 0, 2 / 3, colormap[0],
                    thickness=2, lineType=cv2.LINE_AA)
    return img


if __name__ == '__main__':
    # 创建输出文件夹
    if Path(output_folder).exists():
        shutil.rmtree(output_folder)
    os.mkdir(output_folder)
    # labels和images可能存在不相等的情况，需要额外判断对应关系
    img_index = 0
    label_index = 0
    for _ in tqdm(range(len(label_list))):
        image_path = img_folder + "/" + img_list[img_index]
        label_path = label_folder + "/" + label_list[label_index]
        if img_list[img_index][:-4] != label_list[label_index][:-4]:
            img_index += 1
            continue
        # 读取图像文件
        img = cv2.imread(str(image_path))
        h, w = img.shape[:2]
        # 读取 labels
        with open(label_path, 'r') as f:
            lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32)
        # 绘制每一个目标
        for x in lb:
            # 反归一化并得到左上和右下坐标，画出矩形框
            img = xywh2xyxy(x, w, h, img)
        cv2.imwrite(output_folder + '/' + '.png'.format(image_path.split('/')[-1][:-4]), img)
        img_index += 1
        label_index += 1

参考博客：颜色格式转换网站，代码主体来源