feat: 特殊车牌生成管线 + C++/多进程/色温校正
This commit is contained in:
kirl committed
1 parent
3808b9f035
commit
3101992f51
29 files changed
+2336
-80
No files matched your search
@@ -0,0 +1,83 @@
|
||||
# CCPD_SPE_GEN — 特殊车牌数据集生成工具
|
||||
|
||||
基于 CCPD 真实车辆照片，将普通蓝牌替换为 9 种特殊车牌（警车/军车/港澳/使馆/领馆/教练/新能源）。
|
||||
|
||||
---
|
||||
|
||||
## 目录结构
|
||||
|
||||
```
|
||||
CCPD_SPE_GEN/
|
||||
├── README.md
|
||||
│
|
||||
├── cv_pipeline/ # 传统 CV 管线（无需 GPU/训练）
|
||||
│ ├── plate_swap.py # 主体：车牌替换全流程
|
||||
│ ├── extract_real_plates.py # 从 CCPD 提取真实车牌 ROI
|
||||
│ └── gen_pure_plates.py # 生成纯合成车牌（无背景）
|
||||
│
|
||||
├── gan_pipeline/ # GAN 管线（CUT 纹理迁移）
|
||||
│ ├── dataset.py # L 通道 DataLoader + 推理封装
|
||||
│ ├── models.py # Generator + PatchGAN + PatchNCE Loss
|
||||
│ ├── train.py # 训练脚本
|
||||
│ └── requirements.txt # PyTorch 依赖
|
||||
│
|
||||
└── scripts/ # 一键运行脚本
|
||||
├── quick_test.sh # 快速测试 20 张警车牌
|
||||
├── preview_all.sh # 预览全部 9 种类型
|
||||
└── gan_data_prep.sh # GAN 训练数据准备
|
||||
```
|
||||
|
||||
## 两套算法
|
||||
|
||||
### CV 管线 (`cv_pipeline/`)
|
||||
**经典计算机视觉，零训练成本，即开即用。**
|
||||
|
||||
流程：CCPD 原图 → 解析文件名角点 → 生成特殊车牌 → Lab 色彩对齐 → 透视回贴 → 距离变换羽化 → 噪声/模糊匹配 → JPEG 退化模拟 → 输出
|
||||
|
||||
```bash
|
||||
cd ICE_Dataset
|
||||
python CCPD_SPE_GEN/cv_pipeline/plate_swap.py \
|
||||
--ccpd-dir ./CCPD/CCPD2019/ccpd_base \
|
||||
--plate-type police \
|
||||
--num 500 \
|
||||
--output ./output/police_dataset
|
||||
```
|
||||
|
||||
### GAN 管线 (`gan_pipeline/`)
|
||||
**基于 CUT (Contrastive Unpaired Translation)，需 PyTorch + GPU 训练。**
|
||||
|
||||
仅对车牌 L (亮度) 通道做纹理迁移，保留 AB 颜色通道不变。训练后将 GAN 推理接入 CV 管线，替换硬编码的退化模拟步骤。
|
||||
|
||||
```bash
|
||||
# 训练
|
||||
cd CCPD_SPE_GEN/gan_pipeline
|
||||
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
|
||||
|
||||
# 推理 (在 CV 管线中调用)
|
||||
# from CCPD_SPE_GEN.gan_pipeline.dataset import LChannelInference
|
||||
```
|
||||
|
||||
## 依赖
|
||||
|
||||
| 管线 | 依赖 |
|
||||
|------|------|
|
||||
| CV | `numpy`, `opencv-python`, `tqdm`, `chinese_license_plate_generator/` |
|
||||
| GAN | 以上 + `torch>=1.12`, `torchvision`, `tensorboard` |
|
||||
|
||||
## 支持的特殊车牌类型
|
||||
|
||||
| 类型 | 底板 | 说明 |
|
||||
|------|------|------|
|
||||
| `police` | 白底 | 警用车牌 (末位"警"红字) |
|
||||
| `army` | 白底 | 军车/武警 |
|
||||
| `gangao` | 黑底 | 港澳入出境 |
|
||||
| `shiguan_old` | 黑底 | 使馆旧版 ("使"红字) |
|
||||
| `shiguan_new` | 黑底 | 使馆新版 (白字) |
|
||||
| `lingguan_old` | 黑底 | 领馆旧版 ("领"红字) |
|
||||
| `lingguan_new` | 黑底 | 领馆新版 (白字) |
|
||||
| `coach` | 黄底 | 教练车 (末位"学") |
|
||||
| `new_energy` | 黄绿双拼 | 大型新能源 (8 位) |
|
||||
@@ -0,0 +1,24 @@
|
||||
cmake_minimum_required(VERSION 3.16)
|
||||
project(ice_plate_swap LANGUAGES CXX)
|
||||
|
||||
set(CMAKE_CXX_STANDARD 17)
|
||||
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
||||
|
||||
# OpenCV
|
||||
find_package(OpenCV REQUIRED)
|
||||
|
||||
# 源文件
|
||||
add_executable(plate_swap
|
||||
main.cpp
|
||||
plate_swap.cpp
|
||||
plate_swap.hpp
|
||||
)
|
||||
|
||||
target_link_libraries(plate_swap ${OpenCV_LIBS})
|
||||
|
||||
# 优化选项
|
||||
if(MSVC)
|
||||
target_compile_options(plate_swap PRIVATE /O2 /arch:AVX2 /utf-8)
|
||||
else()
|
||||
target_compile_options(plate_swap PRIVATE -O3 -march=native -fopenmp)
|
||||
endif()
|
||||
@@ -0,0 +1,42 @@
|
||||
# C++ 管线编译脚本
|
||||
# 前置条件: Visual Studio 2022 + C++ workload + OpenCV
|
||||
# OpenCV 安装: 下载 https://opencv.org/releases/ 解压到 C:\opencv
|
||||
|
||||
param(
|
||||
[string]$OpenCVDir = "C:\opencv\build"
|
||||
)
|
||||
|
||||
$ErrorActionPreference = "Stop"
|
||||
|
||||
Write-Host "=== ICE C++ Plate Swap 编译 ===" -ForegroundColor Cyan
|
||||
|
||||
# 检查 OpenCV
|
||||
if (-not (Test-Path "$OpenCVDir\OpenCVConfig.cmake")) {
|
||||
Write-Host "[ERROR] OpenCV 未找到: $OpenCVDir" -ForegroundColor Red
|
||||
Write-Host " 请从 https://opencv.org/releases/ 下载 Windows 版"
|
||||
Write-Host " 解压到 C:\opencv 或使用 -OpenCVDir 指定路径"
|
||||
exit 1
|
||||
}
|
||||
|
||||
# 创建 build 目录
|
||||
$buildDir = "build"
|
||||
if (Test-Path $buildDir) { Remove-Item -Recurse -Force $buildDir }
|
||||
New-Item -ItemType Directory $buildDir | Out-Null
|
||||
|
||||
# CMake 配置 (使用 VS 2022 生成器)
|
||||
Write-Host "CMake 配置..."
|
||||
cmake -B $buildDir -G "Visual Studio 17 2022" -A x64 `
|
||||
-DOpenCV_DIR="$OpenCVDir"
|
||||
|
||||
# 编译 Release
|
||||
Write-Host "编译 Release..."
|
||||
cmake --build $buildDir --config Release --parallel
|
||||
|
||||
Write-Host ""
|
||||
Write-Host "=== 编译完成 ===" -ForegroundColor Green
|
||||
Write-Host "可执行文件: $buildDir\Release\plate_swap.exe"
|
||||
|
||||
# 使用示例
|
||||
Write-Host ""
|
||||
Write-Host "用法示例:"
|
||||
Write-Host " .\build\Release\plate_swap.exe --ccpd-dir ./CCPD/CCPD2019/ccpd_base --plate-type police --plate-pool ./output/plate_pool --num 500 --workers 32 --fast"
|
||||
@@ -0,0 +1,123 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
预生成合成车牌池 (供 C++ 管线使用)
|
||||
===================================
|
||||
为每种特殊车牌类型批量生成纯车牌图，C++ 运行时直接读内存池，
|
||||
避免每张 CCPD 原图都重新渲染车牌。
|
||||
|
||||
用法:
|
||||
python gen_plate_pool.py --num 5000 --output ./output/plate_pool
|
||||
|
||||
输出目录结构:
|
||||
plate_pool/
|
||||
yellow/ 教练车牌 (440x140)
|
||||
green_truck/ 新能源车牌 (480x140)
|
||||
white/ 警车车牌
|
||||
white_army/ 军车车牌
|
||||
black/ 港澳/领馆车牌
|
||||
black_shi/ 使馆车牌
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
|
||||
# 路径
|
||||
_SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
_PROJECT_DIR = _SCRIPT_DIR.parent.parent
|
||||
_GENERATOR_DIR = _PROJECT_DIR / 'chinese_license_plate_generator'
|
||||
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,
|
||||
)
|
||||
|
||||
# 类型配置
|
||||
PLATE_TYPES = [
|
||||
('police', 'white', generate_plate_number_white),
|
||||
('army', 'white_army', generate_plate_number_army),
|
||||
('gangao', 'black', generate_plate_number_black_gangao),
|
||||
('shiguan_old', 'black_shi', generate_plate_number_black_shi_old),
|
||||
('shiguan_new', 'black_shi', generate_plate_number_black_shi_new),
|
||||
('lingguan_old','black', generate_plate_number_black_ling_old),
|
||||
('lingguan_new','black', generate_plate_number_black_ling_new),
|
||||
('coach', 'yellow', generate_plate_number_yellow_xue),
|
||||
('new_energy', 'green_truck', lambda: generate_plate_number_blue(length=8)),
|
||||
]
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='预生成合成车牌池')
|
||||
parser.add_argument('--num', type=int, default=5000,
|
||||
help='每种类型生成数量 (默认: 5000)')
|
||||
parser.add_argument('--output', type=str, default='./output/plate_pool',
|
||||
help='输出根目录')
|
||||
args = parser.parse_args()
|
||||
|
||||
# 初始化生成器
|
||||
pm = str(_GENERATOR_DIR / 'plate_model')
|
||||
fm = str(_GENERATOR_DIR / 'font_model')
|
||||
gen = MultiPlateGenerator(pm, fm)
|
||||
|
||||
for plate_type, bg_color, num_func in PLATE_TYPES:
|
||||
out_dir = Path(args.output) / bg_color
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
existing = len(list(out_dir.glob('*.jpg')))
|
||||
if existing >= args.num:
|
||||
print(f"[SKIP] {plate_type:12s} ({bg_color:12s}): 已有 {existing} 张")
|
||||
continue
|
||||
|
||||
need = args.num - existing
|
||||
print(f"[GEN] {plate_type:12s} ({bg_color:12s}): {need} 张...", end=' ', flush=True)
|
||||
|
||||
success = 0
|
||||
pbar = tqdm(total=need, desc=f'{plate_type}', unit='张', leave=False)
|
||||
attempts = 0
|
||||
while success < need and attempts < need * 3:
|
||||
attempts += 1
|
||||
try:
|
||||
number = num_func()
|
||||
plate = gen.generate_plate_special(
|
||||
number, bg_color, is_double=False, enhance=False
|
||||
)
|
||||
# 用 imencode 支持中文文件名
|
||||
fname = f"{number}.jpg"
|
||||
ok, buf = cv2.imencode('.jpg', plate, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||||
if ok:
|
||||
with open(str(out_dir / fname), 'wb') as f:
|
||||
f.write(buf.tobytes())
|
||||
success += 1
|
||||
pbar.update(1)
|
||||
except Exception:
|
||||
pass
|
||||
pbar.close()
|
||||
print(f"OK ({success} 张)")
|
||||
|
||||
print(f"\n完成! 输出: {args.output}")
|
||||
# 统计
|
||||
for bg_color in ['yellow','green_truck','white','white_army','black','black_shi']:
|
||||
d = Path(args.output) / bg_color
|
||||
count = len(list(d.glob('*.jpg'))) if d.is_dir() else 0
|
||||
print(f" {bg_color:15s}: {count} 张")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
import os
|
||||
os.chdir(str(_PROJECT_DIR))
|
||||
main()
|
||||
@@ -0,0 +1,145 @@
|
||||
# ICE Plate Swap — C++ 高性能特殊车牌替换引擎
|
||||
|
||||
基于 CCPD 真实街景照片，将普通蓝牌替换为 9 种特殊车牌。
|
||||
C++ / OpenCV 实现，32 线程满血，~53 张/秒。
|
||||
|
||||
---
|
||||
|
||||
## 快速开始
|
||||
|
||||
### 1. 环境要求
|
||||
|
||||
| 组件 | 说明 |
|
||||
|------|------|
|
||||
| Windows 10/11 x64 | — |
|
||||
| [Visual Studio 2022 BuildTools](https://visualstudio.microsoft.com/downloads/#build-tools-for-visual-studio-2022) | 含 C++ 桌面开发 workload |
|
||||
| [OpenCV 4.10.0](https://github.com/opencv/opencv/releases/tag/4.10.0) | 解压到 `C:\opencv` |
|
||||
|
||||
### 2. 预编译版 (推荐)
|
||||
|
||||
直接下载 `plate_swap.exe`（已通过 Git LFS 管理）：
|
||||
|
||||
```powershell
|
||||
# 设置 OpenCV DLL 路径
|
||||
$env:PATH = "C:\opencv\build\x64\vc16\bin;$env:PATH"
|
||||
|
||||
# 运行
|
||||
.\install\plate_swap.exe --help
|
||||
```
|
||||
|
||||
### 3. 从源码编译
|
||||
|
||||
```powershell
|
||||
# 打开 VS Developer Command Prompt
|
||||
# 然后:
|
||||
cd CCPD_SPE_GEN\cpp_pipeline
|
||||
cmake -B build -G "Visual Studio 17 2022" -A x64 -DOpenCV_DIR=C:\opencv\build
|
||||
cmake --build build --config Release --parallel
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 用法
|
||||
|
||||
```powershell
|
||||
$env:PATH = "C:\opencv\build\x64\vc16\bin;$env:PATH"
|
||||
|
||||
.\plate_swap.exe `
|
||||
--ccpd-dir ./CCPD/CCPD2019/ccpd_base,./CCPD/CCPD2019/ccpd_tilt `
|
||||
--plate-type police `
|
||||
--plate-pool ./output/plate_pool_v2 `
|
||||
--output ./CCPD_SPE `
|
||||
--num 1000 `
|
||||
--workers 24
|
||||
```
|
||||
|
||||
### 参数
|
||||
|
||||
| 参数 | 必需 | 默认 | 说明 |
|
||||
|------|------|------|------|
|
||||
| `--ccpd-dir` | ✅ | — | CCPD 图片目录，多个用逗号分隔 |
|
||||
| `--plate-type` | ✅ | — | 特殊车牌类型 |
|
||||
| `--plate-pool` | ✅ | — | 预生成车牌池路径（含 `pool.bin`） |
|
||||
| `--output` | ✅ | — | 输出根目录 |
|
||||
| `--num` | | 500 | 生成图片数量 |
|
||||
| `--workers` | | 全部核心 | 并行线程数 |
|
||||
| `--seed` | | 42 | 随机种子 |
|
||||
| `--fast` | | false | 快速模式（跳过退化模拟，不建议） |
|
||||
| `--h-shift` | | 0 | 底板色相偏移 (GUI调色) |
|
||||
| `--s-scale` | | 100 | 底板饱和度缩放% |
|
||||
| `--v-scale` | | 100 | 底板明度缩放% |
|
||||
|
||||
### 支持的类型
|
||||
|
||||
```
|
||||
police 警用车牌 (白底, JING)
|
||||
army 军车/武警 (白底)
|
||||
gangao 港澳入出境 (黑底, GANG/AO)
|
||||
shiguan_old 使馆旧版 (黑底, SHI)
|
||||
shiguan_new 使馆新版 (黑底)
|
||||
lingguan_old 领馆旧版 (黑底, LING)
|
||||
lingguan_new 领馆新版 (黑底)
|
||||
coach 教练车 (黄底, XUE)
|
||||
new_energy 大型新能源 (黄绿双拼, 8位)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 预生成车牌池
|
||||
|
||||
```powershell
|
||||
cd ICE_Dataset
|
||||
python CCPD_SPE_GEN\cpp_pipeline\gen_plate_pool.py --num 1000 --output ./output/plate_pool_v2
|
||||
```
|
||||
|
||||
生成 `pool.bin` (~1.6GB)，含 9 种类型各 1000 张 + 车牌号元数据。
|
||||
|
||||
---
|
||||
|
||||
## 输出结构
|
||||
|
||||
```
|
||||
CCPD_SPE/
|
||||
├── ccpd_base/
|
||||
│ ├── police/
|
||||
│ ├── gangao/
|
||||
│ └── ...
|
||||
├── ccpd_tilt/
|
||||
└── ...
|
||||
```
|
||||
|
||||
文件名: `{CCPD原名}__{类型}__{车牌号拼音}.jpg`
|
||||
|
||||
---
|
||||
|
||||
## 性能
|
||||
|
||||
| 指标 | 值 |
|
||||
|------|-----|
|
||||
| 单线程 | ~7 张/秒 |
|
||||
| 24 线程 | ~53 张/秒 |
|
||||
| 35.5 万张全量 | ~1.9 小时 |
|
||||
|
||||
---
|
||||
|
||||
## 调色 GUI
|
||||
|
||||
```powershell
|
||||
python CCPD_SPE_GEN\cv_pipeline\color_tuner.py
|
||||
```
|
||||
|
||||
H/S/V 滑块实时预览 → 按 `S` 保存 → `plate_color_profile.py` → C++ 管线 `--h-shift` `--s-scale` `--v-scale` 读取。
|
||||
|
||||
---
|
||||
|
||||
## 构建
|
||||
|
||||
```powershell
|
||||
# 安装 OpenCV
|
||||
Invoke-WebRequest "https://github.com/opencv/opencv/releases/download/4.10.0/opencv-4.10.0-windows.exe" -OutFile "$env:TEMP\opencv.exe"
|
||||
Start-Process "$env:TEMP\opencv.exe" -Args "-oC:\opencv -y" -Wait
|
||||
|
||||
# 编译
|
||||
cd CCPD_SPE_GEN\cpp_pipeline
|
||||
.\build.ps1
|
||||
```
|
||||
Binary file not shown.
@@ -0,0 +1,203 @@
|
||||
#include "plate_swap.hpp"
|
||||
#include <iostream>
|
||||
#include <thread>
|
||||
#include <mutex>
|
||||
#include <atomic>
|
||||
#include <chrono>
|
||||
#include <filesystem>
|
||||
|
||||
namespace fs = std::filesystem;
|
||||
using namespace std;
|
||||
using namespace ice;
|
||||
|
||||
// ── 配置 ──
|
||||
struct Config {
|
||||
vector<string> ccpd_dirs;
|
||||
string plate_type;
|
||||
string output_dir;
|
||||
string plate_pool_dir;
|
||||
int num_images = 500;
|
||||
unsigned workers = thread::hardware_concurrency();
|
||||
bool fast = false;
|
||||
int seed = 42;
|
||||
int h_shift = 0, s_scale = 100, v_scale = 100;
|
||||
};
|
||||
|
||||
// ── 参数解析 ──
|
||||
Config parse_args(int argc, char* argv[]) {
|
||||
Config cfg;
|
||||
cfg.workers = thread::hardware_concurrency();
|
||||
|
||||
for (int i = 1; i < argc; i++) {
|
||||
string arg = argv[i];
|
||||
if (arg == "--ccpd-dir" && i+1 < argc) {
|
||||
string dirs = argv[++i];
|
||||
size_t pos = 0;
|
||||
while (pos < dirs.size()) {
|
||||
size_t comma = dirs.find(',', pos);
|
||||
string d = (comma == string::npos) ? dirs.substr(pos) : dirs.substr(pos, comma-pos);
|
||||
// trim
|
||||
while (!d.empty() && d.front() == ' ') d.erase(0,1);
|
||||
while (!d.empty() && d.back() == ' ') d.pop_back();
|
||||
if (!d.empty()) cfg.ccpd_dirs.push_back(d);
|
||||
if (comma == string::npos) break;
|
||||
pos = comma + 1;
|
||||
}
|
||||
}
|
||||
else if (arg == "--plate-type" && i+1 < argc) cfg.plate_type = argv[++i];
|
||||
else if (arg == "--output" && i+1 < argc) cfg.output_dir = argv[++i];
|
||||
else if (arg == "--plate-pool" && i+1 < argc) cfg.plate_pool_dir = argv[++i];
|
||||
else if (arg == "--num" && i+1 < argc) cfg.num_images = stoi(argv[++i]);
|
||||
else if (arg == "--workers" && i+1 < argc) cfg.workers = stoi(argv[++i]);
|
||||
else if (arg == "--seed" && i+1 < argc) cfg.seed = stoi(argv[++i]);
|
||||
else if (arg == "--fast") cfg.fast = true;
|
||||
else if (arg == "--h-shift" && i+1 < argc) cfg.h_shift = stoi(argv[++i]);
|
||||
else if (arg == "--s-scale" && i+1 < argc) cfg.s_scale = stoi(argv[++i]);
|
||||
else if (arg == "--v-scale" && i+1 < argc) cfg.v_scale = stoi(argv[++i]);
|
||||
}
|
||||
|
||||
return cfg;
|
||||
}
|
||||
|
||||
// ── 收集并预解析 CCPD 图片（过滤无效，预存角点） ──
|
||||
struct IndexedImage {
|
||||
string path;
|
||||
CCPDMeta meta;
|
||||
};
|
||||
|
||||
vector<IndexedImage> collect_and_index(const vector<string>& dirs, mt19937& rng) {
|
||||
vector<IndexedImage> result;
|
||||
for (auto& d : dirs) {
|
||||
error_code ec;
|
||||
if (!fs::is_directory(d, ec)) {
|
||||
cerr << "[WARN] Dir not found: " << d << endl;
|
||||
continue;
|
||||
}
|
||||
for (auto& entry : fs::directory_iterator(d, ec)) {
|
||||
if (ec) break;
|
||||
if (!entry.is_regular_file()) continue;
|
||||
string ext = entry.path().extension().string();
|
||||
for (auto& c : ext) c = (char)tolower(c);
|
||||
if (ext != ".jpg" && ext != ".jpeg" && ext != ".png") continue;
|
||||
|
||||
string fname = entry.path().filename().string();
|
||||
CCPDMeta meta = parse_ccpd_filename(fname);
|
||||
if (!meta.valid) continue;
|
||||
|
||||
IndexedImage ii;
|
||||
ii.path = entry.path().string();
|
||||
ii.meta = meta;
|
||||
result.push_back(ii);
|
||||
}
|
||||
}
|
||||
shuffle(result.begin(), result.end(), rng);
|
||||
return result;
|
||||
}
|
||||
|
||||
// ── 主函数 ──
|
||||
int main(int argc, char* argv[]) {
|
||||
Config cfg = parse_args(argc, argv);
|
||||
|
||||
if (cfg.ccpd_dirs.empty() || cfg.plate_type.empty()) {
|
||||
cerr << "Usage: plate_swap --ccpd-dir <dirs> --plate-type <type> [options]" << endl;
|
||||
cerr << " --plate-pool <dir> Plate pool directory (required)" << endl;
|
||||
cerr << " --num <N> Number of images (default 500)" << endl;
|
||||
cerr << " --workers <N> Thread count (default all cores)" << endl;
|
||||
cerr << " --fast Fast mode (skip degradation)" << endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
cout << "=== ICE C++ Plate Swap ===" << endl;
|
||||
cout << "Type: " << cfg.plate_type << endl;
|
||||
cout << "Threads: " << cfg.workers << endl;
|
||||
cout << "Target: " << cfg.num_images << " images" << endl;
|
||||
cout << "Fast: " << (cfg.fast ? "yes" : "no") << endl;
|
||||
|
||||
if (cfg.plate_pool_dir.empty()) {
|
||||
cerr << "[ERROR] --plate-pool is required" << endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Load plate pool
|
||||
auto t0 = chrono::steady_clock::now();
|
||||
cout << "Loading plate pool..." << flush;
|
||||
PlatePool pool;
|
||||
string bin_path = cfg.plate_pool_dir + "/pool.bin";
|
||||
if (fs::exists(bin_path))
|
||||
pool.load_bin(bin_path);
|
||||
else
|
||||
pool.load(cfg.plate_pool_dir);
|
||||
size_t pool_count = pool.count(cfg.plate_type);
|
||||
auto t1 = chrono::steady_clock::now();
|
||||
cout << " (" << pool_count << " available, "
|
||||
<< chrono::duration<double>(t1-t0).count() << "s)" << endl;
|
||||
if (pool_count == 0) {
|
||||
cerr << "[ERROR] No plates of type " << cfg.plate_type << " in pool" << endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Collect & pre-parse CCPD (filter invalid, store corners)
|
||||
mt19937 rng(cfg.seed);
|
||||
cout << "Scanning & indexing CCPD..." << flush;
|
||||
vector<IndexedImage> images = collect_and_index(cfg.ccpd_dirs, rng);
|
||||
auto t2 = chrono::steady_clock::now();
|
||||
cout << " (" << images.size() << " valid, "
|
||||
<< chrono::duration<double>(t2-t1).count() << "s)" << endl;
|
||||
if (images.empty()) {
|
||||
cerr << "[ERROR] No valid CCPD images found" << endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
fs::create_directories(cfg.output_dir);
|
||||
|
||||
SwapConfig swap_cfg;
|
||||
swap_cfg.plate_type = cfg.plate_type;
|
||||
swap_cfg.output_dir = cfg.output_dir;
|
||||
swap_cfg.fast_mode = cfg.fast;
|
||||
swap_cfg.h_shift = cfg.h_shift;
|
||||
swap_cfg.s_scale = cfg.s_scale;
|
||||
swap_cfg.v_scale = cfg.v_scale;
|
||||
|
||||
// Parallel processing with pre-parsed metadata
|
||||
atomic<int> success{0}, failed{0};
|
||||
atomic<int> idx{0};
|
||||
int total = min((int)images.size(), cfg.num_images);
|
||||
mutex cout_mtx;
|
||||
|
||||
auto t_proc = chrono::steady_clock::now();
|
||||
|
||||
vector<thread> threads;
|
||||
for (unsigned t = 0; t < cfg.workers; t++) {
|
||||
threads.emplace_back([&](unsigned seed_offset) {
|
||||
mt19937 local_rng(cfg.seed + seed_offset * 10007);
|
||||
|
||||
while (true) {
|
||||
int i = idx.fetch_add(1);
|
||||
if (i >= total) break;
|
||||
|
||||
bool ok = process_single(images[i].path, swap_cfg, pool,
|
||||
local_rng, &images[i].meta);
|
||||
if (ok) {
|
||||
int s = success.fetch_add(1) + 1;
|
||||
if (s % 100 == 0 || s == total) {
|
||||
lock_guard<mutex> lk(cout_mtx);
|
||||
cout << "\rProgress: " << s << "/" << total << flush;
|
||||
}
|
||||
} else {
|
||||
failed.fetch_add(1);
|
||||
}
|
||||
}
|
||||
}, t);
|
||||
}
|
||||
|
||||
for (auto& th : threads) th.join();
|
||||
|
||||
auto t_done = chrono::steady_clock::now();
|
||||
double elapsed = chrono::duration<double>(t_done - t_proc).count();
|
||||
|
||||
cout << "\rDone! Success: " << success << ", Failed: " << failed << endl;
|
||||
cout << "Process time: " << elapsed << "s (" << (success/max(elapsed,0.01)) << " img/s)" << endl;
|
||||
cout << "Output: " << cfg.output_dir << endl;
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,692 @@
|
||||
#include "plate_swap.hpp"
|
||||
#include <fstream>
|
||||
#include <algorithm>
|
||||
#include <numeric>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <filesystem>
|
||||
|
||||
namespace fs = std::filesystem;
|
||||
using namespace cv;
|
||||
using namespace std;
|
||||
|
||||
namespace ice {
|
||||
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
// CCPD 文件名解析
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
|
||||
CCPDMeta parse_ccpd_filename(const string& filename) {
|
||||
CCPDMeta meta;
|
||||
meta.valid = false;
|
||||
|
||||
// 提取不带扩展名的部分
|
||||
fs::path p(filename);
|
||||
string stem = p.stem().string();
|
||||
|
||||
// 按 '-' 分割
|
||||
vector<string> parts;
|
||||
size_t pos = 0;
|
||||
while (pos < stem.size()) {
|
||||
size_t dash = stem.find('-', pos);
|
||||
if (dash == string::npos) {
|
||||
parts.push_back(stem.substr(pos));
|
||||
break;
|
||||
}
|
||||
parts.push_back(stem.substr(pos, dash - pos));
|
||||
pos = dash + 1;
|
||||
}
|
||||
|
||||
if (parts.size() < 5) return meta;
|
||||
|
||||
try {
|
||||
// parts[3] = 4角点
|
||||
string& coords = parts[3];
|
||||
vector<Point2f> pts;
|
||||
size_t start = 0;
|
||||
while (start < coords.size()) {
|
||||
size_t us = coords.find('_', start);
|
||||
string pt = (us == string::npos) ? coords.substr(start) : coords.substr(start, us - start);
|
||||
size_t amp = pt.find('&');
|
||||
if (amp == string::npos) break;
|
||||
float x = stof(pt.substr(0, amp));
|
||||
float y = stof(pt.substr(amp + 1));
|
||||
pts.push_back({x, y});
|
||||
if (us == string::npos) break;
|
||||
start = us + 1;
|
||||
}
|
||||
|
||||
if (pts.size() != 4) return meta;
|
||||
|
||||
// 几何排序: TL=min(x+y), TR=max(x-y), BR=max(x+y), BL=min(x-y)
|
||||
float sums[4], diffs[4];
|
||||
for (int i = 0; i < 4; i++) {
|
||||
sums[i] = pts[i].x + pts[i].y;
|
||||
diffs[i] = pts[i].x - pts[i].y;
|
||||
}
|
||||
int tl = (int)(min_element(sums, sums+4) - sums);
|
||||
int tr = (int)(max_element(diffs, diffs+4) - diffs);
|
||||
int br = (int)(max_element(sums, sums+4) - sums);
|
||||
int bl = (int)(min_element(diffs, diffs+4) - diffs);
|
||||
|
||||
meta.corners[0] = pts[tl];
|
||||
meta.corners[1] = pts[tr];
|
||||
meta.corners[2] = pts[br];
|
||||
meta.corners[3] = pts[bl];
|
||||
meta.valid = true;
|
||||
} catch (...) {
|
||||
meta.valid = false;
|
||||
}
|
||||
|
||||
return meta;
|
||||
}
|
||||
|
||||
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
// 合成车牌池
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
|
||||
void PlatePool::load(const string& dir) {
|
||||
// Fallback: load from directory
|
||||
for (auto& entry : fs::directory_iterator(dir)) {
|
||||
if (!entry.is_directory()) continue;
|
||||
string bg_color = entry.path().filename().string();
|
||||
vector<PlateEntry>& vec = pools_[bg_color];
|
||||
|
||||
for (auto& f : fs::directory_iterator(entry.path())) {
|
||||
if (!f.is_regular_file()) continue;
|
||||
string ext = f.path().extension().string();
|
||||
if (ext != ".jpg" && ext != ".png" && ext != ".PNG") continue;
|
||||
|
||||
Mat img = imread_unicode(f.path().string());
|
||||
if (img.empty()) continue;
|
||||
|
||||
PlateEntry pe;
|
||||
pe.image = img;
|
||||
vec.push_back(pe);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void PlatePool::load_bin(const string& bin_path) {
|
||||
ifstream f(bin_path, ios::binary);
|
||||
if (!f) { cerr << "[WARN] pool.bin not found" << endl; return; }
|
||||
|
||||
char magic[4];
|
||||
f.read(magic, 4);
|
||||
if (memcmp(magic, "IPPL", 4) != 0) { cerr << "[WARN] bad magic" << endl; return; }
|
||||
|
||||
uint32_t num_types;
|
||||
f.read((char*)&num_types, 4);
|
||||
|
||||
for (uint32_t ti = 0; ti < num_types; ti++) {
|
||||
uint32_t name_len;
|
||||
f.read((char*)&name_len, 4);
|
||||
string plate_type(name_len, 0);
|
||||
f.read(&plate_type[0], name_len);
|
||||
|
||||
uint32_t count;
|
||||
f.read((char*)&count, 4);
|
||||
|
||||
vector<PlateEntry>& vec = pools_[plate_type];
|
||||
vec.reserve(count);
|
||||
|
||||
for (uint32_t pi = 0; pi < count; pi++) {
|
||||
uint32_t w, h, data_len;
|
||||
f.read((char*)&w, 4);
|
||||
f.read((char*)&h, 4);
|
||||
f.read((char*)&data_len, 4);
|
||||
|
||||
vector<uint8_t> buf(data_len);
|
||||
f.read((char*)buf.data(), data_len);
|
||||
|
||||
PlateEntry pe;
|
||||
pe.image = Mat(h, w, CV_8UC3, buf.data()).clone();
|
||||
|
||||
// Read plate number metadata
|
||||
uint32_t num_len;
|
||||
f.read((char*)&num_len, 4);
|
||||
pe.plate_number.resize(num_len);
|
||||
f.read(&pe.plate_number[0], num_len);
|
||||
|
||||
vec.push_back(pe);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const PlateEntry& PlatePool::random(const string& plate_type, mt19937& rng) {
|
||||
auto& vec = pools_[plate_type];
|
||||
if (vec.empty()) {
|
||||
for (auto& [k, v] : pools_) {
|
||||
if (!v.empty()) {
|
||||
uniform_int_distribution<size_t> dist(0, v.size()-1);
|
||||
return v[dist(rng)];
|
||||
}
|
||||
}
|
||||
throw runtime_error("PlatePool is empty");
|
||||
}
|
||||
uniform_int_distribution<size_t> dist(0, vec.size()-1);
|
||||
return vec[dist(rng)];
|
||||
}
|
||||
|
||||
size_t PlatePool::count(const string& plate_type) const {
|
||||
auto it = pools_.find(plate_type);
|
||||
return (it != pools_.end()) ? it->second.size() : 0;
|
||||
}
|
||||
|
||||
|
||||
// Helper: plate_type → bg_color for color matching
|
||||
static string bg_color_of(const string& plate_type) {
|
||||
if (plate_type == "police") return "white";
|
||||
if (plate_type == "army") return "white_army";
|
||||
if (plate_type == "gangao") return "black";
|
||||
if (plate_type == "shiguan_old" || plate_type == "shiguan_new") return "black_shi";
|
||||
if (plate_type == "lingguan_old" || plate_type == "lingguan_new") return "black";
|
||||
if (plate_type == "coach") return "yellow";
|
||||
if (plate_type == "new_energy") return "green_truck";
|
||||
return "white";
|
||||
}
|
||||
|
||||
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
// Unicode 文件读写
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
|
||||
Mat imread_unicode(const string& path) {
|
||||
ifstream f(fs::path(path), ios::binary);
|
||||
if (!f) return Mat();
|
||||
vector<uint8_t> data((istreambuf_iterator<char>(f)), istreambuf_iterator<char>());
|
||||
return imdecode(data, IMREAD_COLOR);
|
||||
}
|
||||
|
||||
bool imwrite_unicode(const string& path, const Mat& img, const vector<int>& params) {
|
||||
vector<uint8_t> buf;
|
||||
fs::path p(path);
|
||||
string ext = p.extension().string();
|
||||
if (!imencode(ext, img, buf, params)) return false;
|
||||
// Use fs::path directly for proper Unicode support on Windows
|
||||
ofstream f(p, ios::binary);
|
||||
if (!f) return false;
|
||||
f.write(reinterpret_cast<const char*>(buf.data()), buf.size());
|
||||
return f.good();
|
||||
}
|
||||
|
||||
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
// 透视提取车牌 ROI
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
|
||||
Mat extract_plate_roi(const Mat& img, const Point2f corners[4],
|
||||
int plate_w, int plate_h) {
|
||||
Point2f dst[4] = {
|
||||
{0, 0},
|
||||
{(float)(plate_w - 1), 0},
|
||||
{(float)(plate_w - 1), (float)(plate_h - 1)},
|
||||
{0, (float)(plate_h - 1)}
|
||||
};
|
||||
Mat M = getPerspectiveTransform(corners, dst);
|
||||
Mat roi;
|
||||
warpPerspective(img, roi, M, Size(plate_w, plate_h),
|
||||
INTER_LINEAR, BORDER_REPLICATE);
|
||||
return roi;
|
||||
}
|
||||
|
||||
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
// 色温匹配
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
|
||||
Mat match_plate_color(const Mat& synthetic, const Mat& orig_roi,
|
||||
const string& bg_color) {
|
||||
Mat syn, orig;
|
||||
synthetic.convertTo(syn, CV_32F);
|
||||
orig_roi.convertTo(orig, CV_32F);
|
||||
|
||||
// 尺寸对齐
|
||||
if (syn.size() != orig.size())
|
||||
resize(syn, syn, orig.size(), 0, 0, INTER_LINEAR);
|
||||
|
||||
if (bg_color == "black" || bg_color == "black_shi") {
|
||||
Mat orig_gray;
|
||||
cvtColor(orig_roi, orig_gray, COLOR_BGR2GRAY);
|
||||
double thresh_val;
|
||||
{
|
||||
Mat sorted;
|
||||
orig_gray.reshape(1, 1).copyTo(sorted);
|
||||
cv::sort(sorted, sorted, SORT_EVERY_ROW + SORT_ASCENDING);
|
||||
int idx = (int)(sorted.cols * 0.90);
|
||||
thresh_val = sorted.at<uint8_t>(0, idx);
|
||||
}
|
||||
Mat white_mask = orig_gray > thresh_val;
|
||||
if (countNonZero(white_mask) < 50) return synthetic;
|
||||
|
||||
Scalar white_ref = mean(orig, white_mask);
|
||||
float ref_mean = (float)(white_ref[0] + white_ref[1] + white_ref[2]) / 3.0f;
|
||||
if (ref_mean < 10) return synthetic;
|
||||
|
||||
Scalar color_cast(white_ref[0]/ref_mean, white_ref[1]/ref_mean, white_ref[2]/ref_mean);
|
||||
Scalar amplified(1.0f + (color_cast[0]-1.0f), 1.0f+(color_cast[1]-1.0f), 1.0f+(color_cast[2]-1.0f));
|
||||
amplified[0] = std::clamp((double)amplified[0], 0.85, 1.15);
|
||||
amplified[1] = std::clamp((double)amplified[1], 0.85, 1.15);
|
||||
amplified[2] = std::clamp((double)amplified[2], 0.85, 1.15);
|
||||
|
||||
Mat syn_gray;
|
||||
cvtColor(synthetic, syn_gray, COLOR_BGR2GRAY);
|
||||
Mat alpha;
|
||||
syn_gray.convertTo(alpha, CV_32F, 1.0/128.0);
|
||||
alpha = max(0.0f, min(1.0f, alpha));
|
||||
|
||||
Mat result = syn.clone();
|
||||
for (int c = 0; c < 3; c++) {
|
||||
Mat ch_syn, ch_res;
|
||||
extractChannel(syn, ch_syn, c);
|
||||
ch_res = ch_syn.mul(1.0f + ((float)amplified[c] - 1.0f) * alpha);
|
||||
ch_res.setTo(0, ch_res < 0);
|
||||
ch_res.setTo(255, ch_res > 255);
|
||||
insertChannel(ch_res, result, c);
|
||||
}
|
||||
result.convertTo(result, CV_8U);
|
||||
return result;
|
||||
}
|
||||
|
||||
// 亮底牌：提取白参考
|
||||
Mat orig_gray;
|
||||
cvtColor(orig_roi, orig_gray, COLOR_BGR2GRAY);
|
||||
double thresh_val;
|
||||
{
|
||||
Mat sorted;
|
||||
orig_gray.reshape(1, 1).copyTo(sorted);
|
||||
cv::sort(sorted, sorted, SORT_EVERY_ROW + SORT_ASCENDING);
|
||||
int idx = (int)(sorted.cols * 0.90);
|
||||
thresh_val = sorted.at<uint8_t>(0, idx);
|
||||
}
|
||||
Mat white_mask = orig_gray > thresh_val;
|
||||
if (countNonZero(white_mask) < 50) return synthetic;
|
||||
|
||||
Scalar white_ref = mean(orig, white_mask);
|
||||
float ref_mean = (float)(white_ref[0] + white_ref[1] + white_ref[2]) / 3.0f;
|
||||
if (ref_mean < 10) return synthetic;
|
||||
|
||||
if (bg_color == "yellow") {
|
||||
Scalar gain = white_ref / 255.0;
|
||||
gain[0] = std::clamp((float)gain[0], 0.6f, 1.3f);
|
||||
gain[1] = std::clamp((float)gain[1], 0.6f, 1.3f);
|
||||
gain[2] = std::clamp((float)gain[2], 0.6f, 1.3f);
|
||||
Mat res;
|
||||
multiply(syn, gain, res);
|
||||
res.setTo(0, res < 0);
|
||||
res.setTo(255, res > 255);
|
||||
res.convertTo(res, CV_8U);
|
||||
return res;
|
||||
}
|
||||
|
||||
// 白/绿牌
|
||||
Scalar color_cast(white_ref[0]/ref_mean, white_ref[1]/ref_mean, white_ref[2]/ref_mean);
|
||||
float strength = (bg_color == "green_truck") ? 1.5f : 1.0f;
|
||||
Scalar amplified(1.0f + (color_cast[0]-1.0f)*strength, 1.0f+(color_cast[1]-1.0f)*strength, 1.0f+(color_cast[2]-1.0f)*strength);
|
||||
amplified[0] = std::clamp((float)amplified[0], 0.85f, 1.15f);
|
||||
amplified[1] = std::clamp((float)amplified[1], 0.85f, 1.15f);
|
||||
amplified[2] = std::clamp((float)amplified[2], 0.85f, 1.15f);
|
||||
|
||||
Mat syn_gray;
|
||||
cvtColor(synthetic, syn_gray, COLOR_BGR2GRAY);
|
||||
Mat alpha;
|
||||
syn_gray.convertTo(alpha, CV_32F, 1.0/128.0);
|
||||
alpha = max(0.0f, min(1.0f, alpha));
|
||||
|
||||
Mat result = syn.clone();
|
||||
for (int c = 0; c < 3; c++) {
|
||||
Mat ch_syn, ch_res;
|
||||
extractChannel(syn, ch_syn, c);
|
||||
ch_res = ch_syn.mul(1.0f + (amplified[c] - 1.0f) * alpha);
|
||||
ch_res.setTo(0, ch_res < 0);
|
||||
ch_res.setTo(255, ch_res > 255);
|
||||
insertChannel(ch_res, result, c);
|
||||
}
|
||||
result.convertTo(result, CV_8U);
|
||||
return result;
|
||||
}
|
||||
|
||||
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
// 底板色校正
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
|
||||
Mat fix_bg_color(const Mat& plate, const string& bg_color,
|
||||
int h_shift, int s_scale, int v_scale) {
|
||||
if (h_shift == 0 && s_scale == 100 && v_scale == 100)
|
||||
return plate;
|
||||
|
||||
int h_lo = 0, h_hi = 180;
|
||||
if (bg_color == "yellow") { h_lo = 15; h_hi = 50; }
|
||||
else if (bg_color == "green_truck") { h_lo = 40; h_hi = 85; }
|
||||
|
||||
Mat hsv;
|
||||
cvtColor(plate, hsv, COLOR_BGR2HSV);
|
||||
hsv.convertTo(hsv, CV_32F);
|
||||
|
||||
vector<Mat> ch(3);
|
||||
split(hsv, ch);
|
||||
Mat &H = ch[0], &S = ch[1], &V = ch[2];
|
||||
|
||||
Mat bg_mask = (H > h_lo) & (H < h_hi) & (S > 80) & (V > 100);
|
||||
if (countNonZero(bg_mask) < 200) return plate;
|
||||
|
||||
// 距离变换羽化
|
||||
Mat bg_u8, dist;
|
||||
bg_mask.convertTo(bg_u8, CV_8U, 255);
|
||||
distanceTransform(bg_u8, dist, DIST_L2, DIST_MASK_PRECISE);
|
||||
Mat soft;
|
||||
dist.convertTo(soft, CV_32F, 1.0/4.0);
|
||||
soft = max(0.0f, min(1.0f, soft));
|
||||
|
||||
H = H + h_shift * soft;
|
||||
S = S.mul(1.0f + (s_scale / 100.0f - 1.0f) * soft);
|
||||
V = V.mul(1.0f + (v_scale / 100.0f - 1.0f) * soft);
|
||||
|
||||
H.setTo(0, H < 0); H.setTo(180, H > 180);
|
||||
S.setTo(0, S < 0); S.setTo(255, S > 255);
|
||||
V.setTo(0, V < 0); V.setTo(255, V > 255);
|
||||
|
||||
Mat result;
|
||||
merge(ch, hsv);
|
||||
hsv.convertTo(hsv, CV_8U);
|
||||
cvtColor(hsv, result, COLOR_HSV2BGR);
|
||||
return result;
|
||||
}
|
||||
|
||||
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
// 透视回贴
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
|
||||
pair<Mat, Mat> warp_plate_back(const Mat& plate, const Point2f dst_corners[4],
|
||||
Size output_size) {
|
||||
int h = plate.rows, w = plate.cols;
|
||||
Point2f src[4] = {
|
||||
{0, 0}, {(float)(w-1), 0}, {(float)(w-1), (float)(h-1)}, {0, (float)(h-1)}
|
||||
};
|
||||
|
||||
Mat M = getPerspectiveTransform(src, dst_corners);
|
||||
Mat warped, mask_warped;
|
||||
warpPerspective(plate, warped, M, output_size, INTER_LINEAR, BORDER_CONSTANT, Scalar(0,0,0));
|
||||
|
||||
Mat mask(h, w, CV_8UC1, Scalar(255));
|
||||
warpPerspective(mask, mask_warped, M, output_size, INTER_LINEAR, BORDER_CONSTANT, Scalar(0));
|
||||
|
||||
return {warped, mask_warped};
|
||||
}
|
||||
|
||||
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
// 羽化贴合
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
|
||||
Mat feather_blend(const Mat& original, const Mat& warped,
|
||||
const Mat& mask, int edge_width) {
|
||||
Mat mask_u8;
|
||||
mask.convertTo(mask_u8, CV_8U);
|
||||
|
||||
if (countNonZero(mask_u8) < 100) {
|
||||
Mat result = original.clone();
|
||||
warped.copyTo(result, mask_u8);
|
||||
return result;
|
||||
}
|
||||
|
||||
// Fast blend: erode mask slightly for edge transition, no expensive distanceTransform
|
||||
Mat eroded;
|
||||
int ksize = max(3, edge_width) | 1; // ensure odd
|
||||
Mat kernel = getStructuringElement(MORPH_ELLIPSE, Size(ksize, ksize));
|
||||
erode(mask_u8, eroded, kernel);
|
||||
|
||||
Mat alpha;
|
||||
eroded.convertTo(alpha, CV_32F, 1.0/255.0);
|
||||
// Soften edges with a small blur
|
||||
GaussianBlur(alpha, alpha, Size(ksize, ksize), ksize/3.0);
|
||||
|
||||
Mat result(original.size(), CV_32FC3);
|
||||
Mat warped_f, orig_f;
|
||||
warped.convertTo(warped_f, CV_32F);
|
||||
original.convertTo(orig_f, CV_32F);
|
||||
|
||||
for (int c = 0; c < 3; c++) {
|
||||
Mat wc, oc;
|
||||
extractChannel(warped_f, wc, c);
|
||||
extractChannel(orig_f, oc, c);
|
||||
Mat rc = wc.mul(alpha) + oc.mul(1.0f - alpha);
|
||||
rc.setTo(0, rc < 0);
|
||||
rc.setTo(255, rc > 255);
|
||||
insertChannel(rc, result, c);
|
||||
}
|
||||
result.convertTo(result, CV_8U);
|
||||
return result;
|
||||
}
|
||||
|
||||
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
// 噪声匹配
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
|
||||
Mat match_noise(const Mat& blended, const Mat& original, const Mat& mask) {
|
||||
Mat mask_u8;
|
||||
mask.convertTo(mask_u8, CV_8U);
|
||||
if (countNonZero(mask_u8) < 100) return blended;
|
||||
|
||||
// Crop plate region for noise estimation (like Python's plate_crop)
|
||||
vector<Point> pts;
|
||||
findNonZero(mask_u8, pts);
|
||||
if (pts.size() < 200) return blended;
|
||||
|
||||
int bw = blended.cols, bh = blended.rows;
|
||||
int xmin = bw, xmax = 0, ymin = bh, ymax = 0;
|
||||
for (auto& p : pts) {
|
||||
if (p.x < xmin) xmin = p.x;
|
||||
if (p.x > xmax) xmax = p.x;
|
||||
if (p.y < ymin) ymin = p.y;
|
||||
if (p.y > ymax) ymax = p.y;
|
||||
}
|
||||
int x1 = max(0, xmin - 5), x2 = min(bw, xmax + 5);
|
||||
int y1 = max(0, ymin - 5), y2 = min(bh, ymax + 5);
|
||||
Mat plate_crop = original(Rect(x1, y1, x2-x1, y2-y1));
|
||||
|
||||
// 1) Noise sigma from original plate crop
|
||||
Mat crop_blurred, crop_residual;
|
||||
GaussianBlur(plate_crop, crop_blurred, Size(5,5), 1.0);
|
||||
Mat crop_f32, blur_f32;
|
||||
plate_crop.convertTo(crop_f32, CV_32F);
|
||||
crop_blurred.convertTo(blur_f32, CV_32F);
|
||||
crop_residual = crop_f32 - blur_f32;
|
||||
Scalar stddev;
|
||||
meanStdDev(crop_residual, Scalar(), stddev);
|
||||
float noise_sigma = (float)max(stddev[0], max(stddev[1], max(stddev[2], 1.5)));
|
||||
|
||||
// 2) Blur sigma from original plate crop
|
||||
Mat crop_gray;
|
||||
cvtColor(plate_crop, crop_gray, COLOR_BGR2GRAY);
|
||||
Mat lap;
|
||||
Laplacian(crop_gray, lap, CV_64F);
|
||||
Scalar lap_m, lap_s;
|
||||
meanStdDev(lap, lap_m, lap_s);
|
||||
float lap_var = (float)(lap_s[0] * lap_s[0]);
|
||||
float blur_sigma = 0;
|
||||
if (lap_var < 50) blur_sigma = 2.0f;
|
||||
else if (lap_var < 200) blur_sigma = 1.2f;
|
||||
else if (lap_var < 500) blur_sigma = 0.5f;
|
||||
|
||||
// Apply blur to plate area
|
||||
Mat result = blended.clone();
|
||||
if (blur_sigma > 0.1f) {
|
||||
Mat blurred_full;
|
||||
GaussianBlur(result, blurred_full, Size(0,0), blur_sigma);
|
||||
blurred_full.copyTo(result, mask_u8);
|
||||
}
|
||||
|
||||
// Add noise to plate area only
|
||||
Mat result_f;
|
||||
result.convertTo(result_f, CV_32F);
|
||||
Mat noise(result.size(), CV_32FC3);
|
||||
randn(noise, Scalar(0,0,0), Scalar(noise_sigma, noise_sigma, noise_sigma));
|
||||
Mat noisy = result_f + noise;
|
||||
|
||||
// Alpha blend: noise in plate area, clean outside
|
||||
Mat alpha;
|
||||
mask_u8.convertTo(alpha, CV_32F, 1.0/255.0);
|
||||
Mat alpha_3;
|
||||
merge(vector<Mat>{alpha, alpha, alpha}, alpha_3);
|
||||
Mat one_minus_a;
|
||||
subtract(Scalar::all(1.0), alpha_3, one_minus_a);
|
||||
noisy = noisy.mul(alpha_3) + result_f.mul(one_minus_a);
|
||||
|
||||
// Clip AFTER blend (matching Python np.clip order)
|
||||
noisy.setTo(0, noisy < 0);
|
||||
noisy.setTo(255, noisy > 255);
|
||||
noisy.convertTo(result, CV_8U);
|
||||
return result;
|
||||
}
|
||||
|
||||
Mat degrade_plate(const Mat& plate) {
|
||||
Mat result = plate.clone();
|
||||
|
||||
// JPEG 极限压缩
|
||||
vector<uint8_t> buf;
|
||||
imencode(".jpg", result, buf, {IMWRITE_JPEG_QUALITY, 2});
|
||||
result = imdecode(buf, IMREAD_COLOR);
|
||||
|
||||
// 降采样升采样
|
||||
int w = result.cols, h = result.rows;
|
||||
Mat small;
|
||||
resize(result, small, Size(w/4, h/4), 0, 0, INTER_LINEAR);
|
||||
resize(small, result, Size(w, h), 0, 0, INTER_NEAREST);
|
||||
|
||||
// 模糊
|
||||
GaussianBlur(result, result, Size(5,5), 1.5);
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
// 单张处理
|
||||
// ═══════════════════════════════════════════════════════════
|
||||
|
||||
bool process_single(const string& img_path, const SwapConfig& cfg,
|
||||
PlatePool& pool, mt19937& rng, const CCPDMeta* pre_meta) {
|
||||
// 1. 解析文件名 (use pre-parsed if available)
|
||||
fs::path fp(img_path);
|
||||
CCPDMeta meta;
|
||||
if (pre_meta && pre_meta->valid) {
|
||||
meta = *pre_meta;
|
||||
} else {
|
||||
meta = parse_ccpd_filename(fp.filename().string());
|
||||
if (!meta.valid) return false;
|
||||
}
|
||||
|
||||
// 2. 读原图
|
||||
Mat img_orig = imread_unicode(img_path);
|
||||
if (img_orig.empty()) return false;
|
||||
|
||||
int h_img = img_orig.rows, w_img = img_orig.cols;
|
||||
|
||||
// 校验角点
|
||||
for (int i = 0; i < 4; i++) {
|
||||
auto& c = meta.corners[i];
|
||||
if (c.x < 0 || c.y < 0 || c.x > w_img*1.5 || c.y > h_img*1.5)
|
||||
return false;
|
||||
}
|
||||
|
||||
// 3. Pick from pool
|
||||
string bg_color = bg_color_of(cfg.plate_type);
|
||||
const PlateEntry& entry = pool.random(cfg.plate_type, rng);
|
||||
Mat synthetic = entry.image.clone();
|
||||
int plate_w = synthetic.cols, plate_h = synthetic.rows;
|
||||
string plate_number = entry.plate_number;
|
||||
|
||||
// 4. Extract original ROI
|
||||
Mat orig_roi = extract_plate_roi(img_orig, meta.corners, plate_w, plate_h);
|
||||
|
||||
|
||||
// 5. Fix bg color
|
||||
synthetic = fix_bg_color(synthetic, bg_color,
|
||||
cfg.h_shift, cfg.s_scale, cfg.v_scale);
|
||||
|
||||
// 6. Color match
|
||||
Mat matched = match_plate_color(synthetic, orig_roi, bg_color);
|
||||
|
||||
// 7. Brightness match
|
||||
Mat orig_gray, matched_gray;
|
||||
cvtColor(orig_roi, orig_gray, COLOR_BGR2GRAY);
|
||||
cvtColor(matched, matched_gray, COLOR_BGR2GRAY);
|
||||
|
||||
auto p95 = [](const Mat& m) -> float {
|
||||
Mat sorted;
|
||||
m.reshape(1,1).copyTo(sorted);
|
||||
cv::sort(sorted, sorted, SORT_EVERY_ROW + SORT_ASCENDING);
|
||||
return sorted.at<uint8_t>(0, (int)(sorted.cols * 0.95));
|
||||
};
|
||||
float target_w = p95(orig_gray);
|
||||
float current_w = p95(matched_gray);
|
||||
|
||||
if (current_w > 10 && abs(target_w - current_w) > 10) {
|
||||
float l_ratio = target_w / current_w;
|
||||
float min_c = 0.5f, max_c = 1.1f;
|
||||
if (bg_color == "green_truck" || bg_color == "yellow") { min_c = 0.4f; max_c = 1.0f; }
|
||||
else if (bg_color == "black" || bg_color == "black_shi") { min_c = 0.3f; max_c = 1.2f; }
|
||||
l_ratio = std::clamp(l_ratio, min_c, max_c);
|
||||
|
||||
Mat alpha;
|
||||
matched_gray.convertTo(alpha, CV_32F, 1.0/128.0);
|
||||
alpha = max(0.3f, min(1.0f, alpha));
|
||||
|
||||
Mat matched_f;
|
||||
matched.convertTo(matched_f, CV_32F);
|
||||
for (int c = 0; c < 3; c++) {
|
||||
Mat ch;
|
||||
extractChannel(matched_f, ch, c);
|
||||
ch = ch.mul(1.0f + (l_ratio - 1.0f) * alpha);
|
||||
ch.setTo(0, ch < 0); ch.setTo(255, ch > 255);
|
||||
insertChannel(ch, matched_f, c);
|
||||
}
|
||||
matched_f.convertTo(matched, CV_8U);
|
||||
}
|
||||
|
||||
// 8. Degradation
|
||||
if (!cfg.fast_mode)
|
||||
matched = degrade_plate(matched);
|
||||
|
||||
// 9-11. Warp + blend + noise
|
||||
auto [warped, mask] = warp_plate_back(matched, meta.corners, Size(w_img, h_img));
|
||||
Mat blended = feather_blend(img_orig, warped, mask);
|
||||
Mat final = match_noise(blended, img_orig, mask);
|
||||
|
||||
// 12. Save with pinyin tags
|
||||
string stem = fp.stem().string();
|
||||
// Convert Chinese plate markers to ASCII tags
|
||||
string tag = plate_number;
|
||||
// Manual pinyin mapping for common plate characters (UTF-8 byte sequences)
|
||||
// 港=E6B8AF 澳=E6BEB3 警=E8ADA6 使=E4BDBF 领=E9A286 学=E5ADA6
|
||||
struct { const char* utf8; const char* py; } map[] = {
|
||||
{"\xE6\xB8\xAF", "GANG"}, {"\xE6\xBE\xB3", "AO"},
|
||||
{"\xE8\xAD\xA6", "JING"}, {"\xE4\xBD\xBF", "SHI"},
|
||||
{"\xE9\xA2\x86", "LING"}, {"\xE5\xAD\xA6", "XUE"},
|
||||
};
|
||||
for (auto& m : map) {
|
||||
size_t pos = 0;
|
||||
while ((pos = tag.find(m.utf8, pos)) != string::npos) {
|
||||
tag.replace(pos, 3, m.py);
|
||||
pos += strlen(m.py);
|
||||
}
|
||||
}
|
||||
// Remove any remaining non-ASCII
|
||||
string clean;
|
||||
for (char c : tag) {
|
||||
if ((unsigned char)c < 128) clean += c;
|
||||
}
|
||||
string out_name = stem + "__" + cfg.plate_type + "__" + clean + ".jpg";
|
||||
string ccpd_subdir = fp.parent_path().filename().string();
|
||||
fs::path out_dir = fs::path(cfg.output_dir) / ccpd_subdir / cfg.plate_type;
|
||||
fs::create_directories(out_dir);
|
||||
|
||||
return imwrite_unicode((out_dir / out_name).string(), final, {IMWRITE_JPEG_QUALITY, 75});
|
||||
}
|
||||
|
||||
} // namespace ice
|
||||
@@ -0,0 +1,80 @@
|
||||
#pragma once
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <random>
|
||||
#include <unordered_map>
|
||||
|
||||
namespace ice {
|
||||
|
||||
// ── CCPD 角点解析 ──
|
||||
struct CCPDMeta {
|
||||
cv::Point2f corners[4]; // TL, TR, BR, BL
|
||||
bool valid = false;
|
||||
};
|
||||
CCPDMeta parse_ccpd_filename(const std::string& filename);
|
||||
|
||||
// ── 合成车牌池 ──
|
||||
struct PlateEntry {
|
||||
cv::Mat image; // BGR
|
||||
std::string plate_number; // 车牌号 (e.g. "粤Z1234港")
|
||||
};
|
||||
class PlatePool {
|
||||
public:
|
||||
void load(const std::string& dir);
|
||||
void load_bin(const std::string& bin_path);
|
||||
const PlateEntry& random(const std::string& plate_type, std::mt19937& rng);
|
||||
size_t count(const std::string& plate_type) const;
|
||||
private:
|
||||
std::unordered_map<std::string, std::vector<PlateEntry>> pools_;
|
||||
};
|
||||
|
||||
// ── 管线核心 ──
|
||||
struct SwapConfig {
|
||||
std::string plate_type;
|
||||
std::string output_dir;
|
||||
int h_shift = 0, s_scale = 100, v_scale = 100;
|
||||
bool fast_mode = false;
|
||||
};
|
||||
|
||||
// 单张处理 (pre-parsed meta 避免重复解析)
|
||||
bool process_single(const std::string& img_path,
|
||||
const SwapConfig& cfg,
|
||||
PlatePool& pool,
|
||||
std::mt19937& rng,
|
||||
const CCPDMeta* pre_meta = nullptr);
|
||||
|
||||
// 透视提取车牌ROI
|
||||
cv::Mat extract_plate_roi(const cv::Mat& img, const cv::Point2f corners[4],
|
||||
int plate_w, int plate_h);
|
||||
|
||||
// 色温匹配
|
||||
cv::Mat match_plate_color(const cv::Mat& synthetic, const cv::Mat& orig_roi,
|
||||
const std::string& bg_color);
|
||||
|
||||
// 底板色校正
|
||||
cv::Mat fix_bg_color(const cv::Mat& plate, const std::string& bg_color,
|
||||
int h_shift, int s_scale, int v_scale);
|
||||
|
||||
// 透视回贴: 返回 (warped, mask)
|
||||
std::pair<cv::Mat, cv::Mat> warp_plate_back(const cv::Mat& plate,
|
||||
const cv::Point2f dst_corners[4],
|
||||
cv::Size output_size);
|
||||
|
||||
// 羽化贴合
|
||||
cv::Mat feather_blend(const cv::Mat& original, const cv::Mat& warped,
|
||||
const cv::Mat& mask, int edge_width = 8);
|
||||
|
||||
// 噪声匹配
|
||||
cv::Mat match_noise(const cv::Mat& blended, const cv::Mat& original,
|
||||
const cv::Mat& mask);
|
||||
|
||||
// 退化模拟
|
||||
cv::Mat degrade_plate(const cv::Mat& plate);
|
||||
|
||||
// ── 工具 ──
|
||||
cv::Mat imread_unicode(const std::string& path);
|
||||
bool imwrite_unicode(const std::string& path, const cv::Mat& img,
|
||||
const std::vector<int>& params = {});
|
||||
|
||||
} // namespace ice
|
||||
@@ -0,0 +1,295 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
车牌底板色手动调色 GUI
|
||||
======================
|
||||
用 HSV 滑块实时调整合成车牌底板色，预览效果，保存参数供管线调用。
|
||||
|
||||
用法:
|
||||
python color_tuner.py
|
||||
|
||||
控件:
|
||||
H 偏移 — 色相偏移量 (-30 ~ +30)
|
||||
S 缩放%% — 饱和度缩放 (10% ~ 150%)
|
||||
V 缩放%% — 明度缩放 (10% ~ 150%)
|
||||
[1-9] — 切换车牌类型
|
||||
[R] — 复位
|
||||
[S] — 保存当前参数到 plate_color_profile.py
|
||||
[ESC] — 退出
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
# ── 路径 ──
|
||||
_SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
_PROJECT_DIR = _SCRIPT_DIR.parent.parent
|
||||
_GENERATOR_DIR = _PROJECT_DIR / 'chinese_license_plate_generator'
|
||||
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,
|
||||
)
|
||||
|
||||
# ── 车牌类型 ──
|
||||
PLATE_TYPES = [
|
||||
('yellow', '教练车', generate_plate_number_yellow_xue),
|
||||
('green_truck','新能源', lambda: generate_plate_number_blue(length=8)),
|
||||
('white', '警车', generate_plate_number_white),
|
||||
('white_army', '军车', generate_plate_number_army),
|
||||
('black', '港澳', generate_plate_number_black_gangao),
|
||||
('black_shi', '使馆旧', generate_plate_number_black_shi_old),
|
||||
('black_shi', '使馆新', generate_plate_number_black_shi_new),
|
||||
('black', '领馆旧', generate_plate_number_black_ling_old),
|
||||
('black', '领馆新', generate_plate_number_black_ling_new),
|
||||
]
|
||||
|
||||
# ── 参考图 ──
|
||||
REF_PATH = Path(r"C:\Users\Kirl\Downloads\R-C (3).jpg")
|
||||
|
||||
# ── 全局状态 ──
|
||||
generator = None
|
||||
current_type_idx = 0
|
||||
synthetic_plate = None
|
||||
ref_image = None
|
||||
|
||||
# 各类型的默认参数 (H偏移, S缩放%, V缩放%)
|
||||
params = {}
|
||||
|
||||
|
||||
def init_generator():
|
||||
global generator
|
||||
pm = str(_GENERATOR_DIR / 'plate_model')
|
||||
fm = str(_GENERATOR_DIR / 'font_model')
|
||||
generator = MultiPlateGenerator(pm, fm)
|
||||
|
||||
|
||||
def load_ref():
|
||||
global ref_image
|
||||
if REF_PATH.is_file():
|
||||
ref = cv2.imread(str(REF_PATH))
|
||||
if ref is not None:
|
||||
h, w = ref.shape[:2]
|
||||
ref = cv2.resize(ref, (int(w * 0.6), int(h * 0.6)))
|
||||
ref_image = ref
|
||||
|
||||
|
||||
def generate_plate(type_idx):
|
||||
global synthetic_plate
|
||||
bg_color, label, gen_func = PLATE_TYPES[type_idx]
|
||||
number = gen_func()
|
||||
plate_w = 480 if len(number) == 8 else 440
|
||||
plate = generator.generate_plate_special(number, bg_color, is_double=False, enhance=False)
|
||||
synthetic_plate = plate
|
||||
return plate, number, label
|
||||
|
||||
|
||||
def load_params(bg_color):
|
||||
"""从已保存的 profile 加载参数"""
|
||||
key = bg_color
|
||||
if key not in params:
|
||||
params[key] = {'h_shift': 0, 's_scale': 100, 'v_scale': 100}
|
||||
return params[key]
|
||||
|
||||
|
||||
def adjust_plate(plate_bgr, bg_color, h_shift, s_scale, v_scale):
|
||||
"""
|
||||
在 HSV 空间调整底板色。
|
||||
仅作用于底板区域（高S + 对应色相范围 + 高V），字符不受影响。
|
||||
"""
|
||||
if h_shift == 0 and s_scale == 100 and v_scale == 100:
|
||||
return plate_bgr.copy()
|
||||
|
||||
hue_ranges = {
|
||||
'yellow': (15, 50), 'green_truck': (40, 85),
|
||||
'white': (0, 180), 'white_army': (0, 180),
|
||||
'black': (0, 180), 'black_shi': (0, 180),
|
||||
}
|
||||
h_lo, h_hi = hue_ranges.get(bg_color, (0, 180))
|
||||
|
||||
hsv = cv2.cvtColor(plate_bgr, cv2.COLOR_BGR2HSV).astype(np.float32)
|
||||
h, s, v = hsv[:, :, 0], hsv[:, :, 1], hsv[:, :, 2]
|
||||
|
||||
# 底板区域 mask
|
||||
hue_in = (h > h_lo) & (h < h_hi)
|
||||
bg_mask = hue_in & (s > 80) & (v > 100)
|
||||
|
||||
if bg_mask.sum() < 200:
|
||||
return plate_bgr.copy()
|
||||
|
||||
# 羽化
|
||||
bg_u8 = bg_mask.astype(np.uint8) * 255
|
||||
dist = cv2.distanceTransform(bg_u8, cv2.DIST_L2, cv2.DIST_MASK_PRECISE)
|
||||
soft = np.clip(dist / 4.0, 0.0, 1.0)
|
||||
|
||||
h_new = h + h_shift * soft
|
||||
s_new = s * (1.0 + (s_scale / 100.0 - 1.0) * soft)
|
||||
v_new = v * (1.0 + (v_scale / 100.0 - 1.0) * soft)
|
||||
|
||||
hsv[:, :, 0] = np.clip(h_new, 0, 180)
|
||||
hsv[:, :, 1] = np.clip(s_new, 0, 255)
|
||||
hsv[:, :, 2] = np.clip(v_new, 0, 255)
|
||||
|
||||
return cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
|
||||
|
||||
|
||||
def refresh():
|
||||
global synthetic_plate
|
||||
_, number, label = generate_plate(current_type_idx)
|
||||
update_display()
|
||||
|
||||
|
||||
def update_display(*_args):
|
||||
bg_color, label, _ = PLATE_TYPES[current_type_idx]
|
||||
p = load_params(bg_color)
|
||||
|
||||
# trackbar 可能在创建过程中尚未就绪，容错读取
|
||||
try:
|
||||
h_shift = cv2.getTrackbarPos('H offset', 'Color Tuner') - 30
|
||||
s_scale = cv2.getTrackbarPos('S scale', 'Color Tuner')
|
||||
v_scale = cv2.getTrackbarPos('V scale', 'Color Tuner')
|
||||
except cv2.error:
|
||||
h_shift = p.get('h_shift', 0)
|
||||
s_scale = p.get('s_scale', 100)
|
||||
v_scale = p.get('v_scale', 100)
|
||||
|
||||
p['h_shift'] = h_shift
|
||||
p['s_scale'] = s_scale
|
||||
p['v_scale'] = v_scale
|
||||
|
||||
adjusted = adjust_plate(synthetic_plate, bg_color, h_shift, s_scale, v_scale)
|
||||
|
||||
# 并排显示：原始 | 调整后
|
||||
h, w = synthetic_plate.shape[:2]
|
||||
display = np.hstack([synthetic_plate, adjusted])
|
||||
|
||||
# 加标签
|
||||
cv2.putText(display, f"Original ({bg_color})", (10, 25),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.55, (200, 200, 200), 1)
|
||||
cv2.putText(display, f"Adjusted H:{h_shift:+d} S:{s_scale} V:{v_scale}",
|
||||
(w + 10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 255, 200), 1)
|
||||
|
||||
# 类型标签
|
||||
cv2.putText(display, f"[{current_type_idx+1}/9] {label}",
|
||||
(10, display.shape[0] - 10),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (180, 180, 180), 1)
|
||||
|
||||
# 如果参考图存在，放在下方
|
||||
if ref_image is not None:
|
||||
rw = display.shape[1]
|
||||
rh = ref_image.shape[0] * rw // ref_image.shape[1]
|
||||
ref_rs = cv2.resize(ref_image, (rw, rh))
|
||||
display = np.vstack([display, np.ones((4, rw, 3), dtype=np.uint8) * 30, ref_rs])
|
||||
cv2.putText(display, "Reference", (10, display.shape[0] - rh + 25),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)
|
||||
|
||||
cv2.imshow('Color Tuner', display)
|
||||
|
||||
|
||||
def save_params():
|
||||
"""保存当前所有类型的参数到 plate_color_profile.py"""
|
||||
profile_path = _SCRIPT_DIR / 'plate_color_profile.py'
|
||||
lines = ['# Auto-generated by color_tuner.py',
|
||||
'# Plate background color correction profiles',
|
||||
'# Format: (h_shift, s_scale_percent, v_scale_percent)',
|
||||
'',
|
||||
'PLATE_COLOR_PROFILE = {']
|
||||
for bg_color in ['yellow', 'green_truck', 'white', 'white_army', 'black', 'black_shi']:
|
||||
p = params.get(bg_color, {'h_shift': 0, 's_scale': 100, 'v_scale': 100})
|
||||
lines.append(f" '{bg_color}': ({p['h_shift']}, {p['s_scale']}, {p['v_scale']}),")
|
||||
lines.append('}')
|
||||
lines.append('')
|
||||
|
||||
profile_path.write_text('\n'.join(lines), encoding='utf-8')
|
||||
print(f"\n[SAVED] {profile_path}")
|
||||
print(f" Params: { {k: v for k, v in params.items()} }")
|
||||
|
||||
|
||||
def on_key(key):
|
||||
global current_type_idx, synthetic_plate
|
||||
if ord('1') <= key <= ord('9'):
|
||||
idx = key - ord('1')
|
||||
if idx < len(PLATE_TYPES):
|
||||
current_type_idx = idx
|
||||
# 切换类型时复位 trackbar 到该类型已保存的参数
|
||||
bg_color = PLATE_TYPES[current_type_idx][0]
|
||||
p = load_params(bg_color)
|
||||
cv2.setTrackbarPos('H offset', 'Color Tuner', p['h_shift'] + 30)
|
||||
cv2.setTrackbarPos('S scale', 'Color Tuner', p['s_scale'])
|
||||
cv2.setTrackbarPos('V scale', 'Color Tuner', p['v_scale'])
|
||||
# 生成新车牌
|
||||
synthetic_plate, _, _ = generate_plate(current_type_idx)
|
||||
update_display()
|
||||
elif key == ord('r') or key == ord('R'):
|
||||
bg_color = PLATE_TYPES[current_type_idx][0]
|
||||
params[bg_color] = {'h_shift': 0, 's_scale': 100, 'v_scale': 100}
|
||||
# Reset trackbars
|
||||
cv2.setTrackbarPos('H offset', 'Color Tuner', 30)
|
||||
cv2.setTrackbarPos('S scale', 'Color Tuner', 100)
|
||||
cv2.setTrackbarPos('V scale', 'Color Tuner', 100)
|
||||
update_display()
|
||||
elif key == ord('s') or key == ord('S'):
|
||||
save_params()
|
||||
elif key == ord('n') or key == ord('N'):
|
||||
refresh()
|
||||
|
||||
|
||||
def main():
|
||||
global ref_image
|
||||
init_generator()
|
||||
load_ref()
|
||||
|
||||
# 初始化第一张车牌
|
||||
_, number, label = generate_plate(current_type_idx)
|
||||
|
||||
# 先建窗口和滑动条，再刷新显示
|
||||
cv2.namedWindow('Color Tuner', cv2.WINDOW_NORMAL)
|
||||
cv2.resizeWindow('Color Tuner', 1000, 500)
|
||||
|
||||
bg_color = PLATE_TYPES[current_type_idx][0]
|
||||
p = load_params(bg_color)
|
||||
cv2.createTrackbar('H offset', 'Color Tuner', p['h_shift'] + 30, 60, update_display)
|
||||
cv2.createTrackbar('S scale', 'Color Tuner', p['s_scale'], 150, update_display)
|
||||
cv2.createTrackbar('V scale', 'Color Tuner', p['v_scale'], 150, update_display)
|
||||
|
||||
update_display()
|
||||
|
||||
print("=" * 60)
|
||||
print(" 车牌底板色手动调色 GUI")
|
||||
print("=" * 60)
|
||||
print(" 按键:")
|
||||
print(" 1-9 切换车牌类型")
|
||||
print(" R 复位当前类型参数")
|
||||
print(" N 生成新随机车牌")
|
||||
print(" S 保存参数到 plate_color_profile.py")
|
||||
print(" ESC 退出")
|
||||
print("=" * 60)
|
||||
print(f" 当前类型: {PLATE_TYPES[current_type_idx][1]}")
|
||||
print("=" * 60)
|
||||
|
||||
while True:
|
||||
key = cv2.waitKey(50) & 0xFF
|
||||
if key == 27: # ESC
|
||||
break
|
||||
if key != 255:
|
||||
on_key(key)
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
import os
|
||||
os.chdir(str(_PROJECT_DIR))
|
||||
main()
|
||||
@@ -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()
|
||||
@@ -0,0 +1,45 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
生成纯合成车牌（无场景背景）作为 CUT Domain A
|
||||
============================================
|
||||
独立于 CCPD 原图，直接调用车牌生成器渲染纯车牌图像。
|
||||
"""
|
||||
import sys
|
||||
import os
|
||||
import random
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
# 路径解析：定位 chinese_license_plate_generator
|
||||
_SCRIPT_DIR = Path(__file__).resolve().parent # .../CCPD_SPE_GEN/cv_pipeline/
|
||||
_PROJECT_DIR = _SCRIPT_DIR.parent.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
|
||||
|
||||
random.seed(42)
|
||||
np.random.seed(42)
|
||||
|
||||
g = MultiPlateGenerator(
|
||||
str(_GENERATOR_DIR / 'plate_model'),
|
||||
str(_GENERATOR_DIR / '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)')
|
||||
@@ -0,0 +1,12 @@
|
||||
# Auto-generated by color_tuner.py
|
||||
# Plate background color correction profiles
|
||||
# Format: (h_shift, s_scale_percent, v_scale_percent)
|
||||
|
||||
PLATE_COLOR_PROFILE = {
|
||||
'yellow': (-2, 150, 100),
|
||||
'green_truck': (0, 100, 100),
|
||||
'white': (0, 100, 100),
|
||||
'white_army': (0, 100, 100),
|
||||
'black': (0, 100, 100),
|
||||
'black_shi': (0, 100, 100),
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -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
|
||||
@@ -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)
|
||||
@@ -0,0 +1,6 @@
|
||||
torch>=1.12
|
||||
torchvision
|
||||
opencv-python>=4.5
|
||||
numpy
|
||||
tqdm
|
||||
tensorboard
|
||||
@@ -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()
|
||||
@@ -0,0 +1,43 @@
|
||||
#!/bin/bash
|
||||
# ============================================================
|
||||
# GAN 数据准备：合成车牌(A) + 真实车牌(B)
|
||||
# 用法: cd ICE_Dataset && bash CCPD_SPE_GEN/scripts/gan_data_prep.sh
|
||||
# ============================================================
|
||||
source /home/huawei/rknn_env/bin/activate # TODO: 替换为你的虚拟环境路径
|
||||
cd /mnt/d/ICEv2/ICE_Dataset # TODO: 替换为你的 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 CCPD_SPE_GEN/cv_pipeline/plate_swap.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 CCPD_SPE_GEN/cv_pipeline/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
|
||||
@@ -0,0 +1,11 @@
|
||||
#!/bin/bash
|
||||
# 一键预览全部 9 种特殊车牌类型 (每种 2 张，保存中间结果)
|
||||
# 用法: cd ICE_Dataset && bash CCPD_SPE_GEN/scripts/preview_all.sh
|
||||
source /home/huawei/rknn_env/bin/activate # TODO: 替换为你的虚拟环境
|
||||
cd /mnt/d/ICEv2/ICE_Dataset # TODO: 替换为你的 ICE_Dataset 路径
|
||||
for p in police army gangao shiguan_old shiguan_new lingguan_old lingguan_new coach new_energy; do
|
||||
echo "=== $p ==="
|
||||
python3 CCPD_SPE_GEN/cv_pipeline/plate_swap.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"
|
||||
@@ -0,0 +1,7 @@
|
||||
#!/bin/bash
|
||||
# 快速测试：生成 20 张警车特殊车牌
|
||||
# 用法: cd ICE_Dataset && bash CCPD_SPE_GEN/scripts/quick_test.sh
|
||||
source /home/huawei/rknn_env/bin/activate # TODO: 替换为你的虚拟环境
|
||||
cd /mnt/d/ICEv2/ICE_Dataset # TODO: 替换为你的 ICE_Dataset 路径
|
||||
python3 CCPD_SPE_GEN/cv_pipeline/plate_swap.py --ccpd-dir ./CCPD/CCPD2019/ccpd_base --plate-type police --num 20 --output ./output/preview/police --seed 42 2>&1 | tail -2
|
||||
echo DONE
|
||||
Reference in new issue
Block a user