"""
CCPD2020 绿牌真图 → YOLO EXP_CUT + 角点TRANS处理
输出: ICE_Dataset/CCPD_GREEN_TRANS/train/
管线: YOLO检出→12%/18%外扩→裁切→角点重映射→warp (与 REAL_TRANS 一致)
"""
import os, sys, cv2
import numpy as np
from pathlib import Path

ICEv2 = Path(r"C:\Users\Kirl\Desktop\ICEv2")
OCR_TOOLS = ICEv2 / "ICE_NNv2/ICE-OCR/script/tools"
sys.path.insert(0, str(OCR_TOOLS))
from dataset import imread_unicode

SRC_DIR = ICEv2 / "ICE_Dataset/CCPD/CCPD2020/ccpd_green"
OUT_DIR = ICEv2 / "ICE_Dataset/CCPD_GREEN_TRANS/train"

YOLO_PATH = str(ICEv2 / "ICE_NNv2/yolov5t/output/1_CCPD_19K/weights/best.pt")
EXPAND_NORMAL, EXPAND_SMALL = 1.12, 1.18
MIN_PAD = 2

def process_one(img_path, model, img_size=640):
    img = imread_unicode(str(img_path))
    if img is None: return None
    h, w = img.shape[:2]
    results = model(img_path, imgsz=img_size, verbose=False)
    if len(results) == 0 or results[0].boxes is None or len(results[0].boxes) == 0:
        return None
    boxes = results[0].boxes
    best_idx = boxes.conf.argmax()
    x1, y1, x2, y2 = boxes.xyxy[best_idx].cpu().numpy().astype(int)
    bw, bh = x2 - x1, y2 - y1
    expand = EXPAND_SMALL if (bw < 16 or bh < 16) else EXPAND_NORMAL
    px = max(MIN_PAD, int(round(bw * (expand - 1.0) / 2.0)))
    py = max(MIN_PAD, int(round(bh * (expand - 1.0) / 2.0)))
    cx1, cy1 = max(0, x1 - px), max(0, y1 - py)
    cx2, cy2 = min(w, x2 + px), min(h, y2 + py)
    if cx2 <= cx1 or cy2 <= cy1: return None
    cropped = img[cy1:cy2, cx1:cx2]
    corners_orig = np.float32([[x1,y1],[x2,y1],[x2,y2],[x1,y2]])
    corners_crop = (corners_orig - np.array([cx1, cy1])).astype(np.float32)
    dst = np.float32([[0,0],[191,0],[191,63],[0,63]])
    M = cv2.getPerspectiveTransform(corners_crop, dst)
    return cv2.warpPerspective(cropped, M, (192, 64))

# --- 收集所有绿牌 ---
all_files = []
for root, dirs, files in os.walk(str(SRC_DIR)):
    for f in files:
        if f.lower().endswith(".jpg"):
            all_files.append(Path(root) / f)
print(f"[1/2] Found {len(all_files)} green plate images in {SRC_DIR}")

# --- YOLO + warp ---
from ultralytics import YOLO
model = YOLO(YOLO_PATH)
OUT_DIR.mkdir(parents=True, exist_ok=True)
done = err = 0
for i, img_path in enumerate(all_files):
    try:
        out_path = OUT_DIR / f"{img_path.stem}_trans.jpg"
        if out_path.exists(): done += 1; continue
        warped = process_one(img_path, model)
        if warped is None: err += 1; continue
        ret, enc = cv2.imencode(".jpg", warped, [cv2.IMWRITE_JPEG_QUALITY, 95])
        if ret:
            with open(str(out_path), "wb") as wf: wf.write(enc.tobytes())
            done += 1
        else: err += 1
    except Exception as e:
        err += 1
        if err <= 3: print(f"  [!] {img_path.name}: {e}")
    if (i+1) % 200 == 0: print(f"  [{i+1}/{len(all_files)}] done={done} err={err}")

print(f"\n[2/2] Done! {done} saved, {err} errors -> {OUT_DIR}")