From 3808b9f03596f9ad5860e8ed353f375fe860e6f7 Mon Sep 17 00:00:00 2001 From: LEE <484836843@qq.com> Date: Sat, 20 Jun 2026 19:49:52 +0800 Subject: [PATCH] =?UTF-8?q?=E7=89=B9=E6=AE=8A=E8=BD=A6=E7=89=8C=E6=95=B0?= =?UTF-8?q?=E6=8D=AE=E9=9B=86=E6=9E=84=E5=BB=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tools/_gan_prep.sh | 42 ++ tools/_p20.sh | 5 + tools/_r3.sh | 9 + tools/cut_lchannel/dataset.py | 131 ++++ tools/cut_lchannel/models.py | 217 +++++++ tools/cut_lchannel/requirements.txt | 6 + tools/cut_lchannel/train.py | 204 +++++++ tools/extract_real_plates.py | 145 +++++ tools/gen_pure_plates.py | 28 + tools/plate_swap_pipeline.py | 909 ++++++++++++++++++++++++++++ 10 files changed, 1696 insertions(+) create mode 100644 tools/_gan_prep.sh create mode 100644 tools/_p20.sh create mode 100644 tools/_r3.sh create mode 100644 tools/cut_lchannel/dataset.py create mode 100644 tools/cut_lchannel/models.py create mode 100644 tools/cut_lchannel/requirements.txt create mode 100644 tools/cut_lchannel/train.py create mode 100644 tools/extract_real_plates.py create mode 100644 tools/gen_pure_plates.py create mode 100644 tools/plate_swap_pipeline.py diff --git a/tools/_gan_prep.sh b/tools/_gan_prep.sh new file mode 100644 index 0000000..ebdb6d1 --- /dev/null +++ b/tools/_gan_prep.sh @@ -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 diff --git a/tools/_p20.sh b/tools/_p20.sh new file mode 100644 index 0000000..f60c3d7 --- /dev/null +++ b/tools/_p20.sh @@ -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 diff --git a/tools/_r3.sh b/tools/_r3.sh new file mode 100644 index 0000000..258d938 --- /dev/null +++ b/tools/_r3.sh @@ -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" \ No newline at end of file diff --git a/tools/cut_lchannel/dataset.py b/tools/cut_lchannel/dataset.py new file mode 100644 index 0000000..d61ce8d --- /dev/null +++ b/tools/cut_lchannel/dataset.py @@ -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 diff --git a/tools/cut_lchannel/models.py b/tools/cut_lchannel/models.py new file mode 100644 index 0000000..6f9bdb8 --- /dev/null +++ b/tools/cut_lchannel/models.py @@ -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) diff --git a/tools/cut_lchannel/requirements.txt b/tools/cut_lchannel/requirements.txt new file mode 100644 index 0000000..38b7f08 --- /dev/null +++ b/tools/cut_lchannel/requirements.txt @@ -0,0 +1,6 @@ +torch>=1.12 +torchvision +opencv-python>=4.5 +numpy +tqdm +tensorboard diff --git a/tools/cut_lchannel/train.py b/tools/cut_lchannel/train.py new file mode 100644 index 0000000..c348ef1 --- /dev/null +++ b/tools/cut_lchannel/train.py @@ -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() diff --git a/tools/extract_real_plates.py b/tools/extract_real_plates.py new file mode 100644 index 0000000..c98d511 --- /dev/null +++ b/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() diff --git a/tools/gen_pure_plates.py b/tools/gen_pure_plates.py new file mode 100644 index 0000000..3090abd --- /dev/null +++ b/tools/gen_pure_plates.py @@ -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)') diff --git a/tools/plate_swap_pipeline.py b/tools/plate_swap_pipeline.py new file mode 100644 index 0000000..01343c4 --- /dev/null +++ b/tools/plate_swap_pipeline.py @@ -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()