Skip to content

Repository files navigation

PlateTruth

基于对比学习的套牌车 / 克隆车识别系统

Python PyTorch ONNX FastAPI License


PlateTruth 是一个深度学习系统,用于判断两张车辆图片是否属于同一辆车(即是否挂同一张车牌)。它使用 ResNet 特征提取器,通过对比学习(contrastive loss) 训练,将车辆图片映射到 512 维特征空间,利用余弦相似度衡量车辆身份的一致性。

应用截图

目录

工作原理

┌────────────┐    ┌──────────────────┐    ┌─────────────────┐
│  输入图片   │───▶│  PlateEncoder    │───▶│  L2 归一化      │
│  (224×224) │    │  (ResNet34)      │    │  512 维特征向量  │
└────────────┘    └──────────────────┘    └────────┬────────┘
                                                   │
                                            ┌──────▼──────┐
                                            │  余弦相似度  │
                                            └──────┬──────┘
                                                   │
                                        ┌──────────▼──────────┐
                                        │  相似度 ≥ 0.85?     │
                                        │  → 同一车辆          │
                                        │  相似度 < 0.85?     │
                                        │  → 不同车辆          │
                                        └─────────────────────┘

流水线概览

阶段 运行时 产出
训练 PyTorch + ResNet best_model.pth(约 247 MB)
导出 torch.onnx.export platetruth.onnx
推理 ONNX Runtime(CPU/GPU) 512 维特征向量
服务 FastAPI + Uvicorn REST API

核心设计决策:

  • 按车牌 ID 划分数据集,而非按图片划分——确保同一辆车的图片不会泄漏到多个集合中。
  • 动态正负样本对生成——PlateDataset 在每个 batch 内实时采样正负样本对,无需预计算。
  • 生产环境使用 ONNX——推理时无需 PyTorch 依赖,镜像体积更小,启动速度更快。

快速开始

环境要求

  • Python 3.12+
  • uv(包管理工具)

安装

git clone https://github.com/noimank/PlateTruth.git
cd PlateTruth

# 安装依赖
uv sync

端到端示例

# 1. 准备并划分数据
python prepare_data.py --raw-dir data/raw --output-dir data/processed

# 2. 训练模型
python train.py --epochs 50 --batch-size 8

# 3. 在测试集上评估
python eval.py --data-dir data/processed/test --threshold 0.85

# 4. 导出 ONNX 模型用于部署
python export_onnx.py --model-path models/best_model.pth --output models/export/platetruth.onnx

# 5. 比较两张图片
python inference_onnx.py models/export/platetruth.onnx img1.jpg img2.jpg --threshold 0.85

数据准备

将车辆图片放入 data/raw/ 目录,每辆车的图片放在以车牌号命名的文件夹中:

data/raw/
├── 沪A12345/
│   ├── 2024-01-15_08-30-00.jpg
│   ├── 2024-01-15_18-45-22.jpg
│   └── 2024-01-16_09-12-03.jpg
├── 沪B67890/
│   ├── 2024-01-15_10-00-00.jpg
│   └── 2024-01-15_16-30-00.jpg
└── ...

运行数据划分脚本:

python prepare_data.py --raw-dir data/raw --output-dir data/processed

按车牌 ID 进行 80/10/10 划分:

集合 比例 用途
train/ 80% 模型训练
val/ 10% 验证集(用于选择最佳模型)
test/ 10% 最终评估

注意: 划分以车牌 ID 文件夹为单位,而非单张图片。这保证了同一车辆的所有图片落在同一个集合中,杜绝数据泄漏,是评估结果可信的前提。

支持的图片格式:.jpg.jpeg.png.bmp

模型训练

# 从头训练(忽略已有 checkpoint)
python train.py --no-resume --epochs 50 --batch-size 8 --lr 1e-4

# 从最近的 checkpoint 恢复训练
python train.py

# 从指定 checkpoint 恢复训练
python train.py --resume-from models/checkpoints/epoch_10.pt

训练行为

  • 验证损失最低的模型自动保存至 models/best_model.pth
  • SAVE_INTERVAL 个 epoch(默认 2)保存一次 checkpoint 到 models/checkpoints/
  • 所有超参数可在 src/config.py 中统一调整。

损失函数

PlateTruth 使用带余弦相似度边距的对比损失(contrastive loss)

$$ \mathcal{L} = \begin{cases} (1 - \text{sim})^2 & \text{同一车牌} \\ \max(0, ,\text{sim} - (1 - m))^2 & \text{不同车牌} \end{cases} $$

其中 sim = cos(embedding₁, embedding₂),边距 m = 0.5

模型评估

python eval.py --data-dir data/processed/test --threshold 0.85

脚本会输出以下指标:

指标 说明
Accuracy(准确率) 整体预测正确率
Precision(精确率) 预测为正例中真正例的比例
Recall(召回率) 实际正例中被正确识别的比例
F1 Score 精确率与召回率的调和平均

推理预测

