910 lines
32 KiB
Python
910 lines
32 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
ICE 特殊车牌数据集生成管线
|
||
============================
|
||
|
||
在真实车辆照片（CCPD 格式）上，将普通蓝牌替换为特殊车牌（警车/军车/港澳/使领馆/教练/新能源）。
|
||
仅依赖 CCPD 文件名中已有的角点标注，无需 YOLO 检测或 ICE-COR 推理。
|
||
|
||
管线:
|
||
CCPD 原图 → 解析文件名角点 → 提取原车牌 ROI → 生成目标特殊车牌
|
||
→ Lab 色彩对齐 → 透视变换回贴 → 泊松融合 → 噪点匹配 → 输出
|
||
|
||
用法:
|
||
# 生成 500 张警车特殊车牌
|
||
python plate_swap_pipeline.py \
|
||
--ccpd-dir ./CCPD/CCPD2019/ccpd_base \
|
||
--plate-type police \
|
||
--num 500 \
|
||
--output ./output/police_dataset
|
||
|
||
# 生成军车数据集（多个 CCPD 子集，开启预览）
|
||
python plate_swap_pipeline.py \
|
||
--ccpd-dir ./CCPD/CCPD2019/ccpd_base,./CCPD/CCPD2019/ccpd_tilt \
|
||
--plate-type army \
|
||
--num 2000 \
|
||
--output ./output/army_dataset \
|
||
--preview
|
||
|
||
支持的特殊车牌类型:
|
||
police 警用车牌（白底, 末位"警"红色）
|
||
army 军车/武警（白底, 首位为军车字母）
|
||
gangao 港澳入出境（黑底, 粤+港/澳）
|
||
shiguan_old 使馆旧版（黑底, "使"在首位, 红色）
|
||
shiguan_new 使馆新版（黑底, XXX·XXX使, 白字）
|
||
lingguan_old 领馆旧版（黑底, 省A·XXXX领, "领"红色）
|
||
lingguan_new 领馆新版（黑底, 省XXX·XX领, 白字）
|
||
coach 教练车（黄底, 末位"学"）
|
||
new_energy 大型新能源（黄绿双拼底, 8位）
|
||
|
||
依赖:
|
||
pip install numpy opencv-python tqdm
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import argparse
|
||
import random
|
||
import math
|
||
from pathlib import Path
|
||
from typing import List, Tuple, Optional, Dict
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from tqdm import tqdm
|
||
|
||
# ============================================================
|
||
# 路径解析：定位 chinese_license_plate_generator
|
||
# ============================================================
|
||
_SCRIPT_DIR = Path(__file__).resolve().parent # ICE_Dataset/tools/
|
||
_PROJECT_DIR = _SCRIPT_DIR.parent # ICE_Dataset/
|
||
_GENERATOR_DIR = _PROJECT_DIR / 'chinese_license_plate_generator'
|
||
|
||
if str(_GENERATOR_DIR) not in sys.path:
|
||
sys.path.insert(0, str(_GENERATOR_DIR))
|
||
|
||
from generate_multi_plate import MultiPlateGenerator
|
||
from plate_number import (
|
||
generate_plate_number_white,
|
||
generate_plate_number_army,
|
||
generate_plate_number_black_gangao,
|
||
generate_plate_number_black_shi_old,
|
||
generate_plate_number_black_shi_new,
|
||
generate_plate_number_black_ling_old,
|
||
generate_plate_number_black_ling_new,
|
||
generate_plate_number_yellow_xue,
|
||
generate_plate_number_blue,
|
||
)
|
||
|
||
# ============================================================
|
||
# 特殊车牌类型配置
|
||
# ============================================================
|
||
|
||
# 每个类型: (bg_color, 号码生成函数, 是否需要8位)
|
||
PLATE_TYPE_CONFIG = {
|
||
'police': {
|
||
'bg_color': 'white',
|
||
'generator': generate_plate_number_white,
|
||
'description': '警用车牌'
|
||
},
|
||
'army': {
|
||
'bg_color': 'white_army',
|
||
'generator': generate_plate_number_army,
|
||
'description': '军车/武警'
|
||
},
|
||
'gangao': {
|
||
'bg_color': 'black',
|
||
'generator': generate_plate_number_black_gangao,
|
||
'description': '港澳入出境'
|
||
},
|
||
'shiguan_old': {
|
||
'bg_color': 'black_shi',
|
||
'generator': generate_plate_number_black_shi_old,
|
||
'description': '使馆旧版'
|
||
},
|
||
'shiguan_new': {
|
||
'bg_color': 'black_shi',
|
||
'generator': generate_plate_number_black_shi_new,
|
||
'description': '使馆新版'
|
||
},
|
||
'lingguan_old': {
|
||
'bg_color': 'black',
|
||
'generator': generate_plate_number_black_ling_old,
|
||
'description': '领馆旧版'
|
||
},
|
||
'lingguan_new': {
|
||
'bg_color': 'black',
|
||
'generator': generate_plate_number_black_ling_new,
|
||
'description': '领馆新版'
|
||
},
|
||
'coach': {
|
||
'bg_color': 'yellow',
|
||
'generator': generate_plate_number_yellow_xue,
|
||
'description': '教练车'
|
||
},
|
||
'new_energy': {
|
||
'bg_color': 'green_truck',
|
||
'generator': lambda: generate_plate_number_blue(length=8),
|
||
'description': '大型新能源'
|
||
},
|
||
}
|
||
|
||
|
||
# ============================================================
|
||
# CCPD 文件名解析
|
||
# ============================================================
|
||
|
||
def parse_ccpd_filename(filename: str) -> Optional[Dict]:
|
||
"""
|
||
解析 CCPD 文件名，提取角点坐标和车牌字符。
|
||
|
||
CCPD 文件名格式:
|
||
025-95_113-154&383_386&473-386&473_177&454_154&383_363&402-0_0_22_27_27_33_16-101-25.jpg
|
||
|
||
各段含义:
|
||
[0] area_ratio 面积比
|
||
[1] tilt_degree 倾斜度
|
||
[2] bbox 外接矩形 TL&BR
|
||
[3] corners 4角点 (CCPD 原始顺序不稳定，用几何排序)
|
||
[4] chars 7个字符的CCPD索引
|
||
[5] brightness 亮度
|
||
[6] blurriness 模糊度
|
||
|
||
Returns:
|
||
dict with keys: corners_px, corners_ordered, bbox, plate_chars, plate_str
|
||
或 None 表示解析失败
|
||
"""
|
||
base = Path(filename).stem
|
||
parts = base.split('-')
|
||
|
||
if len(parts) < 7:
|
||
return None
|
||
|
||
try:
|
||
# ---- 角点 (parts[3]): 提取原始4点，几何排序 ----
|
||
coords_str = parts[3]
|
||
corners_raw = []
|
||
for pt in coords_str.split('_'):
|
||
x_str, y_str = pt.split('&')
|
||
corners_raw.append((float(x_str), float(y_str)))
|
||
|
||
if len(corners_raw) != 4:
|
||
return None
|
||
|
||
# 几何排序: 用 x+y / x-y 确定 TL, TR, BR, BL
|
||
# TL = min(x+y), TR = max(x-y), BR = max(x+y), BL = min(x-y)
|
||
pts_arr = np.array(corners_raw, dtype=np.float32)
|
||
sums = pts_arr[:, 0] + pts_arr[:, 1]
|
||
diffs = pts_arr[:, 0] - pts_arr[:, 1]
|
||
|
||
tl = pts_arr[np.argmin(sums)] # 左上: 最小 x+y
|
||
tr = pts_arr[np.argmax(diffs)] # 右上: 最大 x-y
|
||
br = pts_arr[np.argmax(sums)] # 右下: 最大 x+y
|
||
bl = pts_arr[np.argmin(diffs)] # 左下: 最小 x-y
|
||
|
||
corners_ordered = np.float32([tl, tr, br, bl]) # (4,2) TL TR BR BL
|
||
|
||
# ---- BBox (parts[2]) ----
|
||
bb_parts = parts[2].split('_')
|
||
bb_tl = bb_parts[0].split('&')
|
||
bb_br = bb_parts[1].split('&')
|
||
bbox = (
|
||
float(bb_tl[0]), float(bb_tl[1]),
|
||
float(bb_br[0]), float(bb_br[1]),
|
||
)
|
||
|
||
# ---- 字符 (parts[4]) ----
|
||
char_indices = [int(x) for x in parts[4].split('_')]
|
||
plate_str = _decode_ccpd_chars(char_indices)
|
||
|
||
return {
|
||
'corners_px': corners_ordered, # (4,2) TL TR BR BL (几何排序)
|
||
'corners_raw': corners_raw, # 原始4点 (未排序)
|
||
'bbox': bbox, # TL_x, TL_y, BR_x, BR_y
|
||
'plate_chars': char_indices,
|
||
'plate_str': plate_str,
|
||
'brightness': int(parts[5]) if len(parts) > 5 else 0,
|
||
'blurriness': int(parts[6]) if len(parts) > 6 else 0,
|
||
}
|
||
|
||
except (ValueError, IndexError):
|
||
return None
|
||
|
||
|
||
def _decode_ccpd_chars(indices: List[int]) -> str:
|
||
"""将 CCPD 7 位字符索引解码为车牌字符串"""
|
||
provinces = [
|
||
"皖", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑",
|
||
"苏", "浙", "京", "闽", "赣", "鲁", "豫", "鄂", "湘", "粤",
|
||
"桂", "琼", "川", "贵", "云", "藏", "陕", "甘", "青", "宁", "新"
|
||
]
|
||
ccpd_letters = [
|
||
'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K',
|
||
'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V',
|
||
'W', 'X', 'Y', 'Z'
|
||
]
|
||
digits = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
|
||
|
||
chars = []
|
||
for pos, idx in enumerate(indices):
|
||
if pos == 0:
|
||
if 0 <= idx < len(provinces):
|
||
chars.append(provinces[idx])
|
||
elif pos == 1:
|
||
if 0 <= idx < len(ccpd_letters):
|
||
chars.append(ccpd_letters[idx])
|
||
else:
|
||
if 0 <= idx < len(ccpd_letters):
|
||
chars.append(ccpd_letters[idx])
|
||
elif 24 <= idx < 34:
|
||
chars.append(digits[idx - 24])
|
||
return ''.join(chars)
|
||
|
||
|
||
# ============================================================
|
||
# 色彩匹配模块
|
||
# ============================================================
|
||
|
||
def extract_plate_illumination(image_bgr: np.ndarray,
|
||
corners: np.ndarray,
|
||
output_plate_size=(440, 140)) -> np.ndarray:
|
||
"""
|
||
从原图中提取校正后的车牌 ROI，用于获取光照/色彩参考。
|
||
|
||
Args:
|
||
image_bgr: 原始图像
|
||
corners: 4角点 (TL, TR, BR, BL)
|
||
output_plate_size: 校正输出尺寸
|
||
|
||
Returns:
|
||
校正后的车牌 ROI (BGR)
|
||
"""
|
||
dst_rect = np.float32([
|
||
[0, 0],
|
||
[output_plate_size[0] - 1, 0],
|
||
[output_plate_size[0] - 1, output_plate_size[1] - 1],
|
||
[0, output_plate_size[1] - 1],
|
||
])
|
||
|
||
M = cv2.getPerspectiveTransform(corners, dst_rect)
|
||
rectified = cv2.warpPerspective(
|
||
image_bgr, M, output_plate_size,
|
||
flags=cv2.INTER_CUBIC,
|
||
borderMode=cv2.BORDER_REPLICATE,
|
||
)
|
||
return rectified
|
||
|
||
|
||
def estimate_surrounding_brightness(image_bgr: np.ndarray,
|
||
corners: np.ndarray,
|
||
border=15) -> float:
|
||
"""
|
||
估计车牌周围车体区域的平均亮度（用于匹配合成车牌亮度）。
|
||
|
||
取车牌四角外扩 border 像素的环形区域，排除车牌本身。
|
||
"""
|
||
h, w = image_bgr.shape[:2]
|
||
# 车牌中心
|
||
cx, cy = np.mean(corners, axis=0).astype(int)
|
||
# 车牌大致尺寸
|
||
plate_w = int(np.linalg.norm(corners[1] - corners[0]))
|
||
plate_h = int(np.linalg.norm(corners[3] - corners[0]))
|
||
# 外扩环形区域
|
||
x1 = max(0, cx - plate_w // 2 - border)
|
||
y1 = max(0, cy - plate_h // 2 - border)
|
||
x2 = min(w, cx + plate_w // 2 + border)
|
||
y2 = min(h, cy + plate_h // 2 + border)
|
||
if x2 <= x1 or y2 <= y1:
|
||
return 128.0
|
||
|
||
roi = image_bgr[y1:y2, x1:x2]
|
||
gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
|
||
# 取中位数亮度（排除高光/阴影极端值）
|
||
return float(np.median(gray))
|
||
|
||
|
||
def match_plate_color(synthetic_plate: np.ndarray,
|
||
original_plate_roi: np.ndarray,
|
||
bg_color: str) -> np.ndarray:
|
||
"""
|
||
从原车牌提取场景色温，叠加到合成车牌上。
|
||
|
||
策略:
|
||
1. 从原车牌找"白色参考像素"（白字区域），计算其 RGB 均值
|
||
2. 算出色温偏移: gain = 白参考 / 255
|
||
3. 对合成车牌逐通道乘以 gain → 白牌带上场景色温，同时保护字符色不反转
|
||
|
||
暗底牌（black/black_shi）：从原车牌暗区提取"黑参考"，暗部偏色叠加。
|
||
|
||
Args:
|
||
synthetic_plate: 合成特殊车牌 (BGR)
|
||
original_plate_roi: 从原图提取校正后的车牌 (BGR)
|
||
bg_color: 合成车牌底板颜色
|
||
|
||
Returns:
|
||
色温叠加后的合成车牌 (BGR)
|
||
"""
|
||
syn = synthetic_plate.astype(np.float32)
|
||
orig = original_plate_roi.astype(np.float32)
|
||
|
||
# 尺寸对齐
|
||
h_s, w_s = syn.shape[:2]
|
||
h_o, w_o = orig.shape[:2]
|
||
if h_s != h_o or w_s != w_o:
|
||
syn = cv2.resize(syn, (w_o, h_o), interpolation=cv2.INTER_LINEAR)
|
||
|
||
if bg_color in ('black', 'black_shi'):
|
||
# 暗底牌：仅对白字叠加色温（与白牌相同逻辑）
|
||
orig_gray = cv2.cvtColor(orig.astype(np.uint8), cv2.COLOR_BGR2GRAY)
|
||
thresh = np.percentile(orig_gray, 90)
|
||
white_mask = orig_gray > thresh
|
||
if white_mask.sum() < 50:
|
||
return synthetic_plate
|
||
|
||
white_ref = orig[white_mask].mean(axis=0)
|
||
ref_mean = white_ref.mean()
|
||
if ref_mean < 10:
|
||
return synthetic_plate
|
||
color_cast = white_ref / ref_mean
|
||
amplified = 1.0 + (color_cast - 1.0) * 1.0
|
||
amplified = np.clip(amplified, 0.85, 1.15)
|
||
|
||
syn_gray = cv2.cvtColor(syn.astype(np.uint8), cv2.COLOR_BGR2GRAY)
|
||
alpha = np.clip(syn_gray.astype(np.float32) / 128.0, 0.0, 1.0)
|
||
alpha_3 = alpha[:, :, np.newaxis]
|
||
|
||
result = syn * (1.0 + (amplified - 1.0) * alpha_3)
|
||
return np.clip(result, 0, 255).astype(np.uint8)
|
||
|
||
# ---- 亮底牌共用：提取场景白参考 ----
|
||
orig_gray = cv2.cvtColor(orig.astype(np.uint8), cv2.COLOR_BGR2GRAY)
|
||
thresh = np.percentile(orig_gray, 90)
|
||
white_mask = orig_gray > thresh
|
||
if white_mask.sum() < 50:
|
||
return synthetic_plate
|
||
white_ref = orig[white_mask].mean(axis=0)
|
||
ref_mean = white_ref.mean()
|
||
if ref_mean < 10:
|
||
return synthetic_plate
|
||
|
||
if bg_color == 'yellow':
|
||
# 黄牌：简单全图增益，不破坏黄色
|
||
gain = white_ref / 255.0
|
||
gain = np.clip(gain, 0.6, 1.3)
|
||
result = syn * gain
|
||
return np.clip(result, 0, 255).astype(np.uint8)
|
||
|
||
# 白牌 / 绿牌：亮度加权色温
|
||
result = syn * gain
|
||
return np.clip(result, 0, 255).astype(np.uint8)
|
||
|
||
# 白牌 / 绿牌：亮度加权色温
|
||
color_cast = white_ref / ref_mean
|
||
strength_map = {
|
||
'white': 1.0, 'white_army': 1.0,
|
||
'green_truck': 1.5,
|
||
}
|
||
strength = strength_map.get(bg_color, 2.0)
|
||
|
||
amplified = 1.0 + (color_cast - 1.0) * strength
|
||
# 限制每通道偏差在 ±15% 内
|
||
amplified = np.clip(amplified, 0.85, 1.15)
|
||
|
||
syn_gray = cv2.cvtColor(syn.astype(np.uint8), cv2.COLOR_BGR2GRAY)
|
||
alpha = np.clip(syn_gray.astype(np.float32) / 128.0, 0.0, 1.0)
|
||
alpha_3 = alpha[:, :, np.newaxis]
|
||
|
||
result = syn * (1.0 + (amplified - 1.0) * alpha_3)
|
||
return np.clip(result, 0, 255).astype(np.uint8)
|
||
|
||
|
||
def _estimate_plate_bg_mask(plate_gray: np.ndarray, bg_color: str) -> np.ndarray:
|
||
"""
|
||
估计车牌背景区域 mask（排除字符）。
|
||
|
||
白底牌: 字符为暗色 → 背景是高灰度值区域
|
||
黑底牌: 字符为亮色 → 背景是低灰度值区域
|
||
黄底牌: 字符为暗色 → 背景是高灰度值区域（中等阈值）
|
||
绿底牌: 字符为暗色 → 背景是中高灰度值区域
|
||
"""
|
||
h, w = plate_gray.shape
|
||
|
||
if bg_color in ('white', 'white_army', 'yellow'):
|
||
# 白底/黄底: 背景亮，字符暗
|
||
thresh = 180 if 'white' in bg_color else 140
|
||
_, bg_mask = cv2.threshold(plate_gray, thresh, 255, cv2.THRESH_BINARY)
|
||
elif bg_color in ('black', 'black_shi'):
|
||
# 黑底: 背景暗，字符亮
|
||
_, bg_mask = cv2.threshold(plate_gray, 80, 255, cv2.THRESH_BINARY_INV)
|
||
else:
|
||
# 绿底等: Otsu 自适应
|
||
_, bg_mask = cv2.threshold(plate_gray, 0, 255,
|
||
cv2.THRESH_BINARY + cv2.THRESH_OTSU)
|
||
|
||
# 形态学清理
|
||
kernel = np.ones((5, 5), np.uint8)
|
||
bg_mask = cv2.morphologyEx(bg_mask, cv2.MORPH_OPEN, kernel)
|
||
bg_mask = cv2.morphologyEx(bg_mask, cv2.MORPH_CLOSE, kernel)
|
||
|
||
return bg_mask.astype(bool)
|
||
|
||
|
||
# ============================================================
|
||
# 透视变换回贴
|
||
# ============================================================
|
||
|
||
def warp_plate_back(synthetic_plate: np.ndarray,
|
||
dst_corners: np.ndarray,
|
||
output_size: Tuple[int, int]) -> Tuple[np.ndarray, np.ndarray]:
|
||
"""
|
||
将合成车牌透视变换到原图的车牌位置。
|
||
|
||
Args:
|
||
synthetic_plate: 合成车牌 (BGR)
|
||
dst_corners: 原图上的4角点 (TL, TR, BR, BL)
|
||
output_size: 原图尺寸 (W, H)
|
||
|
||
Returns:
|
||
(warped_plate, plate_mask) — warped 图像和对应的掩膜
|
||
"""
|
||
h_plate, w_plate = synthetic_plate.shape[:2]
|
||
|
||
# 合成车牌 4 角 = 矩形
|
||
src_corners = np.float32([
|
||
[0, 0],
|
||
[w_plate - 1, 0],
|
||
[w_plate - 1, h_plate - 1],
|
||
[0, h_plate - 1],
|
||
])
|
||
|
||
M = cv2.getPerspectiveTransform(src_corners, dst_corners.astype(np.float32))
|
||
warped = cv2.warpPerspective(
|
||
synthetic_plate, M, output_size,
|
||
flags=cv2.INTER_CUBIC,
|
||
borderMode=cv2.BORDER_CONSTANT,
|
||
borderValue=(0, 0, 0),
|
||
)
|
||
|
||
# 生成对应的 mask
|
||
mask = np.ones((h_plate, w_plate), dtype=np.uint8) * 255
|
||
mask_warped = cv2.warpPerspective(
|
||
mask, M, output_size,
|
||
flags=cv2.INTER_LINEAR,
|
||
borderMode=cv2.BORDER_CONSTANT,
|
||
borderValue=0,
|
||
)
|
||
|
||
return warped, mask_warped
|
||
|
||
|
||
# ============================================================
|
||
# 边缘羽化贴合（保护车牌本色，不用泊松融合）
|
||
# ============================================================
|
||
|
||
def jpeg_compress(image_bgr: np.ndarray, quality: int = 10) -> np.ndarray:
|
||
"""
|
||
JPEG 压缩→解压，给合成车牌加上真实相机压缩痕迹。
|
||
quality 越小压缩越狠（1=最差, 100=最佳）。
|
||
"""
|
||
_, encoded = cv2.imencode('.jpg', image_bgr,
|
||
[cv2.IMWRITE_JPEG_QUALITY, quality])
|
||
return cv2.imdecode(encoded, cv2.IMREAD_COLOR)
|
||
|
||
|
||
# ============================================================
|
||
# 边缘羽化贴合
|
||
# ============================================================
|
||
|
||
def feather_blend_plate(original_image: np.ndarray,
|
||
warped_plate: np.ndarray,
|
||
mask: np.ndarray,
|
||
corners: np.ndarray,
|
||
edge_width: int = 8) -> np.ndarray:
|
||
"""
|
||
将 warped 车牌贴合到原图，纯内侧羽化（不引入车体暗影）。
|
||
|
||
只在车牌内侧边缘做渐变过渡，外侧完全不混合，
|
||
避免暗色车体渗入车牌边缘造成黑影。
|
||
|
||
Args:
|
||
original_image: 原图 (H, W, 3) BGR
|
||
warped_plate: 透视变换后的合成车牌
|
||
mask: 车牌区域掩膜 (255=车牌区域)
|
||
corners: 4角点
|
||
edge_width: 内侧羽化宽度 (像素)，默认 15px
|
||
|
||
Returns:
|
||
贴合后的图像
|
||
"""
|
||
mask_u8 = (mask > 128).astype(np.uint8) * 255
|
||
|
||
if mask_u8.sum() < 100:
|
||
result = original_image.copy()
|
||
result[mask_u8 > 0] = warped_plate[mask_u8 > 0]
|
||
return result
|
||
|
||
# 自适应羽化宽度（更保守，防止蓝底露出）
|
||
ys, xs = np.where(mask_u8 > 0)
|
||
plate_px_h = ys.max() - ys.min()
|
||
plate_px_w = xs.max() - xs.min()
|
||
adaptive_w = max(edge_width, int(min(plate_px_w, plate_px_h) * 0.06))
|
||
adaptive_w = min(adaptive_w, 20)
|
||
|
||
# 距离变换：每个像素到最近边缘的距离（仅内侧）
|
||
dist = cv2.distanceTransform(mask_u8, cv2.DIST_L2, cv2.DIST_MASK_PRECISE)
|
||
|
||
# alpha = 内侧 dist/width，clamp 到 [0, 1]，加 0.15 底限防止边缘透明露底
|
||
alpha = np.clip(dist / adaptive_w, 0.15, 1.0)
|
||
|
||
# 外侧 alpha 永远为 0（不做任何外侧混合）
|
||
alpha[mask_u8 == 0] = 0.0
|
||
|
||
alpha_3 = alpha[:, :, np.newaxis]
|
||
|
||
result = (warped_plate.astype(np.float32) * alpha_3 +
|
||
original_image.astype(np.float32) * (1.0 - alpha_3))
|
||
return np.clip(result, 0, 255).astype(np.uint8)
|
||
|
||
|
||
# ============================================================
|
||
# 噪点/像素级纹理匹配
|
||
# ============================================================
|
||
|
||
def estimate_noise_level(image_bgr: np.ndarray) -> float:
|
||
"""
|
||
估计图像噪声水平（使用 Laplacian 方差法）。
|
||
"""
|
||
gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
|
||
laplacian = cv2.Laplacian(gray, cv2.CV_64F)
|
||
return laplacian.var()
|
||
|
||
|
||
def estimate_pixel_noise_sigma(image_bgr: np.ndarray) -> float:
|
||
"""
|
||
估计像素级噪声标准差 σ。
|
||
使用高通滤波残差法：高斯模糊 → 原图减模糊图 → 残差 std。
|
||
"""
|
||
blurred = cv2.GaussianBlur(image_bgr, (5, 5), 1.0)
|
||
residual = image_bgr.astype(np.float32) - blurred.astype(np.float32)
|
||
return float(np.std(residual))
|
||
|
||
|
||
def estimate_blur_sigma(image_bgr: np.ndarray) -> float:
|
||
"""
|
||
估计图像的模糊程度 (Laplacian 方差越低 = 越模糊)。
|
||
返回建议的高斯模糊 σ 用于匹配。
|
||
"""
|
||
gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
|
||
lap_var = cv2.Laplacian(gray, cv2.CV_64F).var()
|
||
|
||
if lap_var > 500:
|
||
return 0.0 # 高清：不模糊
|
||
elif lap_var > 200:
|
||
return 0.5 # 轻微模糊
|
||
elif lap_var > 50:
|
||
return 1.2 # 中等模糊
|
||
else:
|
||
return 2.0 # 严重模糊
|
||
|
||
|
||
def match_image_quality(result: np.ndarray,
|
||
original: np.ndarray,
|
||
mask: np.ndarray) -> np.ndarray:
|
||
"""
|
||
在贴牌区域匹配原图噪声和模糊，仅基于原车牌区域像素。
|
||
|
||
1. 从原车牌区域估计像素噪声 σ
|
||
2. 从原车牌区域估计模糊程度
|
||
3. 在新车牌区域叠加匹配噪声 + 模糊
|
||
"""
|
||
mask_u8 = (mask > 128).astype(np.uint8) * 255
|
||
if mask_u8.sum() < 100:
|
||
return result
|
||
|
||
# ---- 仅从原车牌区域提取像素样本 ----
|
||
orig_plate_pixels = original[mask_u8 > 0] # (N, 3) 原车牌区域所有像素
|
||
|
||
# 1) 像素噪声 σ — 从原车牌区域残差估计
|
||
# 用原车牌区域小块做高斯残差
|
||
h, w = original.shape[:2]
|
||
ys, xs = np.where(mask_u8 > 0)
|
||
if len(ys) < 200:
|
||
return result
|
||
|
||
y1, y2 = max(0, ys.min() - 5), min(h, ys.max() + 5)
|
||
x1, x2 = max(0, xs.min() - 5), min(w, xs.max() + 5)
|
||
plate_crop = original[y1:y2, x1:x2]
|
||
|
||
noise_sigma = estimate_pixel_noise_sigma(plate_crop)
|
||
noise_sigma = max(noise_sigma, 1.5) # 最低噪声底噪
|
||
|
||
# 2) 模糊 σ
|
||
blur_sigma = estimate_blur_sigma(plate_crop)
|
||
|
||
# ---- 仅在新车牌区域应用 ----
|
||
result = result.copy()
|
||
|
||
# 先加模糊
|
||
if blur_sigma > 0.1:
|
||
blurred = cv2.GaussianBlur(result, (0, 0), blur_sigma)
|
||
result[mask_u8 > 0] = blurred[mask_u8 > 0]
|
||
|
||
# 再加像素级高斯噪声
|
||
noise = np.random.normal(0, noise_sigma, result.shape).astype(np.float32)
|
||
noisy = result.astype(np.float32) + noise
|
||
# 仅车牌区域加噪声
|
||
alpha = (mask_u8 / 255.0)[:, :, np.newaxis]
|
||
result = (noisy * alpha + result.astype(np.float32) * (1.0 - alpha))
|
||
result = np.clip(result, 0, 255).astype(np.uint8)
|
||
|
||
return result
|
||
|
||
|
||
# ============================================================
|
||
# 主管线
|
||
# ============================================================
|
||
|
||
class PlateSwapPipeline:
|
||
"""
|
||
车牌替换管线。
|
||
|
||
将 CCPD 真实车辆照片中的普通蓝牌替换为指定类型的特殊车牌。
|
||
"""
|
||
|
||
def __init__(self, plate_type: str = 'police'):
|
||
"""
|
||
Args:
|
||
plate_type: 目标特殊车牌类型
|
||
"""
|
||
if plate_type not in PLATE_TYPE_CONFIG:
|
||
raise ValueError(
|
||
f"不支持的车牌类型: {plate_type}。"
|
||
f"支持: {list(PLATE_TYPE_CONFIG.keys())}"
|
||
)
|
||
|
||
self.plate_type = plate_type
|
||
self.config = PLATE_TYPE_CONFIG[plate_type]
|
||
self.bg_color = self.config['bg_color']
|
||
self.number_generator = self.config['generator']
|
||
|
||
# 初始化车牌生成器
|
||
plate_model_dir = str(_GENERATOR_DIR / 'plate_model')
|
||
font_dir = str(_GENERATOR_DIR / 'font_model')
|
||
self.generator = MultiPlateGenerator(plate_model_dir, font_dir)
|
||
|
||
def generate_plate_number(self) -> str:
|
||
"""生成随机特殊车牌号码"""
|
||
return self.number_generator()
|
||
|
||
def process_single(self,
|
||
image_path: str,
|
||
output_dir: str,
|
||
preview: bool = False) -> Optional[str]:
|
||
"""
|
||
处理单张 CCPD 图片。
|
||
|
||
Args:
|
||
image_path: CCPD 图片路径
|
||
output_dir: 输出目录
|
||
preview: 是否保存中间结果供调试
|
||
|
||
Returns:
|
||
成功返回输出路径，失败返回 None
|
||
"""
|
||
filename = Path(image_path).name
|
||
|
||
# ---- 1. 解析 CCPD 文件名 ----
|
||
meta = parse_ccpd_filename(filename)
|
||
if meta is None:
|
||
return None
|
||
|
||
# ---- 2. 读取原图 ----
|
||
img_orig = cv2.imread(image_path)
|
||
if img_orig is None:
|
||
return None
|
||
|
||
h_img, w_img = img_orig.shape[:2]
|
||
corners = meta['corners_px'] # (4,2) TL TR BR BL
|
||
|
||
# 校验角点范围
|
||
if corners.min() < 0 or corners[:, 0].max() > w_img * 1.5 or \
|
||
corners[:, 1].max() > h_img * 1.5:
|
||
return None
|
||
|
||
# ---- 3. 生成目标特殊车牌（先生成号码，确定尺寸） ----
|
||
new_plate_number = self.generate_plate_number()
|
||
plate_w, plate_h = (440, 140) if len(new_plate_number) == 7 else (480, 140)
|
||
|
||
# ---- 4. 提取原车牌 ROI（用于色彩参考，与合成车牌同尺寸） ----
|
||
orig_plate_roi = extract_plate_illumination(img_orig, corners,
|
||
output_plate_size=(plate_w, plate_h))
|
||
|
||
# ---- 5. 生成目标特殊车牌 ----
|
||
synthetic_plate = self.generator.generate_plate_special(
|
||
new_plate_number, self.bg_color, is_double=False, enhance=False
|
||
)
|
||
|
||
# ---- 6. 色温滤镜 ----
|
||
matched_plate = match_plate_color(synthetic_plate, orig_plate_roi, self.bg_color)
|
||
|
||
# ---- 7. 亮度匹配：用原车牌白字亮度做参考 ----
|
||
orig_gray = cv2.cvtColor(orig_plate_roi, cv2.COLOR_BGR2GRAY)
|
||
target_white = np.percentile(orig_gray, 95)
|
||
|
||
matched_gray = cv2.cvtColor(matched_plate, cv2.COLOR_BGR2GRAY)
|
||
current_white = np.percentile(matched_gray, 95)
|
||
|
||
if current_white > 10 and abs(target_white - current_white) > 10:
|
||
l_ratio = target_white / current_white
|
||
# 不同底板色用不同下限，防止白底变灰、绿牌变亮
|
||
clip_cfg = {
|
||
'white': (0.5, 1.1), 'white_army': (0.5, 1.1),
|
||
'green_truck': (0.4, 1.0),
|
||
}
|
||
min_c, max_c = clip_cfg.get(self.bg_color, (0.15, 1.5))
|
||
l_ratio = np.clip(l_ratio, min_c, max_c)
|
||
|
||
# 仅对亮像素（白字）全量叠加，底色部分减半以保护本色
|
||
alpha = np.clip(matched_gray.astype(np.float32) / 128.0, 0.3, 1.0)
|
||
alpha_3 = alpha[:, :, np.newaxis]
|
||
effective_ratio = 1.0 + (l_ratio - 1.0) * alpha_3
|
||
|
||
matched_plate = (matched_plate.astype(np.float32) * effective_ratio)
|
||
matched_plate = np.clip(matched_plate, 0, 255).astype(np.uint8)
|
||
|
||
# ---- 8. 强退化模拟（原图模糊+马赛克） ----
|
||
# JPEG 极限压缩
|
||
matched_plate = jpeg_compress(matched_plate, quality=2)
|
||
# 降采样→升采样 模拟马赛克/像素化
|
||
h_p, w_p = matched_plate.shape[:2]
|
||
small = cv2.resize(matched_plate, (w_p // 4, h_p // 4), interpolation=cv2.INTER_LINEAR)
|
||
matched_plate = cv2.resize(small, (w_p, h_p), interpolation=cv2.INTER_NEAREST)
|
||
# 较强模糊
|
||
matched_plate = cv2.GaussianBlur(matched_plate, (5, 5), 1.5)
|
||
|
||
# ---- 9. 透视变换回贴 ----
|
||
warped, plate_mask = warp_plate_back(
|
||
matched_plate, corners, output_size=(w_img, h_img)
|
||
)
|
||
|
||
# ---- 10. 边缘羽化贴合 ----
|
||
blended = feather_blend_plate(img_orig, warped, plate_mask, corners)
|
||
|
||
# ---- 11. 像素噪声匹配 ----
|
||
final = match_image_quality(blended, img_orig, plate_mask)
|
||
|
||
# ---- 12. 保存 ----
|
||
basename = Path(filename).stem
|
||
out_name = f"{basename}_{self.plate_type}_{new_plate_number}.jpg"
|
||
out_path = os.path.join(output_dir, out_name)
|
||
cv2.imwrite(out_path, final, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||
|
||
# ---- 可选: 保存调试中间结果 ----
|
||
if preview:
|
||
preview_dir = os.path.join(output_dir, '_preview')
|
||
os.makedirs(preview_dir, exist_ok=True)
|
||
pfx = basename
|
||
cv2.imwrite(os.path.join(preview_dir, f'{pfx}_01_orig.jpg'), img_orig)
|
||
cv2.imwrite(os.path.join(preview_dir, f'{pfx}_02_roi.jpg'), orig_plate_roi)
|
||
cv2.imwrite(os.path.join(preview_dir, f'{pfx}_03_synthetic.jpg'), synthetic_plate)
|
||
cv2.imwrite(os.path.join(preview_dir, f'{pfx}_04_matched.jpg'), matched_plate)
|
||
cv2.imwrite(os.path.join(preview_dir, f'{pfx}_05_warped.jpg'), warped)
|
||
cv2.imwrite(os.path.join(preview_dir, f'{pfx}_06_final.jpg'), final)
|
||
|
||
return out_path
|
||
|
||
def process_batch(self,
|
||
ccpd_dirs: List[str],
|
||
output_dir: str,
|
||
num_images: int = 500,
|
||
preview: bool = False,
|
||
seed: int = 42):
|
||
"""
|
||
批量处理 CCPD 图片目录。
|
||
|
||
Args:
|
||
ccpd_dirs: CCPD 子目录列表
|
||
output_dir: 输出根目录
|
||
num_images: 生成图片总数
|
||
preview: 是否保存调试中间结果
|
||
seed: 随机种子
|
||
"""
|
||
random.seed(seed)
|
||
np.random.seed(seed)
|
||
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# 收集所有 CCPD 图片
|
||
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', '*.JPG', '*.JPEG', '*.PNG'):
|
||
all_images.extend(dd.glob(ext))
|
||
|
||
if not all_images:
|
||
print("[ERROR] 未找到任何 CCPD 图片")
|
||
return
|
||
|
||
# 随机排列
|
||
random.shuffle(all_images)
|
||
|
||
print(f"特殊车牌类型: {self.config['description']} ({self.plate_type})")
|
||
print(f"底板颜色: {self.bg_color}")
|
||
print(f"CCPD 图片总数: {len(all_images)}")
|
||
print(f"目标生成数: {num_images}")
|
||
print(f"输出目录: {output_dir}")
|
||
print("-" * 60)
|
||
|
||
success = 0
|
||
failed = 0
|
||
pbar = tqdm(total=num_images, desc='生成特殊车牌', unit='张')
|
||
|
||
for img_path in all_images:
|
||
if success >= num_images:
|
||
break
|
||
|
||
result = self.process_single(
|
||
str(img_path), output_dir, preview=(preview and success < 10)
|
||
)
|
||
|
||
if result:
|
||
success += 1
|
||
pbar.update(1)
|
||
pbar.set_postfix_str(f'成功:{success} 失败:{failed}')
|
||
else:
|
||
failed += 1
|
||
|
||
pbar.close()
|
||
print(f"\n完成! 成功: {success}, 失败: {failed}")
|
||
print(f"输出目录: {output_dir}")
|
||
|
||
|
||
# ============================================================
|
||
# CLI
|
||
# ============================================================
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description='ICE 特殊车牌数据集生成管线 — 基于 CCPD 角点替换普通车牌为特殊车牌',
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""
|
||
示例:
|
||
python plate_swap_pipeline.py --ccpd-dir ./CCPD/CCPD2019/ccpd_base --plate-type police --num 500
|
||
python plate_swap_pipeline.py --ccpd-dir ./CCPD/CCPD2019/ccpd_base,./CCPD/CCPD2019/ccpd_tilt --plate-type army --num 2000
|
||
python plate_swap_pipeline.py --ccpd-dir ./CCPD --plate-type gangao --num 1000 --preview
|
||
"""
|
||
)
|
||
parser.add_argument('--ccpd-dir', type=str, required=True,
|
||
help='CCPD 图片目录（多个用逗号分隔）')
|
||
parser.add_argument('--plate-type', type=str, required=True,
|
||
choices=list(PLATE_TYPE_CONFIG.keys()),
|
||
help='目标特殊车牌类型')
|
||
parser.add_argument('--num', type=int, default=500,
|
||
help='生成图片数量 (默认: 500)')
|
||
parser.add_argument('--output', type=str, default='./output/special_plates',
|
||
help='输出目录 (默认: ./output/special_plates)')
|
||
parser.add_argument('--preview', action='store_true',
|
||
help='保存前10张的中间结果供调试')
|
||
parser.add_argument('--seed', type=int, default=42,
|
||
help='随机种子 (默认: 42)')
|
||
|
||
args = parser.parse_args()
|
||
|
||
ccpd_dirs = [d.strip() for d in args.ccpd_dir.split(',') if d.strip()]
|
||
pipeline = PlateSwapPipeline(plate_type=args.plate_type)
|
||
pipeline.process_batch(
|
||
ccpd_dirs=ccpd_dirs,
|
||
output_dir=args.output,
|
||
num_images=args.num,
|
||
preview=args.preview,
|
||
seed=args.seed,
|
||
)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|