十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

车辆识别实战:从PyTorch车型分类到YOLO目标检测

车辆识别实战:从PyTorch车型分类到YOLO目标检测 最近刷到一段街头实拍视频深夜的路面上突然出现两台安全车气场直接把旁边的车流压了一头紧接着又有一台挂着“网约车”标识的宝马入镜不少人在评论区问这种场景下拍摄的设备是什么车辆识别是怎么实现的如果要自己写一套能够自动识别车型、跟踪车辆目标的小工具该从哪里入手这篇文章就围绕“车辆识别”这个技术方向展开。我会从概念讲起到环境准备、模型原理、完整训练与预测代码再到 YOLO 检测的进阶实现和真实项目中的排错经验整理成一套可以直接照着做的实战教程。新手可以顺着步骤跑通最小示例有基础的开发者可以直接跳到第 4 节之后看代码和工程建议。1. 车辆识别到底在解决什么问题1.1 从一张街拍图说起先理清一个概念。视频里“惊现一台安全车”“宝马网约车”这类信息对人类来说一眼就能看懂但计算机看到的是由 0 到 255 组成的像素矩阵它不知道图像里哪个区域是一辆车更不知道那辆车属于什么品牌、什么型号。车辆识别Vehicle Recognition要做的就是让计算机自动完成两件事检测在图像或视频帧中定位车辆的位置通常用矩形框标出。分类判断这个框里的车是什么品牌、什么型号或者是否属于某类特殊车辆如安全车、警车、出租车网约车。如果是在连续视频流中做还会多一个任务——跟踪即判断前后两帧里的车是否为同一辆。1.2 车辆识别与传统图像分类的区别很多刚入门的人会混淆“图像分类”和“车辆识别”。图像分类面对的是整张图模型只需要输出“这张图中有什么”例如判断一张图片是猫还是狗。但车辆识别面对的是复杂场景图片里可能有车、有行人、有路灯、有广告牌模型必须先找出车辆所在区域再对区域内容做分类。所以完整流程通常是输入图像 - 目标检测 - 目标裁剪 - 车型分类 - 输出结果其中“目标检测”负责位置“车型分类”负责身份两者缺一不可。1.3 常见应用场景交通卡口和电子警察识别违法车辆的车牌、车型作为执法依据。停车场管理区分小型车、SUV、大货车按车型计费。智慧城市安防在监控视频中搜索特定车辆例如“白色宝马”“黑色SUV”。自动驾驶感知车辆识别是环境感知模块的重要组成部分。数据运营场景对街拍图片、视频素材自动打标签方便后续检索。本文的实战示例以“车型识别 目标检测”为主线适合你想在本地做一套能跑的 Demo也适合作为课程设计或简历项目的起点。2. 环境准备与版本说明本节以常见环境为例。如果你本机已经装好了 Python 和 CUDA可以直接看代码部分。2.1 硬件与操作系统操作系统Windows 10/11、Ubuntu 20.04/22.04、macOS 均可。CPU可以运行小规模数据集和推理训练建议使用 NVIDIA GPU。GPU显存建议 6GB 以上如果显存不足可以使用 Colab 或者 Kaggle Notebook 免费 GPU。2.2 Python 环境与依赖Python 版本建议 3.8 到 3.11。我不建议直接在当前环境装包而是创建虚拟环境。python -m venv vehicle_env source vehicle_env/bin/activate # Linux 或 macOS vehicle_env\Scripts\activate # Windows核心依赖torch torchvision opencv-python matplotlib numpy tqdm scikit-learn ultralytics安装命令pip install torch torchvision opencv-python matplotlib numpy tqdm scikit-learn ultralytics如果你是 NVIDIA GPU 环境建议根据 PyTorch 官网选择对应的 CUDA 版本安装命令例如pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118版本需要根据你的项目实际情况调整本文示例以常见环境为例重点演示配置思路。2.3 项目结构建议先建立如下目录vehicle_recognition/ ├── data/ │ ├── train/ │ │ ├── BMW/ │ │ ├── Benz/ │ │ └── Audi/ │ └── val/ │ ├── BMW/ │ ├── Benz/ │ └── Audi/ ├── dataset.py ├── model.py ├── train.py ├── predict.py └── requirements.txt这样可以把数据、代码、模型分开管理避免项目一变大就乱掉。3. 核心思路拆解检测、分类与数据集3.1 目标检测与图像分类的取舍如果你想从零训练一个模型先要明确任务边界如果图片中只有一辆车且车辆占画面比例大可以直接用图像分类模型例如 ResNet、MobileNet。如果图片是多车辆场景例如街拍视频截图就必须先做目标检测把每辆车从画面中裁剪出来再做分类或直接使用检测模型自带的类别输出。本文会同时覆盖这两条路。3.2 图像分类模型选择对于车型分类常用模型模型特点适用场景ResNet18结构经典训练较快数据量较少时ResNet50精度更高参数量适中中等规模数据MobileNetV3轻量快速适合移动端嵌入式或实时推理EfficientNet兼顾精度与效率有一定调参经验新手建议先跑 ResNet18因为它结构简单不容易出奇怪问题代码也好理解。3.3 目标检测模型选择YOLO 系列在车辆检测任务中YOLOYou Only Look Once系列是实际项目中使用最广泛的方案。它的核心思想是把检测问题转化为回归问题单次前向传播就能输出所有目标框的位置和类别因此推理速度非常快。当前使用较多的是 YOLOv5、YOLOv8 以及最新迭代版本。本文使用 Ultralytics 提供的 YOLOv8 接口代码简洁、文档完整适合快速验证。需要注意YOLO 系列版本迭代很快接口细节可能变化运行前建议先打印模型结构确认一下。3.4 数据集准备是关键公开的车型数据集有 CompCars、Stanford Cars 等。如果你不想下载大数据集可以自己构造一个小型数据集用于学习从网上收集三类车型图片每类 50 到 100 张放在data/train和data/val目录下。这里强调一点数据集的质量直接影响模型效果。图片应尽量多样包含不同角度、光照和背景。如果全部是同一角度、同一背景的图片模型学到的是背景特征而不是车型特征。数据增强也是不可缺少的一环。常见操作有from torchvision import transforms train_transforms transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transforms transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])随机水平翻转、随机旋转、颜色抖动都是常用的数据增强手段可以让模型在有限数据下学会更鲁棒的特征。4. 完整实战一基于 PyTorch 的车型分类模型这一节我们实现一个完整的车型分类训练流程。代码可以直接复制运行。4.1 自定义 Dataset把目录结构中的图片读入并映射为类别索引。这里使用torchvision.datasets.ImageFolder它要求目录结构为“类别名/图片文件”正好和我们之前建立的结构一致。# 文件路径dataset.py from torchvision import datasets, transforms def get_dataloaders(data_root, batch_size32): train_transforms transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transforms transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_dataset datasets.ImageFolder( rootf{data_root}/train, transformtrain_transforms ) val_dataset datasets.ImageFolder( rootf{data_root}/val, transformval_transforms ) train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue) val_loader DataLoader(val_dataset, batch_sizebatch_size, shuffleFalse) return train_loader, val_loader, train_dataset.classes注意ImageFolder会自动按照文件夹名生成类别索引排序与目录名一致。训练结束后预测时需要保持相同的类别映射顺序。4.2 定义模型使用 PyTorch 自带的 ResNet18修改最后一层全连接输出为类别数。# 文件路径model.py import torch.nn as nn from torchvision import models def create_model(num_classes): model models.resnet18(pretrainedTrue) in_features model.fc.in_features model.fc nn.Linear(in_features, num_classes) return modelpretrainedTrue表示加载在 ImageNet 上预训练好的权重。因为我们的数据集较小使用预训练权重可以加快收敛也能提高最终准确率。4.3 训练脚本训练脚本包含训练循环和验证循环并保存验证集准确率最高的模型。# 文件路径train.py import torch import torch.nn as nn import torch.optim as optim from dataset import get_dataloaders from model import create_model from tqdm import tqdm def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_loss 0 correct 0 total 0 for images, labels in tqdm(loader, descTraining): images images.to(device) labels labels.to(device) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() * images.size(0) _, preds torch.max(outputs, 1) correct (preds labels).sum().item() total labels.size(0) return total_loss / total, correct / total def validate(model, loader, criterion, device): model.eval() total_loss 0 correct 0 total 0 with torch.no_grad(): for images, labels in loader: images images.to(device) labels labels.to(device) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) _, preds torch.max(outputs, 1) correct (preds labels).sum().item() total labels.size(0) return total_loss / total, correct / total def main(): device torch.device(cuda if torch.cuda.is_available() else cpu) print(Using device:, device) train_loader, val_loader, classes get_dataloaders(data, batch_size32) model create_model(num_classeslen(classes)).to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-4) best_acc 0.0 for epoch in range(30): train_loss, train_acc train_one_epoch( model, train_loader, criterion, optimizer, device ) val_loss, val_acc validate(model, val_loader, criterion, device) print(fEpoch {epoch 1}/30) print(f Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}) print(f Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), vehicle_model.pth) print(f Best model saved, Acc: {best_acc:.4f}) if __name__ __main__: main()运行方式python train.py预期会输出类似结果Using device: cuda Training: 100%|████████████| 10/10 [00:1200:00] Epoch 1/30 Train Loss: 0.8762, Train Acc: 0.6824 Val Loss: 0.6531, Val Acc: 0.7586 Best model saved, Acc: 0.7586训练过程会因数据集大小、硬件条件而不同重点看 Val Acc 是否逐步上升。4.4 预测脚本训练完成后使用vehicle_model.pth对单张图片进行预测。# 文件路径predict.py import torch from PIL import Image from torchvision import transforms from model import create_model def predict_image(image_path, model, classes, device): transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) image Image.open(image_path).convert(RGB) input_tensor transform(image).unsqueeze(0).to(device) model.eval() with torch.no_grad(): outputs model(input_tensor) probs torch.softmax(outputs, dim1) top_prob, top_class torch.max(probs, 1) class_name classes[top_class.item()] return class_name, top_prob.item() if __name__ __main__: device torch.device(cuda if torch.cuda.is_available() else cpu) classes [Audi, BMW, Benz] # 顺序要和训练时一致 model create_model(num_classeslen(classes)).to(device) model.load_state_dict(torch.load(vehicle_model.pth, map_locationdevice)) pred_class, prob predict_image(test.jpg, model, classes, device) print(f预测结果: {pred_class}, 置信度: {prob:.4f})到这里你已经完成了一个最小可用的车型分类系统。5. 完整实战二用 YOLO 实现车辆检测如果目标是检测街拍图片中的多辆车需要从分类转向检测。5.1 准备工作安装 ultralytics 后先确认版本python -c from ultralytics import YOLO; print(YOLO.__name__)使用 YOLOv8 预训练权重进行推理from ultralytics import YOLO model YOLO(yolov8n.pt) results model(street.jpg, saveTrue) print(results[0].boxes)第一次运行时会自动下载 yolov8n.pt 权重文件。YOLO 默认的 COCO 80 类中包含 car、truck、bus 等类别可以满足常规车辆检测需求。5.2 用 YOLO 输出检测结果并裁剪只检测到车还不够还需要把检测结果保留下来。下面示例将识别出的车辆区域过滤出来并保存。import cv2 from ultralytics import YOLO model YOLO(yolov8n.pt) image_path street.jpg image cv2.imread(image_path) results model(image_path, verboseFalse)[0] vehicle_classes {2: car, 5: bus, 7: truck, 3: motorcycle} save_index 0 for box in results.boxes: cls_id int(box.cls[0]) conf float(box.conf[0]) if cls_id not in vehicle_classes: continue x1, y1, x2, y2 map(int, box.xyxy[0]) crop image[y1:y2, x1:x2] if crop.size 0: continue save_index 1 cv2.imwrite(fcar_{save_index}_{vehicle_classes[cls_id]}_{conf:.2f}.jpg, crop) cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2) label f{vehicle_classes[cls_id]} {conf:.2f} cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) cv2.imwrite(result_street.jpg, image)运行后会得到一张标注了绿框的result_street.jpg同时把每辆车单独裁剪出来保存为car_1_car_0.95.jpg这样的文件。这个流程就可以用来对街拍视频帧做车辆目标提取。5.3 自定义数据集训练 YOLO如果默认 COCO 类别不够满足需求例如需要识别“安全车”“宝马网约车”这种细分类别就需要准备自定义数据集。YOLO 自定义数据集需要两个部分图片文件夹例如images/train、images/val。标签文件夹例如labels/train、labels/val。每张图片对应一个同名.txt文件内容格式如下class_id x_center y_center width height注意坐标值需要归一化到 0 到 1 之间且 x_center、y_center 是矩形框中心点坐标不是左上角坐标。标注工具可以使用 LabelImg 或 X-AnyLabeling。标注完成后编写一个data.yamltrain: D:/vehicle_recognition/dataset/images/train val: D:/vehicle_recognition/dataset/images/val nc: 3 names: [safety_car, bmw_taxi, normal_car]然后开始训练yolo detect train datadata.yaml modelyolov8n.pt epochs100 imgsz640 batch16训练完成后模型权重保存在runs/detect/train/weights/best.pt。使用它替代默认权重即可model YOLO(runs/detect/train/weights/best.pt) results model(street.jpg, saveTrue)关于安全问题这里特别提醒如果你在真实交通场景中做测试应确保使用的是公开数据或自己合法采集的数据避开对特定车辆、个人隐私信息的没有授权的采集和发布。涉及城市道路拍摄、监控数据使用时需要遵守当地法律法规和平台审核规则。5.4 视频车辆检测对视频文件的检测逻辑类似ultralytics 封装了直接推理接口from ultralytics import YOLO model YOLO(yolov8n.pt) results model(street_video.mp4, saveTrue)实现原理是逐帧读取视频在每一帧上执行检测再按顺序写回视频。如果你需要同时保留车辆裁剪结果可以自定义一个视频读取循环在每帧执行检测代码。6. 常见问题与排查思路问题现象常见原因解决思路训练时 loss 不下降学习率过大或过小尝试从 1e-3 和 1e-4 开始调整观察 loss 曲线验证集准确率低但训练集准确率高过拟合增加数据增强、引入 Dropout、减小模型容量或增加数据量CUDA out of memorybatch_size 过大减小 batch_size例如 32 改 16或减小输入图片尺寸ImageFolder 读入的类别数不对目录结构错误确认 train 目录下每个类别文件夹必须包含图片YOLO 下权重失败网络问题或版本不兼容手动下载权重文件放到本地路径或升级 ultralytics视频推理速度慢没有使用 GPU 或 batch 设置不合理确认torch.cuda.is_available()使用device0参数检测框覆盖区域过小真实目标小或分辨率低使用更高分辨率输入例如 imgsz1280或使用更深的模型如果你在跑上面的代码时遇到其他报错建议按以下顺序排查查看完整的错误堆栈确认是数据加载、模型前向传播还是损失计算阶段的问题。打印输入张量的 shape 和模型输出 shape确认维度匹配。单步执行一个小批量数据排除数据损坏问题。搜索报错信息时注意版本环境不要盲目复制高版本的解决方案。7. 最佳实践与工程建议7.1 数据质量优先于模型结构很多新手喜欢先堆模型但实际项目里数据质量对最终效果的影响往往更大。以下是几条实在建议类别样本数量尽量均衡避免某类图片过多导致模型倾斜。预留独立的验证集和测试集不要只靠训练集判断模型好坏。清洗掉带水印、严重模糊、无车辆主体的图片。对类别标签做人工复审避免错误标注。7.2 模型版本管理建议为每个训练实验记录以下信息数据集版本 模型结构 输入图片尺寸 训练超参数学习率、batch_size、epochs 验证集准确率 训练时间可以用 CSV 文件维护也可以用 Weights Biases 这类工具。记录完整的实验信息才能在模型迭代时定位效果变化的原因。7.3 推理性能优化实际部署中推理速度和使用场景强相关服务端批量处理可用 FP16 混合精度加速。边缘设备部署优先选择轻量模型如 MobileNet、YOLOv8n。视频流处理采用多线程或异步队列避免 I/O 阻塞。模型剪枝与量化使用 TensorRT、ONNX Runtime 等工具转换模型格式降低延迟。7.4 合规与隐私边界车辆识别往往涉及公共空间图像数据。在这里要提醒读者使用公开数据集时注意数据集的版权和许可协议。自己采集数据时应避免收集不必要的个人信息例如人脸。在真实道路场景中部署测试应获得相应授权遵守道路交通安全法及相关隐私法规。发布技术文章或演示截图时对车牌、人脸等敏感信息做模糊处理。8. 总结与下一步学习本文从一辆街拍视频中“识别出安全车和网约车”这个真实需求切入讲清楚了车辆识别领域的两个核心任务目标检测和车型分类。你跟着完成了一个基于 PyTorch 的车型分类模型训练与预测也学会了使用 YOLO 对街拍图像做车辆检测与裁剪这些代码可以直接改造成你的课程设计或简历项目。下一步可以继续深入的方向有三个把“检测 分类”串成完整 pipeline先用 YOLO 检测并裁剪车辆再送入分类模型判断品牌型号。学习跟踪算法例如 ByteTrack、DeepSORT让视频中的车辆拥有稳定 ID。研究模型部署把训练好的模型转换为 ONNX 或 TensorRT跑在真实服务或边缘设备上。如果你对某个环节有疑问例如自定义数据集标注、训练参数调优、模型部署报错可以先按第 6 节的排查表过一遍再动手改代码。技术学习最忌讳只看不练建议现在就打开终端把第 4 节的最小示例跑通。
返回列表