feat: 特殊车牌生成管线 + C++/多进程/色温校正
This commit is contained in:
kirl committed
1 parent
3808b9f035
commit
3101992f51
29 files changed
+2336
-80
No files matched your search
@@ -0,0 +1,237 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
CCPD Combined Dataset Builder
|
||||
==============================
|
||||
从 CCPD2019(5k) + CCPD2020(5k) + CCPD_SPE(全量) 构建 YOLO 单类车牌数据集。
|
||||
|
||||
CCPD 文件名格式:
|
||||
025-95_113-154&383_386&473-386&473_177&454_154&383_363&402-...
|
||||
第4段(0-indexed=3): 4个角点 x&y_x&y_x&y_x&y
|
||||
顺序: 左下, 右下, 左上, 右上
|
||||
|
||||
输出:
|
||||
ICE_Dataset/ccpd_combined_5k5k_spe/
|
||||
├── images/train/
|
||||
├── images/val/
|
||||
├── labels/train/
|
||||
├── labels/val/
|
||||
└── dataset.yaml
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import random
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from PIL import Image
|
||||
|
||||
# ── 路径配置 ───────────────────────────────────────────
|
||||
ICE_DATASET = Path(r'C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset')
|
||||
CCPD2019_DIR = ICE_DATASET / 'CCPD' / 'CCPD2019'
|
||||
CCPD2020_DIR = ICE_DATASET / 'CCPD' / 'CCPD2020'
|
||||
CCPD_SPE_DIR = ICE_DATASET / 'CCPD_SPE'
|
||||
OUT_DIR = ICE_DATASET / 'ccpd_combined_5k5k_spe'
|
||||
|
||||
SEED = 42
|
||||
TRAIN_RATIO = 0.9
|
||||
|
||||
|
||||
# ── CCPD 文件名解析 ────────────────────────────────────
|
||||
def parse_ccpd_corners(filename):
|
||||
"""
|
||||
从 CCPD 文件名解析车牌 4 角点 (像素坐标)
|
||||
返回: [(x1,y1), (x2,y2), (x3,y3), (x4,y4)] 或 None
|
||||
顺序: 左下, 右下, 左上, 右上
|
||||
"""
|
||||
base = os.path.splitext(filename)[0]
|
||||
try:
|
||||
parts = base.split('-')
|
||||
if len(parts) < 5:
|
||||
return None
|
||||
coords_str = parts[3]
|
||||
coords = []
|
||||
for pt in coords_str.split('_'):
|
||||
x, y = pt.split('&')
|
||||
coords.append((float(x), float(y)))
|
||||
if len(coords) == 4:
|
||||
return coords
|
||||
except (ValueError, IndexError):
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def corners_to_yolo_bbox(corners, img_w, img_h):
|
||||
"""
|
||||
4 角点 → 外接矩形 → YOLO 归一化 xywh_center
|
||||
返回: (x_center, y_center, w, h) 归一化值 或 None
|
||||
"""
|
||||
xs = [c[0] for c in corners]
|
||||
ys = [c[1] for c in corners]
|
||||
xmin, xmax = min(xs), max(xs)
|
||||
ymin, ymax = min(ys), max(ys)
|
||||
|
||||
bw = xmax - xmin
|
||||
bh = ymax - ymin
|
||||
if bw <= 0 or bh <= 0:
|
||||
return None
|
||||
if xmin < 0 or ymin < 0 or xmax > img_w or ymax > img_h:
|
||||
return None # bbox 越界，跳过
|
||||
|
||||
cx = (xmin + xmax) / 2.0 / img_w
|
||||
cy = (ymin + ymax) / 2.0 / img_h
|
||||
w = bw / img_w
|
||||
h = bh / img_h
|
||||
return (cx, cy, w, h)
|
||||
|
||||
|
||||
# ── 图片收集 ───────────────────────────────────────────
|
||||
def collect_images(root_dir, limit=None):
|
||||
"""递归收集目录下所有 .jpg 文件，返回绝对路径列表"""
|
||||
files = []
|
||||
for dirpath, _, filenames in os.walk(root_dir):
|
||||
for fn in filenames:
|
||||
if fn.lower().endswith('.jpg'):
|
||||
files.append(os.path.join(dirpath, fn))
|
||||
random.seed(SEED)
|
||||
random.shuffle(files)
|
||||
if limit and len(files) > limit:
|
||||
files = files[:limit]
|
||||
return files
|
||||
|
||||
|
||||
# ── 主流程 ─────────────────────────────────────────────
|
||||
def main():
|
||||
print("=" * 60)
|
||||
print("CCPD Combined Dataset Builder")
|
||||
print("=" * 60)
|
||||
|
||||
# 1. 收集图片
|
||||
print("\n[1/4] Collecting images...")
|
||||
|
||||
print(f" CCPD2019: scanning {CCPD2019_DIR} ...")
|
||||
ccpd2019_files = collect_images(str(CCPD2019_DIR), limit=5000)
|
||||
print(f" → sampled {len(ccpd2019_files)}")
|
||||
|
||||
print(f" CCPD2020: scanning {CCPD2020_DIR} ...")
|
||||
ccpd2020_files = collect_images(str(CCPD2020_DIR), limit=5000)
|
||||
print(f" → sampled {len(ccpd2020_files)}")
|
||||
|
||||
print(f" CCPD_SPE: scanning {CCPD_SPE_DIR} ...")
|
||||
ccpd_spe_files = collect_images(str(CCPD_SPE_DIR), limit=None)
|
||||
print(f" → all {len(ccpd_spe_files)}")
|
||||
|
||||
all_files = ccpd2019_files + ccpd2020_files + ccpd_spe_files
|
||||
random.seed(SEED)
|
||||
random.shuffle(all_files)
|
||||
total = len(all_files)
|
||||
print(f"\n Total images: {total}")
|
||||
|
||||
# 2. 解析 & 过滤
|
||||
print("\n[2/4] Parsing bboxes & validating...")
|
||||
valid_samples = [] # (src_path, yolo_bbox_str)
|
||||
skipped_parse = 0
|
||||
skipped_bbox = 0
|
||||
skipped_read = 0
|
||||
|
||||
for src_path in all_files:
|
||||
filename = os.path.basename(src_path)
|
||||
corners = parse_ccpd_corners(filename)
|
||||
if corners is None:
|
||||
skipped_parse += 1
|
||||
continue
|
||||
|
||||
# 获取图片尺寸
|
||||
try:
|
||||
with Image.open(src_path) as img:
|
||||
img_w, img_h = img.size
|
||||
except Exception:
|
||||
skipped_read += 1
|
||||
continue
|
||||
|
||||
bbox = corners_to_yolo_bbox(corners, img_w, img_h)
|
||||
if bbox is None:
|
||||
skipped_bbox += 1
|
||||
continue
|
||||
|
||||
# YOLO 标签: class_id cx cy w h
|
||||
label_line = f"0 {bbox[0]:.6f} {bbox[1]:.6f} {bbox[2]:.6f} {bbox[3]:.6f}"
|
||||
valid_samples.append((src_path, filename, label_line))
|
||||
|
||||
print(f" Valid: {len(valid_samples)}")
|
||||
print(f" Skipped (parse): {skipped_parse}")
|
||||
print(f" Skipped (bbox): {skipped_bbox}")
|
||||
print(f" Skipped (read): {skipped_read}")
|
||||
|
||||
# 3. Train/Val 分割
|
||||
print("\n[3/4] Splitting train/val (90/10)...")
|
||||
random.seed(SEED)
|
||||
random.shuffle(valid_samples)
|
||||
n_train = int(len(valid_samples) * TRAIN_RATIO)
|
||||
train_samples = valid_samples[:n_train]
|
||||
val_samples = valid_samples[n_train:]
|
||||
print(f" Train: {len(train_samples)}, Val: {len(val_samples)}")
|
||||
|
||||
# 4. 输出
|
||||
print(f"\n[4/4] Writing to {OUT_DIR} ...")
|
||||
|
||||
# 清理旧目录
|
||||
for sub in ['images/train', 'images/val', 'labels/train', 'labels/val']:
|
||||
(OUT_DIR / sub).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def write_split(samples, split_name):
|
||||
img_dir = OUT_DIR / 'images' / split_name
|
||||
lbl_dir = OUT_DIR / 'labels' / split_name
|
||||
copied = 0
|
||||
for src_path, filename, label_line in samples:
|
||||
# 处理重名 (不同源可能有同名文件)
|
||||
dst_name = filename
|
||||
counter = 0
|
||||
while (img_dir / dst_name).exists():
|
||||
counter += 1
|
||||
base, ext = os.path.splitext(filename)
|
||||
dst_name = f"{base}_{counter:04d}{ext}"
|
||||
label_name = os.path.splitext(dst_name)[0] + '.txt'
|
||||
|
||||
try:
|
||||
shutil.copy2(src_path, img_dir / dst_name)
|
||||
except Exception as e:
|
||||
print(f" WARNING: copy failed {src_path}: {e}")
|
||||
continue
|
||||
with open(lbl_dir / label_name, 'w') as f:
|
||||
f.write(label_line + '\n')
|
||||
copied += 1
|
||||
return copied
|
||||
|
||||
n_copied_train = write_split(train_samples, 'train')
|
||||
n_copied_val = write_split(val_samples, 'val')
|
||||
print(f" Train images written: {n_copied_train}")
|
||||
print(f" Val images written: {n_copied_val}")
|
||||
|
||||
# 5. 生成 dataset.yaml
|
||||
yaml_path = OUT_DIR / 'dataset.yaml'
|
||||
yaml_content = f"""# CCPD Combined Dataset (single class: license_plate)
|
||||
# CCPD2019 5k + CCPD2020 5k + CCPD_SPE all = {total} images
|
||||
|
||||
path: {OUT_DIR.as_posix()}
|
||||
train: images/train
|
||||
val: images/val
|
||||
|
||||
nc: 1
|
||||
names:
|
||||
0: license_plate
|
||||
"""
|
||||
with open(yaml_path, 'w') as f:
|
||||
f.write(yaml_content)
|
||||
print(f"\n dataset.yaml → {yaml_path}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("DONE!")
|
||||
print(f" Dataset: {OUT_DIR}")
|
||||
print(f" Images: {n_copied_train + n_copied_val}")
|
||||
print(f" YAML: {yaml_path}")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in new issue
Block a user