车辆分类器 - 兄弟，这是什么车？

一个经过微调的 EfficientNet-B4 模型，用于从照片中识别车辆品牌、型号和年份。这是"兄弟，这是什么车？"车辆识别系统的一部分。

模型描述

该模型将车辆分类为 8,949 个独特类别，涵盖品牌、型号和年份的组合。它在 VMMRdb（车辆品牌和型号识别）数据集上进行了训练。

架构

• 基础模型: EfficientNet-B4（通过 timm）
• 输入尺寸: 380 x 380 像素
• 输出: 8,949 个类别（格式：品牌 型号 年份）
• 参数量: 约 1900 万

性能

• Top-1 准确率: 约 50%
• Top-5 准确率: 约 75-80%

注意：由于不同年款车型和配置级别之间存在细微差异，车辆分类具有挑战性。在实际应用中，Top-5 准确率更具参考价值。

预期用途

• 从照片中识别车辆
• 汽车库存管理
• 保险理赔处理
• 执法调查
• 停车和交通系统

文件

快速开始

PyTorch

import torch
import timm
from PIL import Image
from torchvision import transforms

# 加载模型
checkpoint = torch.load("vehicle_classifier.pth", map_location="cpu")
model = timm.create_model("efficientnet_b4", pretrained=False, num_classes=8949)
model.load_state_dict(checkpoint["model_state"])
model.eval()

# 预处理
transform = transforms.Compose([
    transforms.Resize((380, 380)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 预测
image = Image.open("car.jpg").convert("RGB")
input_tensor = transform(image).unsqueeze(0)

with torch.no_grad():
    logits = model(input_tensor)
    probs = torch.softmax(logits, dim=1)
    top5_probs, top5_indices = torch.topk(probs, 5)

# 解码预测结果
class_mapping = checkpoint["class_mapping"]  # {id: "品牌 型号 年份"}
for prob, idx in zip(top5_probs[0], top5_indices[0]):
    print(f"{class_mapping[idx.item()]}: {prob.item()*100:.1f}%")

ONNX Runtime

import onnxruntime as ort
import numpy as np
from PIL import Image

# 加载模型
session = ort.InferenceSession("vehicle_classifier.onnx")

# 预处理（与 PyTorch 相同）
image = Image.open("car.jpg").convert("RGB").resize((380, 380))
img_array = np.array(image).astype(np.float32) / 255.0
img_array = (img_array - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
img_array = img_array.transpose(2, 0, 1)[np.newaxis, ...]

# 预测
outputs = session.run(None, {"input": img_array.astype(np.float32)})
logits = outputs[0]

# 获取 Top 预测结果
probs = np.exp(logits) / np.sum(np.exp(logits), axis=1, keepdims=True)
top5_indices = np.argsort(probs[0])[-5:][::-1]

使用 Hugging Face Hub

from huggingface_hub import hf_hub_download

# 下载模型
model_path = hf_hub_download(repo_id="Jordo23/vehicle-classifier", filename="vehicle_classifier.pth")
onnx_path = hf_hub_download(repo_id="Jordo23/vehicle-classifier", filename="vehicle_classifier.onnx")

训练详情

• 数据集: VMMRdb（车辆品牌和型号识别数据库）
• 训练图像: 约 280,000 张
• 类别: 8,949 个（品牌 + 型号 + 年份组合）
• 训练轮数: 从 ImageNet 预训练权重微调
• 优化器: AdamW
• 损失函数: 交叉熵

局限性

• 在清晰、光线良好的照片上表现最佳
• 可能在以下情况下表现不佳：
  • 严重遮挡的车辆
  • 非常规角度（俯视、极端特写）
  • 训练数据中未包含的非常老旧/稀有车辆
  • 经过大幅改装或定制车辆
• 颜色检测由单独模块处理（不属于此模型）

引用

如果您使用此模型，请引用：

@misc{vehicle-classifier-2024,
  author = {Jordo23},
  title = {Vehicle Classifier - Dude, What's My Car?},
  year = {2024},
  publisher = {Hugging Face},
  url = {https://huggingface.co/Jordo23/vehicle-classifier}
}

许可证

MIT 许可证 - 可免费用于个人和商业用途。

联系方式

如有问题或疑问，请在模型仓库中提交 issue。

"兄弟，这是什么车？"车辆识别系统的一部分 🚗