目标检测-可视化YOLO格式标签

import os
import cv2

class Colors:
    # Ultralytics color palette https://ultralytics.com/
    def __init__(self):
        # hex = matplotlib.colors.TABLEAU_COLORS.values()
        hex = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
               '2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
        # 将hex列表中所有hex格式(十六进制)的颜色转换rgb格式的颜色
        self.palette = [self.hex2rgb('#' + c) for c in hex]
        # 颜色个数
        self.n = len(self.palette)

    def __call__(self, i, bgr=False):
        # 根据输入的index 选择对应的rgb颜色
        c = self.palette[int(i) % self.n]
        # 返回选择的颜色 默认是rgb
        return (c[2], c[1], c[0]) if bgr else c

    @staticmethod
    def hex2rgb(h):  # rgb order (PIL)
        # hex -> rgb
        return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))

def plot_one_box(x, im, color=(128, 128, 128), label=None, line_thickness=3):
    """一般会用在detect.py中在nms之后变量每一个预测框，再将每个预测框画在原图上
    使用opencv在原图im上画一个bounding box
    :params x: 预测得到的bounding box  [x1 y1 x2 y2]
    :params im: 原图 要将bounding box画在这个图上  array
    :params color: bounding box线的颜色
    :params labels: 标签上的框框信息  类别 + score
    :params line_thickness: bounding box的线宽
    """
    # check im内存是否连续
    assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to plot_on_box() input image.'
    # tl = 框框的线宽  要么等于line_thickness要么根据原图im长宽信息自适应生成一个
    tl = line_thickness or round(0.002 * (im.shape[0] + im.shape[1]) / 2) + 1  # line/font thickness
    # c1 = (x1, y1) = 矩形框的左上角   c2 = (x2, y2) = 矩形框的右下角
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    # cv2.rectangle: 在im上画出框框   c1: start_point(x1, y1)  c2: end_point(x2, y2)
    # 注意: 这里的c1+c2可以是左上角+右下角  也可以是左下角+右上角都可以
    cv2.rectangle(im, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
    # 如果label不为空还要在框框上面显示标签label + score
    if label:
        tf = max(tl - 1, 1)  # label字体的线宽 font thickness
        # cv2.getTextSize: 根据输入的label信息计算文本字符串的宽度和高度
        # 0: 文字字体类型  fontScale: 字体缩放系数  thickness: 字体笔画线宽
        # 返回retval 字体的宽高 (width, height), baseLine 相对于最底端文本的 y 坐标
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        # 同上面一样是个画框的步骤  但是线宽thickness=-1表示整个矩形都填充color颜色
        cv2.rectangle(im, c1, c2, color, -1, cv2.LINE_AA)  # filled
        # cv2.putText: 在图片上写文本 这里是在上面这个矩形框里写label + score文本
        # (c1[0], c1[1] - 2)文本左下角坐标  0: 文字样式  fontScale: 字体缩放系数
        # [225, 255, 255]: 文字颜色  thickness: tf字体笔画线宽     lineType: 线样式
        cv2.putText(im, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)

def show_label(path_root_imgs,path_root_labels, save_path_root_imgs):
    #需修改，根据自己的类别，注意一一对应
    class_catagories=['Abrasion', 'Crazing', 'Patches', 'Inclusion', 'Uneven', 'Blowhole', 'Break', 'Crack', 'Crescent_Gap', 'Crease', 'Silk-Spot', 'Water-Spot', 'Weld-Line', 'GC-Inclusion', 'Oil-Spot', 'Rolled-Pit', 'Punching', 'Waist-Folding', 'Bruise', 'Pitted_Surface', 'Rolled-in_Scale', 'Scratches', 'Bubble']

    if not os.path.exists(save_path_root_imgs):
        os.makedirs(save_path_root_imgs)
    type_object = '.txt'
    for ii in os.walk(path_root_imgs):
        for j in ii[2]:
            save_path=save_path_root_imgs
            type = j.split(".")[1]
            if type != 'jpg':
                continue
            save_path=os.path.join(save_path,j)
            path_img = os.path.join(path_root_imgs, j)
            label_name = j[:-4]+type_object
            path_label = os.path.join(path_root_labels, label_name)
            f = open(path_label, 'r+', encoding='utf-8')
            if os.path.exists(path_label) == True:

                img = cv2.imread(path_img)
                w = img.shape[1]
                h = img.shape[0]
                img_tmp = img.copy()
                while True:
                    line = f.readline()
                    if line:
                        temp_list=[]
                        msg = line.split(" ")
                        label_index=int(msg[0])
                        x1 = int((float(msg[1]) - float(msg[3]) / 2) * w)  # x_center - width/2
                        y1 = int((float(msg[2]) - float(msg[4]) / 2) * h)  # y_center - height/2
                        x2 = int((float(msg[1]) + float(msg[3]) / 2) * w)  # x_center + width/2
                        y2 = int((float(msg[2]) + float(msg[4]) / 2) * h)  # y_center + height/2
                        temp_list.append(x1)
                        temp_list.append(y1)
                        temp_list.append(x2)
                        temp_list.append(y2)
                        plot_one_box(x=temp_list,im=img_tmp,color=colors(int(label_index)),label=class_catagories[label_index])
                    else :
                        break
            cv2.imwrite(save_path,img_tmp)
            print(save_path)
            print("succfully")

if __name__ == '__main__':
    # 初始化Colors对象
    colors = Colors()
    # 需修改，检测标签根目录
    path_root_labels = r'D:\desk\asssssssssssssss\txt'
    # 需修改，检测图片根目录
    path_root_imgs = r'D:\desk\asssssssssssssss\img'
    # 需修改，结果保存根目录
    save_path_root_imgs=r'D:\desk\asssssssssssssss\results'

    show_label(path_root_imgs,path_root_labels,save_path_root_imgs)

 需要修改部分

在以下链接的基础上，增加了类别显示功能，仅需修改以下几部分即可使用本代码

（1）第60行

根据自己的类别替换列表

    #需修改，根据自己的类别，注意一一对应

    class_catagories=['Cahua', 'Crazing', 'Patches', 'Inclusion']

（2）第107行

替换为自己的标签根目录

    # 需修改，检测标签根目录
    path_root_labels = r'D:\desk\asssssssssssssss\txt'

（3）第109行

替换为自己的图片根目录

    # 需修改，检测图片根目录
    path_root_imgs = r'D:\desk\asssssssssssssss\img'

（4）第111行

可视化的结果保存根目录

    # 需修改，结果保存根目录
    save_path_root_imgs=r'D:\desk\asssssssssssssss\results'

---

 实验准备

 两张图片

两个txt标签

实验结果

---

参考代码：

【YOLOV5-5.x 源码解读】plots.py_yolov5找不到plot_one_box-CSDN博客

【python】在原图中显示标签(yolo格式)的检测框bbox_python 图片中相框检测_石头城修道的博客-CSDN博客

 本代码主要借鉴了yolov5中的plot.py中的函数，可以根据不同的类别生成不同颜色的标签，并且会根据图片的大小自适应调整标签大小，相比自己定义的标签格式画得更好看些！

本代码仅供本人学习使用