238 lines
7.5 KiB
Python
238 lines
7.5 KiB
Python
#!/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()
|