#!/usr/bin/env python3
"""
批量重新生成 CCPD_SPE 数据集 — 全量版
SPE 9类 × 8000 (ccpd_base) = 72,000 张
蓝绿由独立脚本生成 150K+100K
"""

import os
import sys
import time
from pathlib import Path

ICE_DATASET = Path(r'C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset')
os.chdir(str(ICE_DATASET))

sys.path.insert(0, str(ICE_DATASET / 'CCPD_SPE_GEN' / 'cv_pipeline'))
from plate_swap import PlateSwapPipeline

# 只生成 9 类特殊车牌 (蓝绿由独立脚本处理)
SPE_TYPES = ['police','army','gangao','shiguan_old','shiguan_new',
             'lingguan_old','lingguan_new','coach','new_energy']

COUNT_PER_TYPE = 8000  # 每类 8000, 全部用 ccpd_base 背景

CCPD2019_DIR = ICE_DATASET / 'CCPD' / 'CCPD2019'
SPE_OUTPUT = ICE_DATASET / 'CCPD_SPE'

SEED = 42
WORKERS = 24

def main():
    total_start = time.time()
    total_images = 0

    import shutil
    if SPE_OUTPUT.exists():
        print(f"删除旧 CCPD_SPE: {SPE_OUTPUT}")
        shutil.rmtree(str(SPE_OUTPUT))
    SPE_OUTPUT.mkdir(parents=True, exist_ok=True)

    for plate_type in SPE_TYPES:
        print(f"\n{'='*60}")
        print(f"[{plate_type}] {COUNT_PER_TYPE} 张")
        print(f"{'='*60}")

        pipeline = PlateSwapPipeline(plate_type=plate_type)
        out_dir = SPE_OUTPUT / 'ccpd_base' / plate_type
        out_dir.mkdir(parents=True, exist_ok=True)

        print(f"  生成中...", end=' ', flush=True)
        t0 = time.time()
        pipeline.process_batch(
            ccpd_dirs=[str(CCPD2019_DIR / 'ccpd_base')],
            output_dir=str(out_dir),
            num_images=COUNT_PER_TYPE,
            preview=False,
            seed=SEED,
            workers=WORKERS,
        )
        elapsed = time.time() - t0
        total_images += COUNT_PER_TYPE
        print(f"完成 ({elapsed:.0f}s, {COUNT_PER_TYPE/elapsed:.0f} 张/s)")

    total_t = time.time() - total_start
    print(f"\n{'='*60}")
    print(f"SPE 完成! {total_images} 张, {total_t/60:.1f}min")
    print(f"{'='*60}")

if __name__ == '__main__':
    main()