import os
import cv2

# ----------------------------
# 配置
# ----------------------------
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"    # 输出目录
TARGET_SIZE = 640           # resize 尺寸
SUBSETS = ["train", "val", "test"]

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

# ----------------------------
# 主处理函数
# ----------------------------
def process_images():
    x, y, w, h = GLOBAL_ROI
    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(".jpg") or file.endswith(".png")):
                continue

            img_path = os.path.join(src_dir, file)
            img = cv2.imread(img_path)
            if img is None:
                print(f"❌ 无法读取图片: {img_path}")
                continue

            h_img, w_img = img.shape[:2]
            x1, y1 = max(0, x), max(0, y)
            x2, y2 = min(w_img, x + w), min(h_img, y + h)

            cropped = img[y1:y2, x1:x2]
            if cropped.size == 0:
                print(f"❌ 裁剪结果为空: {file}")
                continue

            resized = cv2.resize(cropped, (TARGET_SIZE, TARGET_SIZE))
            tgt_path = os.path.join(tgt_dir, file)
            cv2.imwrite(tgt_path, resized)
            print(f"✅ 图片处理完成: {subset}/{file}")

if __name__ == "__main__":
    process_images()