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