构建绿牌TRANS数据集

This commit is contained in:
kirl committed 2026-06-26 13:24:24 +08:00
1 parent a64690477a
commit 416f9a0bb9
2 files changed
+81 -1

No files matched your search

+6 -1
View File
@@ -92,7 +92,12 @@ CBLPRD_TRANS/train/*.jpg
CBLPRD_eval/ CBLPRD_eval/
# ============================================================ # ============================================================
# 其他 # CCPD GREEN TRANS (杩愯 build_green_trans.py 鐢熸垚)
# ============================================================
CCPD_GREEN_TRANS/
# ============================================================
# 鍏朵粬
# ============================================================ # ============================================================
ccpd_combined_5k5k_spe/images/ ccpd_combined_5k5k_spe/images/
ccpd_combined_5k5k_spe/labels/ ccpd_combined_5k5k_spe/labels/
+75
View File
@@ -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}")