import os
import cv2
import numpy as np
from collections import Counter, defaultdict
from PIL import Image

# ※ 需要修改输入输出路径 ※ #
input_dir = 'Data/ann_dir'
output_dir_1 = 'Data/labels/train'
output_dir_2 = 'Data/labels/train'

os.makedirs(output_dir_1, exist_ok=True)
os.makedirs(output_dir_2, exist_ok=True)

# 全局统计颜色频率与 class 像素频率
global_color_counter = Counter()
global_class_counter = Counter()
color_class_counter = defaultdict(int)  # (R,G,B) → count
color_to_old_class = {}  # (R,G,B) → class_id
remap_class_dict = {}  # old_class_id → new_class_id

# 自动提取颜色映射（跳过背景）
def extract_color_mapping(img_path):
    img = Image.open(img_path).convert('RGB')
    pixels = list(img.getdata())
    counter = Counter(pixels)
    color_map = {}
    for color, count in counter.items():
        global_color_counter[color] += count
        if color != (0, 0, 0) and color[0] == color[1] == color[2]:
            class_id = color[0] - 1
            if class_id >= 0:
                color_map[color] = class_id
                global_class_counter[class_id] += count
                color_class_counter[color] += count
                color_to_old_class[color] = class_id
    return color_map

# 处理单张图片
def process_image(img_path, save_path_list):
    color_to_class = extract_color_mapping(img_path)
    if not color_to_class:
        print(f"[跳过] {os.path.basename(img_path)} 无有效目标")
        return

    img = cv2.imread(img_path)
    h, w = img.shape[:2]
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    lines = []
    for rgb, old_class_id in color_to_class.items():
        mask = np.all(img_rgb == rgb, axis=-1).astype(np.uint8) * 255
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for contour in contours:
            if contour.shape[0] < 3:
                continue
            norm_pts = contour.squeeze(1).astype(np.float32)
            norm_pts[:, 0] /= w
            norm_pts[:, 1] /= h
            flat = norm_pts.flatten()
            new_class_id = remap_class_dict.get(old_class_id, old_class_id)
            line = f"{new_class_id} " + " ".join([f"{x:.6f}" for x in flat])
            lines.append(line)

    if lines:
        for save_path in save_path_list:
            with open(save_path, 'w') as f:
                for line in lines:
                    f.write(line + "\n")
            print(f"[✔] 转换成功: {os.path.basename(save_path)}，共 {len(lines)} 个实例")
    else:
        print(f"[⚠] {os.path.basename(img_path)} 没有轮廓")

# 第一次遍历图像，仅提取颜色信息用于构建 class 映射
for fname in os.listdir(input_dir):
    if fname.lower().endswith('.png'):
        img_path = os.path.join(input_dir, fname)
        extract_color_mapping(img_path)

# 构建 class_id 重映射表（按像素数从大到小排序）
class_pixel_count = {color_to_old_class[c]: count for c, count in color_class_counter.items()}
sorted_classes = sorted(class_pixel_count.items(), key=lambda x: x[1], reverse=True)
remap_class_dict = {old_cls: new_idx for new_idx, (old_cls, _) in enumerate(sorted_classes)}

# 第二次遍历图像，正式处理并生成标签
for fname in os.listdir(input_dir):
    if fname.lower().endswith('.png'):
        img_path = os.path.join(input_dir, fname)
        base_name = os.path.splitext(fname)[0]
        txt_path_1 = os.path.join(output_dir_1, base_name + '.txt')
        txt_path_2 = os.path.join(output_dir_2, base_name + '.txt')
        process_image(img_path, [txt_path_1, txt_path_2])

# 打印颜色统计和类别映射
print("\n📊 所有图像颜色统计：")
for color, count in global_color_counter.most_common():
    if color == (0, 0, 0):
        print(f"背景颜色 {color} 出现次数: {count}")
    elif color[0] == color[1] == color[2]:
        class_id = color[0] - 1
        print(f"颜色 {color} → class {class_id}，出现次数: {count}")
    else:
        print(f"[⚠] 非灰阶颜色 {color}，跳过")

print("\n✅ 有效类别统计（原 class_id → 新 class_id → 总像素数）：")
for old_id, new_id in remap_class_dict.items():
    print(f"class {old_id} → {new_id}: {class_pixel_count[old_id]} pixels")

print(f"\n✅ 全部图像处理完毕。标签输出目录：{output_dir_1}、{output_dir_2}")
