#!/usr/bin/env python3
import os, cv2, numpy as np

# ========== 配置 ==========
roi_label_dir = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/label"  # 640x640 ROI 标签
original_image_dir = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/class1"
roi_coords_file = "/home/hx/yolo/zhuangtai_class_cls/roi_coordinates/1_rois.txt"  # 每行: x,y,w,h
output_label_dir = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/classdata3_orig_norm_labels"
debug_vis = True
debug_vis_dir = output_label_dir + "_vis"
# ==========================

os.makedirs(output_label_dir, exist_ok=True)
if debug_vis:
    os.makedirs(debug_vis_dir, exist_ok=True)

# ----------- 加载 ROI -----------
with open(roi_coords_file, 'r') as f:
    roi_x, roi_y, roi_w, roi_h = map(int, f.readline().strip().split(','))

# ----------- 处理标签 -----------
for label_file in os.listdir(roi_label_dir):
    if not label_file.lower().endswith('.txt'):
        continue
    base_name = os.path.splitext(label_file)[0]

    # 对应原图
    orig_img_path = None
    for ext in ('.jpg','.jpeg','.png','.bmp'):
        tmp_path = os.path.join(original_image_dir, base_name + ext)
        if os.path.exists(tmp_path):
            orig_img_path = tmp_path
            break
    if orig_img_path is None:
        print(f"[SKIP] 未找到原图: {base_name}")
        continue

    img = cv2.imread(orig_img_path)
    if img is None:
        print(f"[SKIP] 无法读取原图: {orig_img_path}")
        continue
    h_img, w_img = img.shape[:2]

    # 读取 ROI 标签
    in_path = os.path.join(roi_label_dir, label_file)
    lines = [ln.strip() for ln in open(in_path, 'r') if ln.strip()]
    out_lines = []

    for ln in lines:
        parts = ln.split()
        cls = parts[0]
        coords = list(map(float, parts[1:]))
        mapped = []
        for i in range(0, len(coords), 2):
            x_n, y_n = coords[i], coords[i+1]   # ROI 0~1
            x_abs = roi_x + x_n * roi_w
            y_abs = roi_y + y_n * roi_h
            x_abs = min(max(x_abs, 0.0), w_img-1)
            y_abs = min(max(y_abs, 0.0), h_img-1)
            mapped.append((x_abs / w_img, y_abs / h_img))
        out_line = cls + " " + " ".join(f"{x:.6f} {y:.6f}" for x, y in mapped)
        out_lines.append(out_line)

    out_path = os.path.join(output_label_dir, label_file)
    with open(out_path, 'w') as fw:
        fw.write("\n".join(out_lines))
    print(f"[OK] {out_path}")

    # 可视化
    if debug_vis and out_lines:
        vis = img.copy()
        cv2.rectangle(vis, (roi_x, roi_y), (roi_x+roi_w-1, roi_y+roi_h-1), (0,165,255), 2)
        for ol in out_lines:
            pts = np.array([[int(float(parts[i])*w_img), int(float(parts[i+1])*h_img)]
                            for parts in [ol.split()] for i in range(1, len(parts), 2)], np.int32)
            if len(pts) >= 3:
                cv2.polylines(vis, [pts], True, (0,0,255), 2)
        vis_name = os.path.join(debug_vis_dir, os.path.basename(orig_img_path))
        cv2.imwrite(vis_name, vis)

print("🎉 ROI 标签已还原到原图归一化完成！")