diff --git a/.gitignore b/.gitignore index 2e55693..ccda614 100644 --- a/.gitignore +++ b/.gitignore @@ -92,7 +92,12 @@ CBLPRD_TRANS/train/*.jpg CBLPRD_eval/ # ============================================================ -# 其他 +# CCPD GREEN TRANS (杩愯 build_green_trans.py 鐢熸垚) +# ============================================================ +CCPD_GREEN_TRANS/ + +# ============================================================ +# 鍏朵粬 # ============================================================ ccpd_combined_5k5k_spe/images/ ccpd_combined_5k5k_spe/labels/ diff --git a/scripts/build_green_trans.py b/scripts/build_green_trans.py new file mode 100644 index 0000000..6219192 --- /dev/null +++ b/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}")