目标检测标签可视化工具(yolo、coco、voc)

通用目标检测标签可视化工具

在目标检测的任务中，我们通常需要对检测结果进行可视化，以便更好地理解模型的性能。本文将介绍一个通用的目标检测标签可视化工具，该工具支持读取VOC格式（XML）、COCO格式（JSON）和YOLO格式（TXT）的标签，并将这些标签以边界框的形式绘制在对应的图像上。

工具简介

本文提供的代码可以自动读取指定目录下的图片和标签文件，并将标签信息绘制在图片上，然后将处理后的图片保存到指定目录。代码使用了OpenCV库进行图像处理，并提供了颜色管理类Colors，可以方便地为不同类别的对象分配不同的颜色。

代码实现

依赖库

首先，需要确保安装了以下依赖库：

pip install opencv-python

代码实现

以下是实现代码：

import os
import cv2
import xml.etree.ElementTree as ET
import json

class Colors:
    def __init__(self):
        hex = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
               '2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
        self.palette = [self.hex2rgb('#' + c) for c in hex]
        self.n = len(self.palette)

    def __call__(self, i, bgr=False):
        c = self.palette[int(i) % self.n]
        return (c[2], c[1], c[0]) if bgr else c

    @staticmethod
    def hex2rgb(h):
        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):
    assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to plot_on_box() input image.'
    tl = line_thickness or round(0.002 * (im.shape[0] + im.shape[1]) / 2) + 1
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    cv2.rectangle(im, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
    if label:
        tf = max(tl - 1, 1)
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        cv2.rectangle(im, c1, c2, color, -1, cv2.LINE_AA)
        cv2.putText(im, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)

def parse_voc_annotation(xml_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    objects = []
    for obj in root.findall('object'):
        cls = obj.find('name').text
        bbox = obj.find('bndbox')
        x1 = int(bbox.find('xmin').text)
        y1 = int(bbox.find('ymin').text)
        x2 = int(bbox.find('xmax').text)
        y2 = int(bbox.find('ymax').text)
        objects.append((cls, x1, y1, x2, y2))
    return objects

def parse_coco_annotation(json_file, img_id):
    with open(json_file, 'r') as f:
        data = json.load(f)
    objects = []
    for ann in data['annotations']:
        cls_id = ann['category_id']
        bbox = ann['bbox']
        x1 = int(bbox[0])
        y1 = int(bbox[1])
        x2 = int(bbox[0] + bbox[2])
        y2 = int(bbox[1] + bbox[3])
        objects.append((cls_id, x1, y1, x2, y2))
    return objects

def parse_yolo_annotation(txt_file, img_shape):
    with open(txt_file, 'r') as f:
        lines = f.readlines()
    w, h = img_shape[1], img_shape[0]
    objects = []
    for line in lines:
        parts = line.strip().split()
        cls_id = int(parts[0])
        x_center = float(parts[1]) * w
        y_center = float(parts[2]) * h
        width = float(parts[3]) * w
        height = float(parts[4]) * h
        x1 = int(x_center - width / 2)
        y1 = int(y_center - height / 2)
        x2 = int(x_center + width / 2)
        y2 = int(y_center + height / 2)
        objects.append((cls_id, x1, y1, x2, y2))
    return objects

def show_label(path_root_imgs, path_root_labels, save_path_root_imgs):
    class_categories = ['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)

    colors = Colors()

    for root, _, files in os.walk(path_root_imgs):
        for file in files:
            if file.endswith('.jpg'):
                img_path = os.path.join(root, file)
                label_path = os.path.join(path_root_labels, file.replace('.jpg', '.txt'))
                save_path = os.path.join(save_path_root_imgs, file)

                img = cv2.imread(img_path)
                if img is None:
                    continue

                objects = []
                if os.path.exists(label_path):
                    objects = parse_yolo_annotation(label_path, img.shape)
                else:
                    xml_path = os.path.join(path_root_labels, file.replace('.jpg', '.xml'))
                    json_path = os.path.join(path_root_labels, file.replace('.jpg', '.json'))
                    if os.path.exists(xml_path):
                        objects = parse_voc_annotation(xml_path)
                    elif os.path.exists(json_path):
                        img_id = int(file.split('.')[0])
                        objects = parse_coco_annotation(json_path, img_id)

                img_tmp = img.copy()
                for obj in objects:
                    if isinstance(obj[0], int):
                        label = class_categories[obj[0]]
                    else:
                        label = obj[0]
                    plot_one_box(obj[1:], img_tmp, color=colors(obj[0]), label=label)

                cv2.imwrite(save_path, img_tmp)
                print(f"Saved: {save_path}")

if __name__ == '__main__':
    path_root_labels = r'路径到标签文件夹'
    path_root_imgs = r'路径到图片文件夹'
    save_path_root_imgs = r'保存结果的文件夹路径'
    
    show_label(path_root_imgs, path_root_labels, save_path_root_imgs)

代码使用说明

准备工作

1. 安装依赖：确保安装了OpenCV库，使用以下命令安装：

   pip install opencv-python
   

2. 文件组织：将图片文件和标签文件分别放置在两个不同的文件夹中。

参数设置

在运行代码前，需要根据自己的文件路径进行参数设置：

• path_root_labels：标签文件夹路径。
• path_root_imgs：图片文件夹路径。
• save_path_root_imgs：保存结果的文件夹路径。

运行代码

配置好路径后，直接运行代码即可：

python your_script_name.py

功能说明

• Colors类：用于管理颜色的分配，每个类别的对象会分配一个唯一的颜色。
• plot_one_box函数：在图像上绘制一个边界框，并在边界框上显示标签信息。
• parse_voc_annotation函数：解析VOC格式（XML）的标签文件。
• parse_coco_annotation函数：解析COCO格式（JSON）的标签文件。
• parse_yolo_annotation函数：解析YOLO格式（TXT）的标签文件。
• show_label函数：读取指定目录下的图片和标签文件，调用上述函数解析标签并绘制在图片上，最后将结果保存到指定目录。