Dataset: add READMEs, organize scripts, update .gitignore
- Add README.md to all 16 dataset directories with script references - Centralize generation scripts under scripts/ with full README - Update .gitignore to cover all data directories, keep only metadata - Remove redundant CCPD/.gitignore (parent covers same rules) - Remove stale scripts/test/ (254MB pipeline residue) - Fix label parser reference: scripts are authoritative source of truth
This commit is contained in:
kirl committed
1 parent
7867c3bfaf
commit
a64690477a
43 files changed
+2691
-205
No files matched your search
@@ -0,0 +1,156 @@
|
||||
"""
|
||||
CCPD原始真图 去偏采样 + YOLO EXP_CUT + 角点TRANS处理
|
||||
输出: ICE_Dataset/CCPD_REAL_TRANS/train/
|
||||
省份均衡: 每省最多 MAX_PER_PROVINCE 张
|
||||
管线: YOLO检出→12%/18%外扩→裁切→角点重映射→warp (与322K EXP_CUT一致)
|
||||
"""
|
||||
import os, sys, cv2, random, shutil
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from collections import defaultdict
|
||||
|
||||
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 parse_ccpd_filename, imread_unicode
|
||||
|
||||
CCPD_DIRS = [
|
||||
ICEv2 / "ICE_Dataset/CCPD/CCPD2019",
|
||||
ICEv2 / "ICE_Dataset/CCPD/CCPD2020",
|
||||
]
|
||||
OUT_DIR = ICEv2 / "ICE_Dataset/CCPD_REAL_TRANS/train"
|
||||
MAX_PER_PROVINCE = 300
|
||||
SEED = 42
|
||||
|
||||
YOLO_PATH = str(ICEv2 / "ICE_NNv2/yolov5t/output/1_CCPD_19K/weights/best.pt")
|
||||
EXPAND_NORMAL = 1.12 # 常规外扩 12%
|
||||
EXPAND_SMALL = 1.18 # 小车牌外扩 18%
|
||||
SMALL_BOX_W = SMALL_BOX_H = 16
|
||||
MIN_PAD = 2
|
||||
|
||||
random.seed(SEED)
|
||||
|
||||
|
||||
def process_one(img_path, plate, model, img_size=640):
|
||||
"""YOLO→外扩→裁切→角点重映射→warp (仿 EXP_CUT 管线)"""
|
||||
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
|
||||
# 取置信度最高的框 (YOLO 输出 xyxy 格式)
|
||||
best_idx = boxes.conf.argmax()
|
||||
x1, y1, x2, y2 = boxes.xyxy[best_idx].cpu().numpy().astype(int)
|
||||
bw, bh = x2 - x1, y2 - y1
|
||||
|
||||
# 外扩 (与 EXP_CUT 完全一致)
|
||||
expand_factor = EXPAND_SMALL if (bw < SMALL_BOX_W or bh < SMALL_BOX_H) else EXPAND_NORMAL
|
||||
pad_x = max(MIN_PAD, int(round(bw * (expand_factor - 1.0) / 2.0)))
|
||||
pad_y = max(MIN_PAD, int(round(bh * (expand_factor - 1.0) / 2.0)))
|
||||
|
||||
cx1 = max(0, int(x1) - pad_x)
|
||||
cy1 = max(0, int(y1) - pad_y)
|
||||
cx2 = min(w, int(x2) + pad_x)
|
||||
cy2 = min(h, int(y2) + pad_y)
|
||||
if cx2 <= cx1 or cy2 <= cy1:
|
||||
return None
|
||||
|
||||
cropped = img[cy1:cy2, cx1:cx2]
|
||||
|
||||
# 角点: 用 YOLO 检出框的四角 (原始图坐标), 重映射到裁切图坐标
|
||||
corners_orig = np.float32([[x1, y1], [x2, y1], [x2, y2], [x1, y2]]) # TL, TR, BR, BL
|
||||
corners_crop = (corners_orig - np.array([cx1, cy1])).astype(np.float32)
|
||||
|
||||
# 透视变换 → 192x64
|
||||
OW, OH = 192, 64
|
||||
dst = np.float32([[0, 0], [OW - 1, 0], [OW - 1, OH - 1], [0, OH - 1]])
|
||||
M = cv2.getPerspectiveTransform(corners_crop, dst)
|
||||
warped = cv2.warpPerspective(cropped, M, (OW, OH))
|
||||
return warped
|
||||
|
||||
|
||||
# ---- Step 0: 清空旧输出 ----
|
||||
if OUT_DIR.exists():
|
||||
shutil.rmtree(str(OUT_DIR))
|
||||
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
# ---- Step 1: 收集CCPD真图, 按省份分组 ----
|
||||
print("[1/3] Scanning CCPD directories...")
|
||||
prov_files = defaultdict(list) # province -> [(full_path, stem, plate)]
|
||||
total = skipped = 0
|
||||
|
||||
for ccpd_dir in CCPD_DIRS:
|
||||
for root, dirs, files in os.walk(str(ccpd_dir)):
|
||||
for f in files:
|
||||
if not f.lower().endswith(".jpg"):
|
||||
continue
|
||||
indices, plate = parse_ccpd_filename(f)
|
||||
if indices is None:
|
||||
skipped += 1
|
||||
continue
|
||||
total += 1
|
||||
prov = plate[0]
|
||||
full_path = os.path.join(root, f)
|
||||
stem = Path(f).stem
|
||||
prov_files[prov].append((full_path, stem, plate))
|
||||
|
||||
print(f" Valid: {total}, Skipped: {skipped}, Provinces: {len(prov_files)}")
|
||||
for prov in sorted(prov_files, key=lambda p: -len(prov_files[p])):
|
||||
print(f" {prov}: {len(prov_files[prov])}")
|
||||
|
||||
|
||||
# ---- Step 2: 去偏采样 ----
|
||||
print(f"\n[2/3] Sampling max {MAX_PER_PROVINCE} per province...")
|
||||
sampled = []
|
||||
for prov, items in prov_files.items():
|
||||
if len(items) <= MAX_PER_PROVINCE:
|
||||
sampled.extend(items)
|
||||
else:
|
||||
sampled.extend(random.sample(items, MAX_PER_PROVINCE))
|
||||
random.shuffle(sampled)
|
||||
print(f" Sampled: {len(sampled)}")
|
||||
prov_dist = defaultdict(int)
|
||||
for _, _, plate in sampled:
|
||||
prov_dist[plate[0]] += 1
|
||||
for prov in sorted(prov_dist):
|
||||
print(f" {prov}: {prov_dist[prov]}")
|
||||
|
||||
|
||||
# ---- Step 3: YOLO EXP_CUT + warp + 保存 ----
|
||||
print(f"\n[3/3] YOLO EXP_CUT (12%/18%) + warp {len(sampled)} images -> {OUT_DIR}...")
|
||||
print(f" YOLO: {YOLO_PATH}")
|
||||
from ultralytics import YOLO
|
||||
model = YOLO(YOLO_PATH)
|
||||
done = err = 0
|
||||
for i, (full_path, stem, plate) in enumerate(sampled):
|
||||
try:
|
||||
warped = process_one(full_path, plate, model)
|
||||
if warped is None:
|
||||
err += 1; continue
|
||||
|
||||
out_name = f"{stem}_real_blue_{plate}_trans.jpg"
|
||||
out_path = OUT_DIR / out_name
|
||||
|
||||
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 <= 5:
|
||||
print(f" [!] {Path(full_path).name}: {e}")
|
||||
|
||||
if (i + 1) % 200 == 0:
|
||||
print(f" [{i+1}/{len(sampled)}] done={done}, err={err}")
|
||||
|
||||
print(f"\nDone! Saved: {done}, Errors: {err}")
|
||||
print(f"Output: {OUT_DIR} ({len(list(OUT_DIR.glob('*.jpg')))} files)")
|
||||
Reference in new issue
Block a user