按yolo格式标签可视化图像

import os
import cv2
import numpy as np
from glob import glob
import matplotlib.pyplot as plt

def visualize_classified_annotations(classified_root, max_per_class=5):
    """
    可视化分类后的标注数据（直接使用文件夹名称作为类名）
    
    参数:
        classified_root: 分类后的根目录（结构：类别文件夹/images/, labels/）
        max_per_class: 每类最大可视化样本数
    """
    # 获取所有类别文件夹（排序保证顺序一致）
    class_folders = sorted([d for d in os.listdir(classified_root) 
                         if os.path.isdir(os.path.join(classified_root, d))])
    
    # 定义颜色（固定8种鲜艳颜色循环使用）
    colors = [
        (255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0),
        (255, 0, 255), (0, 255, 255), (255, 128, 0), (128, 0, 255)
    ]
    
    # 处理每个类别
    for class_idx, class_folder in enumerate(class_folders):
        # 构建路径
        img_dir = os.path.join(classified_root, class_folder, "images")
        label_dir = os.path.join(classified_root, class_folder, "labels")
        vis_dir = os.path.join(classified_root, class_folder, "visualizations")
        
        # 创建visualizations目录
        os.makedirs(vis_dir, exist_ok=True)
        
        print(f"\n处理类别: {class_folder}")
        print(f"图像目录: {img_dir}")
        print(f"标签目录: {label_dir}")
        
        # 获取图像文件列表（限制数量）
        img_files = []
        for ext in ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tif']:
            img_files.extend(glob(os.path.join(img_dir, ext)))
        img_files = sorted(img_files)[:max_per_class]
        
        # 处理每个图像
        for img_path in img_files:
            base_name = os.path.splitext(os.path.basename(img_path))[0]
            label_path = os.path.join(label_dir, f"{base_name}.txt")
            output_path = os.path.join(vis_dir, f"{base_name}_annotated.jpg")
            
            # 检查标签文件
            if not os.path.exists(label_path):
                print(f"警告: 缺失标签 {os.path.basename(label_path)}")
                continue
            
            # 读取并绘制标注
            img = cv2.imread(img_path)
            if img is None:
                print(f"错误: 无法读取图像 {os.path.basename(img_path)}")
                continue
            
            h, w = img.shape[:2]
            img_display = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
            with open(label_path, 'r') as f:
                for line in f:
                    parts = line.strip().split()
                    if len(parts) < 5: continue
                    
                    try:
                        class_id, xc, yc, bw, bh = map(float, parts[:5])
                        x1 = int((xc - bw/2) * w)
                        y1 = int((yc - bh/2) * h)
                        x2 = int((xc + bw/2) * w)
                        y2 = int((yc + bh/2) * h)
                        
                        # 确保坐标不越界
                        x1, y1 = max(0, x1), max(0, y1)
                        x2, y2 = min(w-1, x2), min(h-1, y2)
                        
                        # 绘制边界框和标签
                        color = colors[int(class_id) % len(colors)]
                        cv2.rectangle(img_display, (x1, y1), (x2, y2), color, 2)
                        cv2.putText(img_display, class_folder, (x1, y1-10),
                                  cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)
                    except:
                        continue
            
            # 保存可视化结果
            plt.figure(figsize=(12, 8))
            plt.imshow(img_display)
            plt.title(f"{class_folder} | {os.path.basename(img_path)}")
            plt.axis('off')
            plt.savefig(output_path, bbox_inches='tight', dpi=100)
            plt.close()
            
            print(f"生成: {os.path.relpath(output_path, classified_root)}")

if __name__ == "__main__":
    visualize_classified_annotations(
        classified_root="path/to/your/classified_data",
        max_per_class=3  # 每类可视化3个样本
    )