import cv2
import os
import random

# ================= 配置区域 =================
# 1. 图片文件夹路径
IMG_DIR = r'20251204\images' 

# 2. 标签文件夹路径 (YOLO格式的txt)
LABEL_DIR = r'20251204\labels'

# 3. 结果保存路径 (如果不填，默认在当前目录下生成 'output' 文件夹)
OUTPUT_DIR = r'output_result'

# 4. 类别名称列表 (按 classes.txt 的顺序或者是你的模型定义的顺序)
# 如果不知道名字，可以留空，代码会直接显示 ID (如 '0', '1')
CLASSES = ['fire', 'fust'] 
# CLASSES = [] # 如果不想显示名字，保持为空列表即可

# 5. 是否只显示有标签的图片？(True: 跳过无标签图片; False: 显示所有)
ONLY_SHOW_LABELED = True
# ===========================================

def get_color(cls_id):
    """为每个类别生成固定的随机颜色"""
    random.seed(cls_id * 10) # 保证同一个类别的颜色在不同图片中是一样的
    b = random.randint(0, 255)
    g = random.randint(0, 255)
    r = random.randint(0, 255)
    return (b, g, r)

def plot_one_box(x, img, color=None, label=None, line_thickness=None):
    """画框的主函数"""
    tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1  # 线宽
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    cv2.rectangle(img, 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(img, c1, c2, color, -1, cv2.LINE_AA)  # 填充文字背景
        cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)

def yolo_to_xyxy(yolo_line, img_w, img_h):
    """将YOLO格式 (class x_center y_center w h) 转换为 (x1, y1, x2, y2)"""
    parts = yolo_line.strip().split()
    cls_id = int(parts[0])
    x_c = float(parts[1])
    y_c = float(parts[2])
    w = float(parts[3])
    h = float(parts[4])

    # 还原为像素坐标
    x1 = int((x_c - w / 2) * img_w)
    y1 = int((y_c - h / 2) * img_h)
    x2 = int((x_c + w / 2) * img_w)
    y2 = int((y_c + h / 2) * img_h)
    
    # 边界保护，防止画出图片外
    x1 = max(0, x1)
    y1 = max(0, y1)
    x2 = min(img_w, x2)
    y2 = min(img_h, y2)

    return cls_id, x1, y1, x2, y2

def main():
    if not os.path.exists(OUTPUT_DIR):
        os.makedirs(OUTPUT_DIR)
    
    print(f"开始处理... 结果将保存到: {OUTPUT_DIR}")
    
    img_files = [f for f in os.listdir(IMG_DIR) if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp'))]
    
    for img_file in img_files:
        img_path = os.path.join(IMG_DIR, img_file)
        label_file = os.path.splitext(img_file)[0] + '.txt'
        label_path = os.path.join(LABEL_DIR, label_file)
        
        # 读取图片
        img = cv2.imread(img_path)
        if img is None:
            continue
        h, w, _ = img.shape
        
        has_label = False
        
        # 检查是否有对应的txt文件
        if os.path.exists(label_path):
            with open(label_path, 'r') as f:
                lines = f.readlines()
                
            for line in lines:
                if not line.strip(): continue
                has_label = True
                
                # 解析坐标
                try:
                    cls_id, x1, y1, x2, y2 = yolo_to_xyxy(line, w, h)
                    
                    # 获取颜色
                    color = get_color(cls_id)
                    
                    # 获取标签名称
                    if CLASSES and 0 <= cls_id < len(CLASSES):
                        label_name = f"{CLASSES[cls_id]}"
                    else:
                        label_name = f"Class {cls_id}"
                    
                    # 画框
                    plot_one_box([x1, y1, x2, y2], img, color=color, label=label_name)
                except Exception as e:
                    print(f"解析错误 {label_file}: {e}")

        # 决定是否保存
        if ONLY_SHOW_LABELED and not has_label:
            continue
            
        save_path = os.path.join(OUTPUT_DIR, "vis_" + img_file)
        cv2.imwrite(save_path, img)
        # 如果想实时看，取消下面两行的注释（按任意键看下一张，按q退出）
        # cv2.imshow('Check', img)
        # if cv2.waitKey(0) == ord('q'): break

    print("处理完成！")
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()