基于深度学习的闯红灯检测系统的PyTorch

闯红灯检测是一个重要的交通安全问题，基于深度学习的闯红灯检测系统能够实时检测交通违规行为，提高交通安全水平。本文将介绍基于PyTorch实现的闯红灯检测系统。

## 数据集

首先需要准备数据集，可以使用公开数据集或者自己收集数据。数据集应包含正常行驶、闯红灯、转弯等多种场景的视频或图像。可以使用视频帧提取工具将视频转换为图像。对于每个图像，需要手动标注其是否为闯红灯行为，标注的结果保存在XML文件中。

## 模型训练

本文使用的是Faster R-CNN模型，该模型在目标检测领域取得了很好的效果。在PyTorch中，可以使用torchvision.models.detection中的FasterRCNN模型实现。

首先需要定义数据集的类，继承自torch.utils.data.Dataset，用于加载图像和标注数据，并将其转换为PyTorch所需的格式。

import torch.utils.data
from PIL import Image

class RedLightDataset(torch.utils.data.Dataset):
    def __init__(self, img_list, transform=None):
        self.img_list = img_list
        self.transform = transform

    def __getitem__(self, idx):
        img_path = self.img_list[idx]
        img = Image.open(img_path).convert("RGB")
        label_path = img_path.replace(".jpg", ".xml")
        label = parse_label(label_path)
        if self.transform:
            img, label = self.transform(img, label)
        return img, label

    def __len__(self):
        return len(self.img_list)

接下来需要定义模型，使用FasterRCNN模型。在模型定义中，需要指定输入图像的大小、输出类别数等参数。

import torchvision.models.detection as detection

class RedLightDetector(torch.nn.Module):
    def __init__(self):
        super(RedLightDetector, self).__init__()
        self.model = detection.fasterrcnn_resnet50_fpn(pretrained=True)
        self.model.roi_heads.box_predictor.cls_score.out_features = 2

    def forward(self, x):
        return self.model(x)

定义完模型后，可以开始训练模型。训练过程中需要定义损失函数、优化器等参数。在训练过程中使用DataLoader将数据集分批加载，每个batch的大小可以根据GPU内存大小进行调整。

import torch.optim as optim
import torch.nn.functional as F

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def train(model, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss_dict = output["loss_classifier"] + output["loss_box_reg"]
        losses = sum(loss_dict.values())
        losses.backward()
        optimizer.step()
        if batch_idx % 10 == 0:
            print("Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), losses.item()))

def main():
    train_list = glob.glob("train/*.jpg")
    train_dataset = RedLightDataset(train_list, transform=train_transforms)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=4, shuffle=True)

    model = RedLightDetector().to(device)
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    for epoch in range(1, 10):
        train(model, train_loader, optimizer, epoch)

## 模型测试

模型训练完成后，可以使用测试集对模型进行测试，并计算准确率、召回率等指标。测试集的格式与训练集类似，也需要定义一个数据集类来加载数据。

class RedLightTestDataset(torch.utils.data.Dataset):
    def __init__(self, img_list, transform=None):
        self.img_list = img_list
        self.transform = transform

    def __getitem__(self, idx):
        img_path = self.img_list[idx]
        img = Image.open(img_path).convert("RGB")
        if self.transform:
            img = self.transform(img)
        return img, img_path

    def __len__(self):
        return len(self.img_list)

def test(model, test_loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            _, predicted = torch.max(output.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
    print("Accuracy: {:.2f}%".format(100. * correct / total))

使用测试集对模型进行测试：

test_list = glob.glob("test/*.jpg")
test_dataset = RedLightTestDataset(test_list, transform=test_transforms)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=4, shuffle=False)

test(model, test_loader)

## 结论

本文介绍了基于PyTorch实现的闯红灯检测系统。使用Faster R-CNN模型进行训练和测试，可以在一定程度上检测闯红灯行为。但是，模型在一些复杂场景下仍存在误检测的问题，需要进一步优化和改进。