1. 项目概述
交通标志识别是计算机视觉领域的一个重要应用方向,尤其在自动驾驶和智能交通系统中扮演着关键角色。今天我要分享的是一个基于TT100k数据集的交通标志识别项目,使用YOLOv3模型和Python实现。这个项目特别适合刚入门计算机视觉的朋友,因为它不仅涵盖了完整的实现流程,还针对160多个交通标志类别进行了优化,在各种光照和视角条件下都能保持较好的识别效果。
提示:虽然YOLOv5等新模型已经出现,但YOLOv3仍然是学习目标检测的优秀选择,它的结构相对简单但性能依然强大,非常适合教学和入门项目。
2. 数据集准备与处理
2.1 TT100k数据集介绍
TT100k(Tsinghua-Tencent 100K)是一个大规模交通标志数据集,包含10万张图像和3万个交通标志实例,涵盖了中国道路上的多种交通标志。数据集中的图像采集自不同天气条件、光照环境和视角,非常适合训练具有鲁棒性的识别模型。
数据集中的交通标志被分为45个主要类别和160多个子类别,包括禁令标志、指示标志、警告标志等多种类型。每个标志都提供了精确的边界框标注,标注格式支持多种目标检测框架。
2.2 数据集目录结构
为了高效管理数据集,我建议采用以下目录结构:
code复制tt100k/
├── annotations/ # 原始标注文件
├── images/ # 原始图像
├── yolov3_format/ # YOLOv3格式数据
│ ├── train/
│ │ ├── images/ # 训练图像
│ │ └── labels/ # 训练标注
│ └── val/
│ ├── images/ # 验证图像
│ └── labels/ # 验证标注
├── classes.txt # 类别列表
└── split.py # 数据集划分脚本
2.3 数据格式转换
TT100k原始标注是JSON格式,需要转换为YOLOv3要求的TXT格式。转换过程需要注意以下几点:
- 坐标归一化:YOLOv3要求标注框中心坐标和宽高都是相对于图像宽高的比例值(0-1)
- 类别ID映射:需要将原始类别ID映射为连续的整数(0-159)
- 数据增强:可以考虑在转换过程中加入随机水平翻转等简单增强
以下是转换脚本的核心部分:
python复制def convert_annotation(json_file, output_dir):
with open(json_file) as f:
data = json.load(f)
for img_info in data['imgs'].values():
img_id = img_info['id']
img_w = img_info['width']
img_h = img_info['height']
txt_path = os.path.join(output_dir, f'{img_id}.txt')
with open(txt_path, 'w') as f_txt:
for obj in img_info['objects']:
category = obj['category']
class_id = class_mapping[category]
# 获取边界框坐标并归一化
x1, y1, x2, y2 = obj['bbox']['xmin'], obj['bbox']['ymin'], \
obj['bbox']['xmax'], obj['bbox']['ymax']
x_center = ((x1 + x2) / 2) / img_w
y_center = ((y1 + y2) / 2) / img_h
width = (x2 - x1) / img_w
height = (y2 - y1) / img_h
f_txt.write(f'{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n')
3. 模型构建与训练
3.1 YOLOv3模型架构
YOLOv3采用Darknet-53作为骨干网络,结合多尺度预测机制,能够在保持较高检测速度的同时获得不错的准确率。模型的主要特点包括:
- 多尺度预测:在3个不同尺度的特征图上进行预测,分别对应大、中、小目标的检测
- 锚框(Anchor)机制:使用K-means聚类得到的先验框尺寸,提高检测精度
- 残差连接:Darknet-53中大量使用残差连接,缓解深层网络梯度消失问题
3.2 模型实现
虽然可以完全从头实现YOLOv3,但为了简化开发,我们可以使用PyTorch的预训练模型作为基础:
python复制import torch
import torch.nn as nn
from torchvision.models import resnet50
class YOLOv3(nn.Module):
def __init__(self, num_classes=160):
super(YOLOv3, self).__init__()
# 使用ResNet50作为骨干网络(替代Darknet-53)
self.backbone = resnet50(pretrained=True)
# 移除最后的全连接层
self.backbone = nn.Sequential(*list(self.backbone.children())[:-2])
# 添加YOLOv3特有的检测头
self.detection_head = self._make_detection_head(2048, num_classes)
def _make_detection_head(self, in_channels, num_classes):
return nn.Sequential(
nn.Conv2d(in_channels, 512, kernel_size=3, padding=1),
nn.BatchNorm2d(512),
nn.LeakyReLU(0.1),
nn.Conv2d(512, (5 + num_classes) * 3, kernel_size=1) # 3个锚框
)
def forward(self, x):
features = self.backbone(x)
predictions = self.detection_head(features)
return predictions
注意:这里使用ResNet50替代了原版的Darknet-53,虽然性能略有差异,但更便于快速实现和训练。对于追求更高准确率的场景,建议实现原版Darknet-53。
3.3 训练策略
训练YOLOv3模型需要特别注意以下几点:
- 学习率调度:采用余弦退火学习率,初始学习率设为0.001
- 损失函数:YOLOv3使用复合损失,包括坐标损失、置信度损失和分类损失
- 数据增强:随机裁剪、颜色抖动、马赛克增强等
以下是训练循环的核心代码:
python复制def train(model, train_loader, optimizer, epoch):
model.train()
total_loss = 0
for batch_idx, (images, targets) in enumerate(train_loader):
images = images.to(device)
targets = targets.to(device)
optimizer.zero_grad()
outputs = model(images)
# 计算YOLOv3的复合损失
loss = compute_loss(outputs, targets)
loss.backward()
optimizer.step()
total_loss += loss.item()
if batch_idx % 100 == 0:
print(f'Train Epoch: {epoch} [{batch_idx}/{len(train_loader)}]\tLoss: {loss.item():.6f}')
avg_loss = total_loss / len(train_loader)
print(f'====> Epoch: {epoch} Average loss: {avg_loss:.4f}')
return avg_loss
def compute_loss(predictions, targets):
# 实现YOLOv3的损失函数
# 包括坐标损失、置信度损失和分类损失
# 这里简化了实现,实际需要更复杂的处理
coord_loss = F.mse_loss(predictions[..., :4], targets[..., :4])
conf_loss = F.binary_cross_entropy(predictions[..., 4], targets[..., 4])
cls_loss = F.cross_entropy(predictions[..., 5:], targets[..., 5:].argmax(-1))
return coord_loss + conf_loss + cls_loss
4. 模型优化与调参
4.1 锚框聚类
YOLOv3使用K-means算法对训练集中的标注框进行聚类,得到适合数据集的先验框尺寸。对于TT100k数据集,我们可以这样实现:
python复制from sklearn.cluster import KMeans
def cluster_annotations(annotation_dir, num_clusters=9):
boxes = []
# 收集所有标注框的宽高
for txt_file in os.listdir(annotation_dir):
with open(os.path.join(annotation_dir, txt_file)) as f:
for line in f:
_, _, _, w, h = map(float, line.strip().split())
boxes.append([w, h])
# 使用K-means聚类
kmeans = KMeans(n_clusters=num_clusters)
kmeans.fit(boxes)
# 按面积排序聚类中心
anchors = kmeans.cluster_centers_
anchors = sorted(anchors, key=lambda x: x[0]*x[1])
return anchors
4.2 多尺度训练
为了提高模型对不同尺寸交通标志的检测能力,可以采用多尺度训练策略:
- 每隔几个epoch改变输入图像的尺寸
- 典型尺寸包括320×320、416×416、608×608等
- 需要在改变尺寸时同步调整网络结构
python复制def random_resize(image, targets, min_size=320, max_size=608):
# 随机选择新的训练尺寸
new_size = random.choice(range(min_size, max_size+1, 32))
# 调整图像尺寸
h, w = image.shape[1:]
ratio = new_size / max(h, w)
new_h, new_w = int(h * ratio), int(w * ratio)
resized_image = F.interpolate(image.unsqueeze(0), size=(new_h, new_w), mode='bilinear')
# 调整目标框坐标
targets[..., 1:5] *= ratio
return resized_image.squeeze(0), targets
4.3 学习率策略
采用渐进式学习率预热和余弦退火策略:
python复制from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 前5个epoch线性增加学习率
warmup_scheduler = LinearLR(optimizer, start_factor=0.1, total_iters=5)
# 之后使用余弦退火
cosine_scheduler = CosineAnnealingLR(optimizer, T_max=95, eta_min=0.0001)
for epoch in range(100):
train(...)
if epoch < 5:
warmup_scheduler.step()
else:
cosine_scheduler.step()
5. 模型评估与推理
5.1 评估指标
交通标志识别模型的评估通常采用以下指标:
- mAP(mean Average Precision):在不同IoU阈值下的平均精度
- FPS(Frames Per Second):推理速度
- 各类别的精确率和召回率
python复制def evaluate(model, val_loader):
model.eval()
all_detections = []
all_annotations = []
with torch.no_grad():
for images, targets in val_loader:
images = images.to(device)
outputs = model(images)
# 处理模型输出
detections = postprocess(outputs)
all_detections.extend(detections)
all_annotations.extend(targets)
# 计算mAP
ap_per_class, mAP = compute_map(all_detections, all_annotations)
print(f'mAP: {mAP:.4f}')
for class_id, ap in enumerate(ap_per_class):
print(f'Class {class_id}: AP = {ap:.4f}')
return mAP
5.2 推理优化
为了提高推理速度,可以采用以下优化技术:
- TensorRT加速:将模型转换为TensorRT引擎
- 半精度推理:使用FP16减少计算量和内存占用
- 批处理推理:同时处理多张图像
python复制def inference(model, image_path, conf_thresh=0.5, nms_thresh=0.4):
# 图像预处理
image = cv2.imread(image_path)
h, w = image.shape[:2]
input_tensor = preprocess(image).to(device)
# 模型推理
with torch.no_grad():
predictions = model(input_tensor.unsqueeze(0))
# 后处理
boxes, scores, classes = postprocess(predictions, conf_thresh, nms_thresh)
# 绘制结果
for box, score, cls in zip(boxes, scores, classes):
x1, y1, x2, y2 = box
x1, y1, x2, y2 = int(x1 * w), int(y1 * h), int(x2 * w), int(y2 * h)
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(image, f'{cls_names[cls]}:{score:.2f}',
(x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
return image
6. 实际应用与部署
6.1 模型轻量化
为了在边缘设备上部署,可以考虑以下轻量化方法:
- 知识蒸馏:使用大模型指导小模型训练
- 通道剪枝:移除不重要的卷积通道
- 量化:将FP32模型转换为INT8
python复制# 量化示例
quantized_model = torch.quantization.quantize_dynamic(
model, # 原始模型
{torch.nn.Linear, torch.nn.Conv2d}, # 要量化的模块类型
dtype=torch.qint8 # 量化类型
)
6.2 部署方案
根据应用场景可以选择不同的部署方式:
- 服务器部署:使用Flask或FastAPI创建REST API
- 移动端部署:转换为Core ML或TFLite格式
- 嵌入式部署:使用TensorRT或OpenVINO优化
以下是简单的Flask API实现:
python复制from flask import Flask, request, jsonify
import cv2
import numpy as np
app = Flask(__name__)
model = load_model('best_model.pth')
@app.route('/predict', methods=['POST'])
def predict():
if 'file' not in request.files:
return jsonify({'error': 'No file uploaded'})
file = request.files['file']
image = cv2.imdecode(np.frombuffer(file.read(), np.uint8), cv2.IMREAD_COLOR)
# 预处理和推理
input_tensor = preprocess(image)
with torch.no_grad():
predictions = model(input_tensor.unsqueeze(0))
# 后处理
results = process_predictions(predictions, image.shape)
return jsonify(results)
if __name__ == '__main__':
app.run(host='0.0.0.0', port=5000)
7. 常见问题与解决方案
7.1 小目标检测效果差
交通标志中有些目标较小,检测效果可能不理想。解决方法包括:
- 增加小目标的训练样本
- 使用更高分辨率的输入图像
- 调整锚框尺寸,增加适合小目标的锚框
7.2 夜间或低光照条件性能下降
提高模型在低光照条件下的鲁棒性:
- 数据增强时加入亮度、对比度调整
- 在训练集中增加夜间场景样本
- 预处理阶段加入自动白平衡或直方图均衡化
7.3 类别不平衡问题
某些交通标志类别样本较少,导致检测效果不佳:
- 使用类别加权损失函数
- 对少数类别样本进行过采样
- 应用焦点损失(Focal Loss)处理难易样本不平衡
python复制class FocalLoss(nn.Module):
def __init__(self, alpha=0.25, gamma=2):
super(FocalLoss, self).__init__()
self.alpha = alpha
self.gamma = gamma
def forward(self, inputs, targets):
BCE_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
pt = torch.exp(-BCE_loss)
focal_loss = self.alpha * (1-pt)**self.gamma * BCE_loss
return focal_loss.mean()
8. 项目扩展与优化方向
8.1 多任务学习
除了检测交通标志,还可以同时完成其他相关任务:
- 交通标志状态识别(如红绿灯状态)
- 交通标志文字识别
- 道路场景分割
python复制class MultiTaskModel(nn.Module):
def __init__(self, num_classes):
super().__init__()
self.backbone = resnet50(pretrained=True)
# 检测头
self.detection_head = DetectionHead(2048, num_classes)
# 分割头
self.segmentation_head = SegmentationHead(2048)
def forward(self, x):
features = self.backbone(x)
detections = self.detection_head(features)
segmentation = self.segmentation_head(features)
return detections, segmentation
8.2 模型融合
结合多个模型的优势:
- 使用YOLOv3进行快速初筛
- 对不确定的检测框用更精确的模型(如Faster R-CNN)二次验证
- 融合多个模型的预测结果
8.3 实时视频处理
将模型应用于实时视频流:
- 使用多线程处理:一个线程负责图像采集,一个线程负责推理
- 应用目标跟踪算法减少计算量
- 实现基于运动检测的感兴趣区域提取
python复制def process_video(video_path, model):
cap = cv2.VideoCapture(video_path)
tracker = cv2.TrackerCSRT_create() # 创建跟踪器
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# 每隔N帧或跟踪失败时重新检测
if frame_count % 30 == 0 or not tracking_success:
detections = model.detect(frame)
if detections:
tracker.init(frame, detections[0]['bbox'])
# 更新跟踪器
tracking_success, bbox = tracker.update(frame)
# 绘制结果
if tracking_success:
draw_box(frame, bbox)
cv2.imshow('Result', frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
在实际项目中,我发现交通标志识别系统的性能很大程度上依赖于数据质量。特别是在中国道路场景下,交通标志的种类繁多,加上各种天气和光照条件的影响,构建一个鲁棒的识别系统确实需要大量的工程实践和经验积累。建议初学者先从这个小项目开始,逐步深入理解目标检测的各个环节,然后再尝试更复杂的应用场景。