【YOLO脚本】yolo格式数据集可视化检测脚本

变量记录两个路径，图片路径和标签路径。

随机抽取三个图片，并将标签的框可视化出来，可视化图片上方有对应的图片名。
效果:

import os
import random
import cv2
import matplotlib.pyplot as plt

# 设置中文字体，确保中文正常显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]

def visualize_yolo_dataset(image_dir, label_dir, num_samples=3):
    """
    可视化YOLO格式的目标检测数据集
    
    参数:
    image_dir (str): 图片文件夹路径
    label_dir (str): 标签文件夹路径
    num_samples (int): 随机抽取的样本数量
    """
    # 获取所有图片文件
    image_files = [f for f in os.listdir(image_dir) if f.endswith(('.jpg', '.jpeg', '.png', '.bmp'))]
    
    # 如果没有图片，直接返回
    if not image_files:
        print("未找到图片文件")
        return
    
    # 随机选择指定数量的图片
    if len(image_files) < num_samples:
        num_samples = len(image_files)
    selected_images = random.sample(image_files, num_samples)
    
    # 创建子图
    fig, axes = plt.subplots(1, num_samples, figsize=(5 * num_samples, 5))
    if num_samples == 1:
        axes = [axes]
    
    # 遍历选中的图片
    for i, img_file in enumerate(selected_images):
        # 读取图片
        img_path = os.path.join(image_dir, img_file)
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # 转换颜色空间
        
        # 获取对应的标签文件
        label_file = os.path.splitext(img_file)[0] + '.txt'
        label_path = os.path.join(label_dir, label_file)
        
        # 读取标签
        if os.path.exists(label_path):
            with open(label_path, 'r') as f:
                lines = f.readlines()
            
            # 获取图像尺寸
            h, w, _ = img.shape
            
            # 绘制边界框
            for line in lines:
                parts = line.strip().split()
                if len(parts) < 5:  # 确保标签格式正确
                    continue
                
                # 解析标签数据
                class_id = int(parts[0])
                x_center = float(parts[1])
                y_center = float(parts[2])
                width = float(parts[3])
                height = float(parts[4])
                
                # 转换为像素坐标
                x_min = int((x_center - width / 2) * w)
                y_min = int((y_center - height / 2) * h)
                x_max = int((x_center + width / 2) * w)
                y_max = int((y_center + height / 2) * h)
                
                # 绘制边界框
                cv2.rectangle(img, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
                
                # 绘制类别标签
                cv2.putText(img, f'Class {class_id}', (x_min, y_min - 10),
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        
        # 显示图片和标题
        axes[i].imshow(img)
        axes[i].set_title(img_file)
        axes[i].axis('off')
    
    plt.tight_layout()
    plt.show()

# 使用示例
if __name__ == "__main__":
    # 请修改为你的图片和标签路径
    IMAGE_DIR = "path/to/images"
    LABEL_DIR = "path/to/labels"
    
    visualize_yolo_dataset(IMAGE_DIR, LABEL_DIR)