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,138 @@
|
||||
"""
|
||||
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')))}")
|
||||
Reference in new issue
Block a user