import os

# ----------------------------
# 配置
# ----------------------------
SOURCE_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/resize_seg2"       # 原始图片目录
TARGET_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/resize_seg3"    # 输出目录
SUBSETS = ["train", "val", "test"]

# 全局 ROI (x, y, w, h)
GLOBAL_ROI = [562, 798, 1287, 453]

# ----------------------------
# 分割标签处理函数
# ----------------------------
def adjust_seg_labels(label_path, crop_coords, orig_size):
    x1, y1, x2, y2 = crop_coords
    crop_w = x2 - x1
    crop_h = y2 - y1
    new_labels = []

    if not os.path.exists(label_path):
        print(f"⚠️ 标签源文件不存在: {label_path}")
        return new_labels

    with open(label_path, 'r') as f:
        lines = f.readlines()

    orig_w, orig_h = orig_size

    for line in lines:
        parts = line.strip().split()
        if len(parts) < 3 or len(parts) % 2 == 0:
            continue  # 至少类别 + 一个点(x,y)

        cls = parts[0]
        coords = list(map(float, parts[1:]))

        new_coords = []
        for i in range(0, len(coords), 2):
            x = coords[i] * orig_w
            y = coords[i+1] * orig_h
            x -= x1
            y -= y1
            x_new = x / crop_w
            y_new = y / crop_h
            new_coords.extend([x_new, y_new])

        new_line = cls + " " + " ".join([f"{c:.6f}" for c in new_coords])
        new_labels.append(new_line)

    return new_labels

# ----------------------------
# 主处理函数
# ----------------------------
def process_seg_labels():
    x, y, w, h = GLOBAL_ROI
    fixed_roi = (x, y, x + w, y + h)

    for subset in SUBSETS:
        src_dir = os.path.join(SOURCE_DIR, subset)
        tgt_dir = os.path.join(TARGET_DIR, subset)
        os.makedirs(tgt_dir, exist_ok=True)

        for file in os.listdir(src_dir):
            if not (file.endswith(".txt")):
                continue

            label_path = os.path.join(src_dir, file)
            # 读取原图尺寸
            img_name = os.path.splitext(file)[0] + ".jpg"
            img_path = os.path.join(src_dir, img_name)
            if not os.path.exists(img_path):
                print(f"❌ 无法读取图片以获取尺寸: {img_path}")
                continue

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

            new_labels = adjust_seg_labels(label_path, fixed_roi, (w_img, h_img))
            tgt_label_path = os.path.join(tgt_dir, file)
            with open(tgt_label_path, 'w') as f:
                f.write("\n".join(new_labels))

            print(f"✅ 标签处理完成: {subset}/{file}, 条数 {len(new_labels)}")

if __name__ == "__main__":
    process_seg_labels()