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

# =========================
# 强制使用非 GUI 后端（关键！）
# =========================
import matplotlib

matplotlib.use('Agg')  # 必须在 import pyplot 之前设置


def visualize_obb(image_path, label_path, output_dir="output_visualizations"):
    """
    可视化图片及其 OBB 标签，并保存结果图像到指定目录。

    :param image_path: 图片路径
    :param label_path: 标签路径
    :param output_dir: 输出目录（自动创建）
    """
    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        print(f"❌ 无法读取图像: {image_path}")
        return

    h, w = image.shape[:2]
    print(f"✅ 正在处理图像: {os.path.basename(image_path)} | 尺寸: {w} x {h}")

    # 创建用于绘图的副本（BGR → 绘图用）
    img_draw = image.copy()

    # 读取标签
    try:
        with open(label_path, 'r') as f:
            lines = f.readlines()
    except Exception as e:
        print(f"❌ 无法读取标签文件 {label_path}: {e}")
        return

    for line in lines:
        parts = line.strip().split()
        if len(parts) < 9:
            print(f"⚠️ 跳过无效标签行: {line}")
            continue

        # 解析：class_id x1 y1 x2 y2 x3 y3 x4 y4
        try:
            points = np.array([float(x) for x in parts[1:9]]).reshape(4, 2)
        except:
            print(f"⚠️ 坐标解析失败: {line}")
            continue

        # 归一化坐标 → 像素坐标
        points[:, 0] *= w  # x
        points[:, 1] *= h  # y
        points = np.int32(points)

        # 绘制四边形（绿色）
        cv2.polylines(img_draw, [points], isClosed=True, color=(0, 255, 0), thickness=3)

        # 绘制顶点（红色圆圈）
        for (x, y) in points:
            cv2.circle(img_draw, (x, y), 6, (0, 0, 255), -1)  # 红色实心圆

    # 转为 RGB 用于 matplotlib 保存
    img_rgb = cv2.cvtColor(img_draw, cv2.COLOR_BGR2RGB)

    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    # 生成输出路径
    filename = os.path.splitext(os.path.basename(image_path))[0] + "_vis.png"
    output_path = os.path.join(output_dir, filename)

    # 使用 matplotlib 保存图像（不显示）
    plt.figure(figsize=(16, 9), dpi=100)
    plt.imshow(img_rgb)
    plt.title(f"OBB Visualization - {os.path.basename(image_path)}", fontsize=14)
    plt.axis('off')
    plt.tight_layout()
    plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
    plt.close()  # 释放内存

    print(f"✅ 可视化结果已保存: {output_path}")


def process_directory(directory):
    """
    遍历目录，处理所有图片和对应的 .txt 标签文件
    """
    print(f"🔍 开始处理目录: {directory}")
    count = 0

    for filename in os.listdir(directory):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff')):
            image_path = os.path.join(directory, filename)
            label_path = os.path.splitext(image_path)[0] + ".txt"

            if os.path.exists(label_path):
                visualize_obb(image_path, label_path)
                count += 1
            else:
                print(f"🟡 跳过 (无标签): {filename}")

    print(f"🎉 处理完成！共处理 {count} 张图像。")


# =========================
# 主程序入口
# =========================
if __name__ == "__main__":
    # 设置你的数据目录
    directory = '/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/obb4/labels'

    if not os.path.exists(directory):
        raise FileNotFoundError(f"目录不存在: {directory}")

    process_directory(directory)