import os
import cv2

# ----------------------------
# 配置
# ----------------------------
SOURCE_ROOT_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/11c"       # 原始图片根目录
TARGET_ROOT_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/11cc"    # 输出根目录
CLASSES = ["class0", "class1", "class2", "class3", "class4"]  # 类别列表
TARGET_SIZE = 640           # resize 尺寸
SUBSETS = ["train", "val", "test"]

# ----------------------------
# 全局 ROI (x, y, w, h)
# ----------------------------
#GLOBAL_ROI = [445,540,931,319]#cls_yiliao
GLOBAL_ROI = [604,182,594,252]
# ----------------------------
# 主处理函数
# ----------------------------
def process_images():
    x, y, w, h = GLOBAL_ROI
    for subset in SUBSETS:
        for class_dir in CLASSES:
            src_dir = os.path.join(SOURCE_ROOT_DIR, subset, class_dir)
            tgt_dir = os.path.join(TARGET_ROOT_DIR, subset, class_dir)
            os.makedirs(tgt_dir, exist_ok=True)

            if not os.path.exists(src_dir):
                print(f"警告: 源目录 {src_dir} 不存在，跳过")
                continue

            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}/{class_dir}/{file}")

if __name__ == "__main__":
    process_images()