#!/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()