JiChuang/ICE_Dataset
3
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

特殊车牌数据集构建

Browse Source
This commit is contained in:
lee committed 2026-06-20 19:49:52 +08:00
1 parent 003f9efa6d
commit 3808b9f035
10 files changed
+1696
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tools/_gan_prep.sh
+42
View File
@@ -0,0 +1,42 @@
#!/bin/bash
# ============================================================
# GAN 数据准备：合成车牌(A) + 真实车牌(B)
# ============================================================
source /home/huawei/rknn_env/bin/activate
cd /mnt/d/ICEv2/ICE_Dataset
OUT_A=./gan_data/domain_A_synthetic
OUT_B=./gan_data/domain_B_real
CCPD=./CCPD/CCPD2019/ccpd_base,./CCPD/CCPD2019/ccpd_tilt,./CCPD/CCPD2019/ccpd_rotate,./CCPD/CCPD2019/ccpd_fn
echo "============================================"
echo " Phase 1: 生成合成车牌 (9类型 x 1000张)"
echo "============================================"
for p in police army gangao shiguan_old shiguan_new lingguan_old lingguan_new coach new_energy; do
echo "--- $p ---"
python3 tools/plate_swap_pipeline.py \
--ccpd-dir ./CCPD/CCPD2019/ccpd_base \
--plate-type $p \
--num 1000 \
--output $OUT_A/$p \
--seed 42 2>&1 | tail -2
done
echo ""
echo "============================================"
echo " Phase 2: 提取真实车牌ROI (1000张)"
echo "============================================"
python3 tools/extract_real_plates.py \
--ccpd-dir "$CCPD" \
--num 1000 \
--output $OUT_B \
--seed 42
echo ""
echo "============================================"
echo " GAN 数据准备完成!"
echo " Domain A (合成): $OUT_A"
echo " Domain B (真实): $OUT_B"
echo "============================================"
ls -l $OUT_A/*/ | wc -l
ls -l $OUT_B/ | wc -l
tools/_p20.sh
+5
View File
@@ -0,0 +1,5 @@
#!/bin/bash
source /home/huawei/rknn_env/bin/activate
cd /mnt/d/ICEv2/ICE_Dataset
python3 tools/plate_swap_pipeline.py --ccpd-dir ./CCPD/CCPD2019/ccpd_base --plate-type police --num 20 --output ./output/preview/police --seed 42 2>&1 | tail -2
echo DONE
tools/_r3.sh
+9
View File
@@ -0,0 +1,9 @@
#!/bin/bash
source /home/huawei/rknn_env/bin/activate
cd /mnt/d/ICEv2/ICE_Dataset
for p in police army gangao shiguan_old shiguan_new lingguan_old lingguan_new coach new_energy; do
echo "=== $p ==="
python3 tools/plate_swap_pipeline.py --ccpd-dir ./CCPD/CCPD2019/ccpd_base --plate-type $p --num 2 --output ./output/preview/$p --preview --seed 42 2>&1 | tail -3
echo ""
done
echo "DONE"
tools/cut_lchannel/dataset.py
+131
View File
@@ -0,0 +1,131 @@
# -*- coding: utf-8 -*-
"""
L通道车牌数据集加载器
=====================
Domain A: 合成车牌 (140×440) → 提取 L 通道
Domain B: 真实车牌 (140×440) → 提取 L 通道
"""
import os
import random
from pathlib import Path
from typing import Tuple
import cv2
import numpy as np
import torch
from torch.utils.data import Dataset
class LChannelPlateDataset(Dataset):
"""
车牌 L 通道数据集。
Args:
root: 图片根目录
img_size: 统一输出尺寸 (W, H)
is_real: True=真实车牌域, False=合成车牌域
"""
def __init__(self, root: str, img_size=(440, 140), is_real=False):
self.root = Path(root)
self.img_size = img_size
self.is_real = is_real
self.paths = []
if self.root.is_dir():
for ext in ('*.jpg', '*.jpeg', '*.png'):
self.paths.extend(self.root.glob(ext))
def __len__(self):
return len(self.paths)
def __getitem__(self, idx) -> torch.Tensor:
path = self.paths[idx]
img = cv2.imread(str(path))
if img is None:
return self.__getitem__((idx + 1) % len(self))
# 统一尺寸
img = cv2.resize(img, self.img_size, interpolation=cv2.INTER_LINEAR)
# 提取 L 通道: BGR → LAB → L
lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
l_channel = lab[:, :, 0].astype(np.float32) # (H, W) [0, 255]
# 归一化到 [-1, 1]
l_channel = (l_channel / 127.5) - 1.0
return torch.from_numpy(l_channel).unsqueeze(0) # (1, H, W)
def make_loader(root: str, batch_size=4, img_size=(440, 140),
is_real=False, shuffle=True, num_workers=2):
"""创建 DataLoader"""
from torch.utils.data import DataLoader
ds = LChannelPlateDataset(root, img_size, is_real)
return DataLoader(ds, batch_size=batch_size, shuffle=shuffle,
num_workers=num_workers, pin_memory=True,
drop_last=True)
class LChannelInference:
"""
L通道 CUT 推理封装：合成车牌 → 真实风格车牌。
用法:
infer = LChannelInference('checkpoint.pth')
result_bgr = infer(synthetic_plate_bgr)
"""
def __init__(self, checkpoint_path: str, device=None, n_res=6):
from models import LChannelGenerator, init_weights
self.device = device or torch.device(
"cuda" if torch.cuda.is_available() else "cpu"
)
self.G = LChannelGenerator(n_res=n_res).to(self.device)
ckpt = torch.load(checkpoint_path, map_location=self.device)
self.G.load_state_dict(ckpt.get('G_state', ckpt))
self.G.eval()
# 期望输入尺寸
self.img_size = (440, 140)
def __call__(self, plate_bgr: np.ndarray) -> np.ndarray:
"""
Args:
plate_bgr: (H, W, 3) 合成车牌 BGR
Returns:
真实风格车牌 BGR (H, W, 3)
"""
h, w = plate_bgr.shape[:2]
img = cv2.resize(plate_bgr, self.img_size, interpolation=cv2.INTER_LINEAR)
# BGR → LAB
lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB).astype(np.float32)
# L → tensor → G → tensor → L
l_tensor = torch.from_numpy(lab[:, :, 0] / 127.5 - 1.0)
l_tensor = l_tensor.unsqueeze(0).unsqueeze(0).to(self.device) # (1,1,H,W)
with torch.no_grad():
l_out = self.G(l_tensor).squeeze().cpu().numpy() # (H, W)
# 还原 L 到 [0, 255]
l_out = (l_out + 1.0) * 127.5
l_out = np.clip(l_out, 0, 255)
lab[:, :, 0] = l_out
# a, b 通道不变 → 保护颜色
result = cv2.cvtColor(lab.astype(np.uint8), cv2.COLOR_LAB2BGR)
# 还原原始尺寸
if (h, w) != (self.img_size[1], self.img_size[0]):
result = cv2.resize(result, (w, h), interpolation=cv2.INTER_LINEAR)
return result
tools/cut_lchannel/models.py
+217
View File
@@ -0,0 +1,217 @@
# -*- coding: utf-8 -*-
"""
CUT L通道 模型定义
==================
轻量级 Generator + PatchGAN Discriminator + PatchNCE 对比损失。
仅处理单通道 L (亮度) 图像，输入尺寸 140×440。
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
# ============================================================
# Generator: UNet 风格 encoder-decoder
# ============================================================
class ResBlock(nn.Module):
"""残差块"""
def __init__(self, ch):
super().__init__()
self.conv = nn.Sequential(
nn.ReflectionPad2d(1),
nn.Conv2d(ch, ch, 3),
nn.InstanceNorm2d(ch),
nn.ReLU(True),
nn.ReflectionPad2d(1),
nn.Conv2d(ch, ch, 3),
nn.InstanceNorm2d(ch),
)
def forward(self, x):
return x + self.conv(x)
class LChannelGenerator(nn.Module):
"""
L通道生成器: 1ch → 1ch
输入: (B, 1, 140, 440) 合成车牌L通道
输出: (B, 1, 140, 440) 真实风格L通道
"""
def __init__(self, n_res=6):
super().__init__()
# Encoder
self.enc1 = self._conv_block(1, 64, 7, 1, 3) # 140×440
self.enc2 = self._conv_block(64, 128, 3, 2, 1) # 70×220
self.enc3 = self._conv_block(128, 256, 3, 2, 1) # 35×110
# Bottleneck
self.res_blocks = nn.Sequential(*[ResBlock(256) for _ in range(n_res)])
# Decoder
self.dec2 = self._upconv_block(256, 128) # 70×220
self.dec1 = self._upconv_block(128, 64) # 140×440
self.out = nn.Sequential(
nn.ReflectionPad2d(3),
nn.Conv2d(64, 1, 7),
nn.Tanh(),
)
def _conv_block(self, ci, co, k, s, p):
return nn.Sequential(
nn.ReflectionPad2d(p),
nn.Conv2d(ci, co, k, s),
nn.InstanceNorm2d(co),
nn.ReLU(True),
)
def _upconv_block(self, ci, co):
return nn.Sequential(
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
nn.ReflectionPad2d(1),
nn.Conv2d(ci, co, 3),
nn.InstanceNorm2d(co),
nn.ReLU(True),
)
def forward(self, x, return_features=False):
"""
Args:
x: (B, 1, H, W) 归一化到 [-1, 1]
return_features: 返回中间特征用于 PatchNCE loss
"""
features = []
e1 = self.enc1(x) # (B, 64, 140, 440)
features.append(e1)
e2 = self.enc2(e1) # (B, 128, 70, 220)
features.append(e2)
e3 = self.enc3(e2) # (B, 256, 35, 110)
features.append(e3)
b = self.res_blocks(e3) # (B, 256, 35, 110)
d2 = self.dec2(b) # (B, 128, 70, 220)
d1 = self.dec1(d2) # (B, 64, 140, 440)
out = self.out(d1) # (B, 1, 140, 440)
if return_features:
return out, features
return out
# ============================================================
# Discriminator: PatchGAN 70×70
# ============================================================
class PatchGANDiscriminator(nn.Module):
"""
PatchGAN 判别器，70×70 感受野。
输入: (B, 1, 140, 440) L通道
输出: (B, 1, H/16, W/16) 真假概率图
"""
def __init__(self):
super().__init__()
self.model = nn.Sequential(
nn.Conv2d(1, 64, 4, 2, 1), # 70×220
nn.LeakyReLU(0.2, True),
nn.Conv2d(64, 128, 4, 2, 1), # 35×110
nn.InstanceNorm2d(128),
nn.LeakyReLU(0.2, True),
nn.Conv2d(128, 256, 4, 2, 1), # 17×55
nn.InstanceNorm2d(256),
nn.LeakyReLU(0.2, True),
nn.Conv2d(256, 512, 4, 1, 1), # 16×54
nn.InstanceNorm2d(512),
nn.LeakyReLU(0.2, True),
nn.Conv2d(512, 1, 4, 1, 1), # 15×53
)
def forward(self, x):
return self.model(x)
# ============================================================
# PatchNCE 对比学习
# ============================================================
class PatchNCELoss(nn.Module):
"""
PatchNCE 对比损失。
对编码器的每一层特征图做 patch-wise 对比学习。
"""
def __init__(self, nce_layers=(0, 1, 2), nce_t=0.07, n_patches=256):
super().__init__()
self.nce_layers = nce_layers
self.nce_t = nce_t
self.n_patches = n_patches
self.cross_entropy = nn.CrossEntropyLoss()
def forward(self, feat_q, feat_k):
"""
Args:
feat_q: 生成图编码特征列表 [(B,C,H,W), ...]
feat_k: 原图编码特征列表 [(B,C,H,W), ...]
Returns:
loss: scalar
"""
total_loss = 0.0
n_layers = 0
for layer_idx in self.nce_layers:
if layer_idx >= len(feat_q) or layer_idx >= len(feat_k):
continue
fq = feat_q[layer_idx] # (B, C, H, W)
fk = feat_k[layer_idx] # (B, C, H, W)
B, C, H, W = fq.shape
n_patches = min(self.n_patches, H * W)
# 随机采样 patch 位置
idx_h = torch.randint(0, H, (n_patches,), device=fq.device)
idx_w = torch.randint(0, W, (n_patches,), device=fq.device)
# query patches
q = fq[:, :, idx_h, idx_w] # (B, C, N)
q = q.permute(0, 2, 1).reshape(B * n_patches, C) # (B*N, C)
# key patches (positive pairs)
k = fk[:, :, idx_h, idx_w] # (B, C, N)
k = k.permute(0, 2, 1).reshape(B * n_patches, C) # (B*N, C)
# 负样本：同一 batch 内其他图片的同一位置
q_norm = F.normalize(q, dim=1) # (B*N, C)
k_norm = F.normalize(k, dim=1) # (B*N, C)
# 计算相似度矩阵
logits = torch.mm(q_norm, k_norm.T) / self.nce_t # (B*N, B*N)
labels = torch.arange(B * n_patches, device=q.device)
loss = self.cross_entropy(logits, labels)
total_loss += loss
n_layers += 1
return total_loss / max(n_layers, 1)
# ============================================================
# 模型工具
# ============================================================
def init_weights(model):
"""Xavier 初始化"""
for m in model.modules():
if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
nn.init.xavier_normal_(m.weight, 0.02)
if m.bias is not None:
nn.init.constant_(m.bias, 0)
def count_params(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
tools/cut_lchannel/requirements.txt
+6
View File
@@ -0,0 +1,6 @@
torch>=1.12
torchvision
opencv-python>=4.5
numpy
tqdm
tensorboard
tools/cut_lchannel/train.py
+204
View File
@@ -0,0 +1,204 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
CUT L通道 训练脚本
==================
仅对车牌 L (亮度) 通道做纹理迁移训练，a,b 颜色通道不动。
用法:
python train.py \
--domain-a ./gan_data/domain_A_synthetic/police \
--domain-b ./gan_data/domain_B_real \
--epochs 100 \
--batch 4 \
--output ./output/cut_police
"""
import os
import sys
import argparse
import time
from pathlib import Path
import torch
import torch.nn as nn
import torch.optim as optim
# 添加 tools 目录到路径
sys.path.insert(0, str(Path(__file__).resolve().parent))
from models import (
LChannelGenerator, PatchGANDiscriminator,
PatchNCELoss, init_weights, count_params,
)
from dataset import make_loader
def train(args):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")
# ---- 数据 ----
loader_a = make_loader(args.domain_a, args.batch, (args.img_w, args.img_h),
is_real=False, num_workers=args.workers)
loader_b = make_loader(args.domain_b, args.batch, (args.img_w, args.img_h),
is_real=True, num_workers=args.workers)
print(f"Domain A (合成): {len(loader_a.dataset)} 张")
print(f"Domain B (真实): {len(loader_b.dataset)} 张")
print(f"Batch size: {args.batch}")
print(f"Epochs: {args.epochs}")
# ---- 模型 ----
G = LChannelGenerator(n_res=args.n_res).to(device)
D = PatchGANDiscriminator().to(device)
init_weights(G)
init_weights(D)
print(f"G params: {count_params(G):,}")
print(f"D params: {count_params(D):,}")
# ---- 损失 & 优化器 ----
criterion_gan = nn.MSELoss() # LSGAN
criterion_nce = PatchNCELoss(
nce_layers=args.nce_layers,
nce_t=args.nce_t,
n_patches=args.n_patches,
).to(device)
criterion_idt = nn.L1Loss()
opt_G = optim.Adam(G.parameters(), lr=args.lr, betas=(0.5, 0.999))
opt_D = optim.Adam(D.parameters(), lr=args.lr, betas=(0.5, 0.999))
# ---- 日志 ----
os.makedirs(args.output, exist_ok=True)
# ---- 训练循环 ----
for epoch in range(1, args.epochs + 1):
epoch_start = time.time()
G_loss_total = 0.0
D_loss_total = 0.0
n_batches = 0
for real_a, real_b in zip(loader_a, loader_b):
real_a = real_a.to(device) # (B, 1, H, W) 合成L
real_b = real_b.to(device) # (B, 1, H, W) 真实L
B = real_a.size(0)
# ---- 1. 更新 D ----
fake_b = G(real_a).detach()
# G 不用返回特征
pred_real = D(real_b)
pred_fake = D(fake_b)
loss_D_real = criterion_gan(pred_real,
torch.ones_like(pred_real))
loss_D_fake = criterion_gan(pred_fake,
torch.zeros_like(pred_fake))
loss_D = (loss_D_real + loss_D_fake) * 0.5
opt_D.zero_grad()
loss_D.backward()
opt_D.step()
# ---- 2. 更新 G ----
fake_b, feat_q = G(real_a, return_features=True)
# GAN loss
pred_fake = D(fake_b)
loss_G_gan = criterion_gan(pred_fake,
torch.ones_like(pred_fake))
# PatchNCE loss (对比学习)
_, feat_k = G(real_b, return_features=True)
loss_G_nce = criterion_nce(feat_q, feat_k)
# Identity loss: 裁剪输出到输入尺寸（UNet可能多出几像素）
identity = G(real_a)
if identity.shape != real_a.shape:
identity = identity[:, :, :real_a.shape[2], :real_a.shape[3]]
loss_G_idt = criterion_idt(identity, real_a) * args.lambda_idt
loss_G = loss_G_gan + loss_G_nce * args.lambda_nce + loss_G_idt
opt_G.zero_grad()
loss_G.backward()
opt_G.step()
G_loss_total += loss_G.item()
D_loss_total += loss_D.item()
n_batches += 1
# ---- Epoch 统计 ----
G_avg = G_loss_total / max(n_batches, 1)
D_avg = D_loss_total / max(n_batches, 1)
elapsed = time.time() - epoch_start
print(f"Epoch {epoch:3d}/{args.epochs} | "
f"G: {G_avg:.4f} D: {D_avg:.4f} | "
f"{elapsed:.1f}s")
pass # 日志仅print，不使用tensorboard
# ---- 保存 ----
if epoch % args.save_every == 0:
ckpt = {
'epoch': epoch,
'G_state': G.state_dict(),
'D_state': D.state_dict(),
'opt_G': opt_G.state_dict(),
'opt_D': opt_D.state_dict(),
}
torch.save(ckpt, os.path.join(args.output, f'epoch_{epoch}.pth'))
# 最终保存
final_ckpt = {
'G_state': G.state_dict(),
'D_state': D.state_dict(),
}
torch.save(final_ckpt, os.path.join(args.output, 'final.pth'))
print(f"\n训练完成! 模型保存在: {args.output}")
def main():
parser = argparse.ArgumentParser(description='CUT L通道训练')
# 数据
parser.add_argument('--domain-a', type=str, required=True,
help='合成车牌目录')
parser.add_argument('--domain-b', type=str, required=True,
help='真实车牌目录')
# 训练
parser.add_argument('--epochs', type=int, default=100)
parser.add_argument('--batch', type=int, default=4)
parser.add_argument('--lr', type=float, default=2e-4)
parser.add_argument('--workers', type=int, default=2)
# 模型
parser.add_argument('--n-res', type=int, default=6,
help='残差块数量')
# 损失权重
parser.add_argument('--lambda-nce', type=float, default=1.0)
parser.add_argument('--lambda-idt', type=float, default=0.5)
# PatchNCE
parser.add_argument('--nce-layers', type=int, nargs='+', default=[0, 1, 2])
parser.add_argument('--nce-t', type=float, default=0.07)
parser.add_argument('--n-patches', type=int, default=128)
# 保存
parser.add_argument('--output', type=str, default='./output/cut_train')
parser.add_argument('--img-w', type=int, default=220)
parser.add_argument('--img-h', type=int, default=70)
parser.add_argument('--save-every', type=int, default=20)
args = parser.parse_args()
train(args)
if __name__ == '__main__':
main()
tools/extract_real_plates.py
+145
View File
@@ -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()
tools/gen_pure_plates.py
+28
View File
@@ -0,0 +1,28 @@
#!/usr/bin/env python3
"""生成纯合成车牌（无场景背景）作为 CUT Domain A"""
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'chinese_license_plate_generator'))
from generate_multi_plate import MultiPlateGenerator
import cv2, random
import numpy as np
random.seed(42)
np.random.seed(42)
g = MultiPlateGenerator(
'chinese_license_plate_generator/plate_model',
'chinese_license_plate_generator/font_model'
)
out = 'gan_data/domain_A_synthetic/police_pure'
os.makedirs(out, exist_ok=True)
for i in range(500):
num, bg, _ = g.generate_plate_number()
# 强制police类型
if '警' not in num:
num = num[:6] + '警'
img = g.generate_plate_special(num, 'white', False)
cv2.imwrite(f'{out}/synth_{i:04d}.jpg', img)
print(f'Done: {out} ({len(os.listdir(out))} images)')
tools/plate_swap_pipeline.py
+909
View File
@@ -0,0 +1,909 @@
#!/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()