import os
import shutil
from pathlib import Path

# ================== 配置 ==================
SOURCE_DIR = "/home/dy/dataset/camera0211/d"  # 原始图片文件夹
OUTPUT_A = "/home/dy/dataset/camera0211/a"  # 前一半
OUTPUT_B = "/home/dy/dataset/camera0211/b"  # 后一半

# 支持的图片扩展名
IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}

# ================== 创建输出目录 ==================
os.makedirs(OUTPUT_A, exist_ok=True)
os.makedirs(OUTPUT_B, exist_ok=True)

# ================== 获取并排序图片 ==================
source_path = Path(SOURCE_DIR)
if not source_path.exists():
    raise FileNotFoundError(f"源文件夹不存在: {SOURCE_DIR}")

image_files = sorted([
    f for f in source_path.iterdir()
    if f.is_file() and f.suffix.lower() in IMG_EXTENSIONS
])

total = len(image_files)
if total == 0:
    print("⚠️ 源文件夹中没有找到图片！")
    exit()

print(f"共找到 {total} 张图片")

# 计算分割点：前一半（向上取整）
mid = (total + 1) // 2  # 奇数时 a 多一张；若想 b 多一张，用 total // 2

first_half = image_files[:mid]
second_half = image_files[mid:]

print(f"→ a 文件夹: {len(first_half)} 张")
print(f"→ b 文件夹: {len(second_half)} 张")

# ================== 移动文件 ==================
for img_path in first_half:
    dest = os.path.join(OUTPUT_A, img_path.name)
    shutil.move(img_path, dest)
    print(f"移动到 a: {img_path.name}")

for img_path in second_half:
    dest = os.path.join(OUTPUT_B, img_path.name)
    shutil.move(img_path, dest)
    print(f"移动到 b: {img_path.name}")

print("\n✅ 平分移动完成！")