构建绿牌TRANS数据集

scripts/build_green_trans.py

@@ -0,0 +1,75 @@
+"""
+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}")