PyTorch 推理(需 best_model.pth

python inference_test.py img1.jpg img2.jpg

ONNX 推理(需 platetruth.onnx

python inference_onnx.py models/export/platetruth.onnx img1.jpg img2.jpg --threshold 0.85

Python API 调用

from src.onnx_inferencer import ONNXInferencer

model = ONNXInferencer("models/export/platetruth.onnx")
vec1 = model.extract(image_path_1)  # np.ndarray (512,)
vec2 = model.extract(image_path_2)
similarity = model.compare(image_path_1, image_path_2)
is_same = model.is_same_vehicle(image_path_1, image_path_2, threshold=0.85)

判定阈值

默认阈值为 0.85。余弦相似度取值范围为 [-1, 1]:

余弦相似度 ≥ 0.85  →  同一车辆
余弦相似度 < 0.85  →  不同车辆

可根据实际场景调整——提高阈值减少误报,降低阈值减少漏报。

模型导出

将 PyTorch 模型导出为 ONNX 格式用于生产部署:

python export_onnx.py --model-path models/best_model.pth --output models/export/platetruth.onnx
格式 依赖 适用场景
.pth PyTorch + TorchVision 训练、实验
.onnx 仅需 ONNX Runtime 生产部署、Docker 容器

API 服务

启动 FastAPI 服务:

uvicorn api.main:app --host 0.0.0.0 --port 8000 --reload

接口列表

方法 路径 说明
GET /api/v1/health 健康检查,返回模型元信息
POST /api/v1/extract 从单张图片提取 512 维特征向量
POST /api/v1/compare 比较两张图片,返回相似度
POST /api/v1/batch-compare 将一张参考图与多张目标图批量比对
GET / 内置 Web 测试界面

请求示例

健康检查:

curl http://localhost:8000/api/v1/health
{
  "status": "ok",
  "model": "platetruth",
  "backend": "onnx",
  "feature_dim": 512,
  "threshold": 0.85
}

比较两张图片:

curl -X POST http://localhost:8000/api/v1/compare \
  -F "image1=@car_a.jpg" \
  -F "image2=@car_b.jpg"
{
  "similarity": 0.923,
  "is_same_vehicle": true,
  "threshold": 0.85
}

提取特征向量:

curl -X POST http://localhost:8000/api/v1/extract \
  -F "image=@car.jpg"
{
  "features": [0.12, -0.34, ...],
  "shape": [512],
  "image_size": [224, 224]
}

交互式 API 文档

访问 http://localhost:8000/docs 查看自动生成的 Swagger UI。

Docker 部署

快速体验

已发布预构建镜像,可直接拉取运行,无需本地编译:

# 拉取镜像并启动
docker run -d -p 8000:8000 --name platetruth noimankdocker/platetruth:latest

# 验证服务是否就绪
curl http://localhost:8000/api/v1/health

# 打开浏览器访问 Web 测试界面
# http://localhost:8000

镜像已内置 ONNX 模型,开箱即用。

构建镜像

docker build -t platetruth .

启动容器

docker run -d -p 8000:8000 --name platetruth platetruth

子路径部署

当通过反向代理部署在子路径下时(如 https://example.com/platetruth/),设置 BASE_URL_PREFIX 环境变量:

docker run -d -p 8000:8000 \
  -e BASE_URL_PREFIX=/platetruth \
  --name platetruth platetruth

容器内部始终在 / 提供服务,反向代理负责剥离前缀后转发:

location /platetruth/ {
    proxy_pass http://127.0.0.1:8000/;
    proxy_set_header Host $host;
    proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
}

配置说明

所有配置集中在 src/config.py

# 模型
MODEL_NAME = "resnet34"          # resnet18 | resnet34 | resnet50
FEATURE_DIM = 512                # 输出特征维度
IMAGE_SIZE = (224, 224)          # 输入图片尺寸

# 训练
BATCH_SIZE = 8
NUM_EPOCHS = 50
LEARNING_RATE = 1e-4
WEIGHT_DECAY = 1e-4
NUM_WORKERS = 4
SAVE_INTERVAL = 2                # checkpoint 保存间隔(epoch)

# 对比损失
MARGIN = 0.5                     # 正负样本对边距

# 推理
SIMILARITY_THRESHOLD = 0.85      # 判定阈值
DEVICE = "cpu"                   # 设为 "cuda" 启用 GPU 推理

# 可复现性
SEED = 42

项目结构

platetruth/
├── src/
│   ├── config.py                # 集中配置
│   ├── model.py                 # PlateEncoder + ContrastiveLoss
│   ├── dataset.py               # PlateDataset(动态正负样本对生成)
│   ├── transforms.py            # 图像预处理与数据增强
│   ├── trainer.py               # 训练循环与 checkpoint 管理
│   ├── inferencer.py            # PyTorch 推理封装
│   ├── onnx_inferencer.py       # ONNX Runtime 推理封装
│   └── export.py                # PyTorch → ONNX 导出
├── api/
│   ├── main.py                  # FastAPI 应用(含 lifespan)
│   ├── routes.py                # API 接口实现
│   ├── schemas.py               # Pydantic 请求/响应模型
│   └── static/
│       └── index.html           # 内置 Web 测试界面
├── train.py                     # 训练入口
├── eval.py                      # 评估入口
├── prepare_data.py              # 数据划分入口
├── inference_test.py            # PyTorch 推理 CLI
├── inference_onnx.py            # ONNX 推理 CLI
├── export_onnx.py               # 模型导出 CLI
├── tests/                       # 测试套件
│   ├── test_dataset.py
│   ├── test_model.py
│   └── test_inferencer.py
├── models/
│   ├── best_model.pth           # 最佳模型
│   ├── checkpoints/             # 训练 checkpoint
│   └── export/                  # 导出 ONNX 模型
├── data/
│   ├── raw/                     # 原始数据(按车牌 ID 分组)
│   └── processed/               # 划分后数据(train/val/test)
├── pyproject.toml
├── Dockerfile
└── README.md

开源协议

本项目采用 MIT 协议开源。


作者: noimank · noimank@163.com

About

PlateTruth 是一个基于对比学习的套牌车识别系统。它将车辆图片通过 ResNet 提取为 512 维特征向量,利用余弦相似度判断两张图片是否属于同一辆车。项目提供从数据准备、模型训练、ONNX 导出到 FastAPI 部署的完整流水线,支持 PyTorch 训练与 ONNX Runtime 推理两种模式,可直接 Docker 部署。适用于安防监控、卡口布控、克隆车稽查等场景。

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages