特殊车牌数据集构建

tools/extract_real_plates.py

@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+CCPD 真实车牌 ROI 提取工具
+==========================
+从 CCPD 图片中透视校正并裁剪车牌区域，保存为统一尺寸的 ROI 图片。
+用于 CycleGAN/CUT 风格对齐训练的真实车牌域数据。
+
+用法:
+  python extract_real_plates.py --ccpd-dir ./CCPD/CCPD2019/ccpd_base --num 500 --output ./gan_data/real_plates
+
+输出:
+  440×140 的校正车牌 ROI (BGR .jpg)
+"""
+
+import os
+import sys
+import argparse
+import random
+from pathlib import Path
+from typing import Optional, Tuple
+
+import cv2
+import numpy as np
+from tqdm import tqdm
+
+
+def parse_ccpd_corners(filename: str) -> Optional[np.ndarray]:
+    """
+    从 CCPD 文件名解析4角点，几何排序为 [TL, TR, BR, BL]。
+    """
+    base = Path(filename).stem
+    parts = base.split('-')
+    if len(parts) < 5:
+        return None
+
+    try:
+        coords_str = parts[3]
+        pts = []
+        for pt in coords_str.split('_'):
+            x, y = pt.split('&')
+            pts.append((float(x), float(y)))
+
+        if len(pts) != 4:
+            return None
+
+        a = np.array(pts)
+        sums = a[:, 0] + a[:, 1]
+        diffs = a[:, 0] - a[:, 1]
+
+        tl = a[np.argmin(sums)]
+        tr = a[np.argmax(diffs)]
+        br = a[np.argmax(sums)]
+        bl = a[np.argmin(diffs)]
+
+        return np.float32([tl, tr, br, bl])
+    except (ValueError, IndexError):
+        return None
+
+
+def extract_plate_roi(image_bgr: np.ndarray,
+                       corners: np.ndarray,
+                       output_size: Tuple[int, int] = (440, 140)) -> Optional[np.ndarray]:
+    """
+    透视校正提取车牌 ROI。
+    """
+    h, w = image_bgr.shape[:2]
+    # 校验角点在图像范围内
+    if corners.min() < -10 or corners[:, 0].max() > w + 10 or corners[:, 1].max() > h + 10:
+        return None
+
+    dst = np.float32([
+        [0, 0],
+        [output_size[0] - 1, 0],
+        [output_size[0] - 1, output_size[1] - 1],
+        [0, output_size[1] - 1],
+    ])
+
+    M = cv2.getPerspectiveTransform(corners, dst)
+    roi = cv2.warpPerspective(image_bgr, M, output_size,
+                               flags=cv2.INTER_CUBIC,
+                               borderMode=cv2.BORDER_REPLICATE)
+    return roi
+
+
+def main():
+    parser = argparse.ArgumentParser(description='CCPD 真实车牌 ROI 提取')
+    parser.add_argument('--ccpd-dir', type=str, required=True,
+                        help='CCPD 图片目录（多个用逗号分隔）')
+    parser.add_argument('--num', type=int, default=500,
+                        help='提取数量 (默认: 500)')
+    parser.add_argument('--output', type=str, default='./gan_data/real_plates',
+                        help='输出目录')
+    parser.add_argument('--seed', type=int, default=42)
+
+    args = parser.parse_args()
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+
+    # 收集图片
+    ccpd_dirs = [d.strip() for d in args.ccpd_dir.split(',') if d.strip()]
+    all_images = []
+    for d in ccpd_dirs:
+        dd = Path(d)
+        if not dd.is_dir():
+            print(f"[WARN] 目录不存在: {d}")
+            continue
+        for ext in ('*.jpg', '*.jpeg', '*.png'):
+            all_images.extend(dd.glob(ext))
+
+    random.shuffle(all_images)
+    print(f"CCPD 图片总数: {len(all_images)}")
+    print(f"目标提取数:   {args.num}")
+    print(f"输出目录:     {args.output}")
+    print("-" * 50)
+
+    os.makedirs(args.output, exist_ok=True)
+
+    success = 0
+    for img_path in tqdm(all_images, desc='提取真实车牌ROI', unit='张'):
+        if success >= args.num:
+            break
+
+        corners = parse_ccpd_corners(img_path.name)
+        if corners is None:
+            continue
+
+        img = cv2.imread(str(img_path))
+        if img is None:
+            continue
+
+        roi = extract_plate_roi(img, corners)
+        if roi is None:
+            continue
+
+        out_name = f"real_{img_path.stem}.jpg"
+        cv2.imwrite(os.path.join(args.output, out_name), roi,
+                     [cv2.IMWRITE_JPEG_QUALITY, 95])
+        success += 1
+
+    print(f"\n完成! 提取: {success} 张")
+
+
+if __name__ == '__main__':
+    main()