kirl
a64690477a
- 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
139 lines
4.5 KiB
Python
139 lines
4.5 KiB
Python
"""
|
|
CBLPRD → CCPD_TRANS 数据集构建 (ICE-COR 矫正)
|
|
输出: ICE_Dataset/CBLPRD_TRANS/train/
|
|
"""
|
|
import sys, os, cv2, numpy as np, torch
|
|
from pathlib import Path
|
|
import zipfile
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from threading import Lock
|
|
|
|
ICEv2 = Path(r"C:\Users\Kirl\Desktop\ICEv2")
|
|
COR_DIR = ICEv2 / "ICE_NNv2/ICE-COR"
|
|
MODEL_COR = COR_DIR / "output/1_CCPD_19K_EXP_CUT/best.pth"
|
|
ZIP_PATH = r"C:\Users\Kirl\Downloads\CBLPRD-330k_v1.zip"
|
|
OUT_DIR = ICEv2 / "ICE_Dataset/CBLPRD_TRANS/train"
|
|
|
|
sys.path.insert(0, str(COR_DIR))
|
|
from corner_rectify import load_corner_model, predict_plate_corners, order_corners
|
|
|
|
# 类型→文件名标记
|
|
TYPE_MARKER = {
|
|
'普通蓝牌': 'blue',
|
|
'新能源小型车': 'green',
|
|
'新能源大型车': 'new_energy',
|
|
'单层黄牌': 'yellow',
|
|
'黑色车牌': 'black',
|
|
}
|
|
SUPPORTED = set(TYPE_MARKER.keys())
|
|
|
|
print(f"COR: {MODEL_COR}")
|
|
dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
cor_model, _ = load_corner_model(str(MODEL_COR), dev)
|
|
|
|
print(f"Output: {OUT_DIR}")
|
|
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
# 读取 CBLPRD train.txt
|
|
zf = zipfile.ZipFile(ZIP_PATH, 'r')
|
|
labels = []
|
|
for line in zf.read('train.txt').decode('utf-8').strip().split('\n'):
|
|
parts = line.strip().split(' ')
|
|
if len(parts) < 3 or parts[2] not in SUPPORTED:
|
|
continue
|
|
fname = parts[0].split('/')[-1]
|
|
labels.append((fname, parts[1], parts[2]))
|
|
|
|
# 同时加入 val.txt (合并训练)
|
|
for line in zf.read('val.txt').decode('utf-8').strip().split('\n'):
|
|
parts = line.strip().split(' ')
|
|
if len(parts) < 3 or parts[2] not in SUPPORTED:
|
|
continue
|
|
fname = parts[0].split('/')[-1]
|
|
labels.append((fname, parts[1], parts[2]))
|
|
|
|
print(f"Total labels: {len(labels)}")
|
|
from collections import Counter
|
|
for t, c in Counter(l[2] for l in labels).most_common():
|
|
print(f" {t}: {c}")
|
|
|
|
# Build name→zip_path index
|
|
name_idx = {}
|
|
for n in zf.namelist():
|
|
if n.endswith('.jpg'):
|
|
name_idx[n.split('/')[-1]] = n
|
|
print(f"Zip index: {len(name_idx)} images")
|
|
|
|
# 统计已完成的
|
|
existing = set(f.stem.split('_trans')[0] if '_trans' in f.stem else '' for f in OUT_DIR.glob('*.jpg'))
|
|
print(f"Already done: {len(existing)}")
|
|
|
|
# 过滤已完成 (按原始文件名前缀)
|
|
pending = []
|
|
for fname, plate, ptype in labels:
|
|
if fname in name_idx and fname.rsplit('.',1)[0] not in existing:
|
|
pending.append((fname, plate, ptype, name_idx[fname]))
|
|
|
|
print(f"Pending: {len(pending)}")
|
|
|
|
lock = Lock()
|
|
count = [0]
|
|
total = len(pending)
|
|
|
|
def process_one(item):
|
|
fname, plate, ptype, zip_path = item
|
|
try:
|
|
data = zf.read(zip_path)
|
|
img = cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR)
|
|
if img is None:
|
|
return f"READ_ERR:{fname}"
|
|
|
|
# ICE-COR
|
|
corners = predict_plate_corners(img, cor_model, dev)
|
|
ordered = order_corners(corners)
|
|
OW, OH = 192, 64
|
|
cx, cy = ordered.mean(axis=0)
|
|
exp = ordered + (ordered - np.array([[cx, cy]])) * np.array([0.08, 0.05])
|
|
exp[:, 0] = np.clip(exp[:, 0], 0, img.shape[1] - 1)
|
|
exp[:, 1] = np.clip(exp[:, 1], 0, img.shape[0] - 1)
|
|
dst = np.float32([[0, 0], [OW - 1, 0], [OW - 1, OH - 1], [0, OH - 1]])
|
|
M = cv2.getPerspectiveTransform(exp.astype(np.float32), dst)
|
|
warped = cv2.warpPerspective(img, M, (OW, OH))
|
|
except Exception as e:
|
|
return f"COR_ERR:{fname}:{e}"
|
|
|
|
# 文件名: {原始名}_{marker}_{plate}_trans.jpg
|
|
marker = TYPE_MARKER[ptype]
|
|
stem = fname.rsplit('.', 1)[0]
|
|
out_name = f"{stem}_{marker}_{plate}_trans.jpg"
|
|
out_path = OUT_DIR / out_name
|
|
|
|
# Unicode safe write
|
|
ret, enc = cv2.imencode('.jpg', warped, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
|
if ret:
|
|
with open(str(out_path), 'wb') as f:
|
|
f.write(enc.tobytes())
|
|
else:
|
|
return f"WRITE_ERR:{fname}"
|
|
|
|
with lock:
|
|
count[0] += 1
|
|
if count[0] % 500 == 0:
|
|
print(f" [{count[0]}/{total}] {fname} -> {marker}:{plate}")
|
|
return f"OK:{fname}"
|
|
|
|
# 多线程处理
|
|
WORKERS = 8
|
|
print(f"\nProcessing {total} images with {WORKERS} workers...")
|
|
with ThreadPoolExecutor(max_workers=WORKERS) as executor:
|
|
futures = {executor.submit(process_one, item): item for item in pending}
|
|
for future in as_completed(futures):
|
|
result = future.result()
|
|
if result.startswith("READ_ERR") or result.startswith("WRITE_ERR"):
|
|
print(f" [!] {result}")
|
|
|
|
zf.close()
|
|
print(f"\nDone! Total saved: {count[0]}/{total}")
|
|
print(f"Output: {OUT_DIR}")
|
|
print(f"Files: {len(list(OUT_DIR.glob('*.jpg')))}")
|