### 基于 PyTorch 的 RCNN 和 LSTM 实现车辆与行人识别的完整代码 以下是使用 PyTorch 结合 Faster R-CNN 和 LSTM 来实现车辆和行人识别的完整代码示例。该方法适用于带有时序特性的图像数据集，并支持从 XML 文件解析标注信息的功能。 --- #### 1. 数据加载与预处理 首先编写函数来读取图像序列及其对应的 XML 标注文件，同时将这些数据整理为适合模型输入的形式。 ```python import os import torch from PIL import Image from torchvision.transforms import functional as F from lxml import etree def parse_xml(file_path): """ 解析单个XML文件的内容。 :param file_path: XML文件路径 :return: 包含边界框坐标和类别的字典列表 [{'boxes': [...], 'labels': [...]}, ...] """ with open(file_path, 'r') as f: tree = etree.parse(f) objects = [] for obj in tree.findall('object'): bbox = obj.find('bndbox') xmin = int(bbox.find('xmin').text) ymin = int(bbox.find('ymin').text) xmax = int(bbox.find('xmax').text) ymax = int(bbox.find('ymax').text) label = obj.find('name').text if label not in ['vehicle', 'pedestrian']: continue objects.append({ 'boxes': [xmin, ymin, xmax, ymax], 'labels': 1 if label == 'vehicle' else 2 }) return objects def collate_fn(batch): """ 自定义批处理函数，适配不同长度的时间序列。 """ imgs, targets = zip(*batch) return list(imgs), list(targets) class CustomDataset(torch.utils.data.Dataset): def __init__(self, image_dir, annotation_dir, transform=None, sequence_length=5): self.image_dir = image_dir self.annotation_dir = annotation_dir self.transform = transform self.sequence_length = sequence_length self.images = sorted(os.listdir(image_dir)) self.annotations = sorted(os.listdir(annotation_dir)) def __len__(self): return len(self.images) - self.sequence_length + 1 def __getitem__(self, idx): frames = [] target_list = [] for t in range(idx, idx + self.sequence_length): img_path = os.path.join(self.image_dir, self.images[t]) anno_path = os.path.join(self.annotation_dir, self.annotations[t]) img = Image.open(img_path).convert("RGB") target = parse_xml(anno_path) if self.transform: img = self.transform(img) frames.append(F.to_tensor(img)) target_list.extend(target) combined_target = { 'boxes': torch.tensor([t['boxes'] for t in target_list], dtype=torch.float32), 'labels': torch.tensor([t['labels'] for t in target_list], dtype=torch.int64) } return torch.stack(frames), combined_target ``` --- #### 2. 构建模型 接下来创建一个结合了 Faster R-CNN 和 LSTM 的神经网络架构。 ```python import torchvision from torchvision.models.detection.faster_rcnn import FastRCNNPredictor from torch.nn import LSTM, Linear class RCNN_LSTM(torch.nn.Module): def __init__(self, num_classes=3, hidden_dim=256, lstm_layers=2): super(RCNN_LSTM, self).__init__() # 初始化 Faster R-CNN 模型 backbone = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True) in_features = backbone.roi_heads.box_predictor.cls_score.in_features backbone.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) self.rcnn = backbone # 添加 LSTM 层 self.lstm = LSTM(input_size=in_features, hidden_size=hidden_dim, num_layers=lstm_layers, batch_first=True) self.fc = Linear(hidden_dim, num_classes - 1) # 减去背景类别 def forward(self, images, targets=None): features = [] predictions = [] for img in images: output = self.rcnn([img]) # 对每张图片单独运行R-CNN features.append(output[-1]['pool'].flatten(start_dim=1)) # 提取ROI池化后的特征向量 feature_seq = torch.stack(features) lstm_out, _ = self.lstm(feature_seq) final_output = self.fc(lstm_out[:, -1, :]) # 取最后一个时间步的结果 return final_output ``` --- #### 3. 训练流程 最后设置训练循环以优化模型性能。 ```python device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') dataset = CustomDataset('/path/to/images', '/path/to/xml_annotations', transform=F.to_tensor()) data_loader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=True, collate_fn=collate_fn) model = RCNN_LSTM(num_classes=3).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.001) for epoch in range(10): # 运行10轮迭代 model.train() total_loss = 0 for imgs, targets in data_loader: imgs = [img.to(device) for img in imgs] targets = [{k: v.to(device) for k, v in t.items()} for t in targets] optimizer.zero_grad() loss_dict = model(imgs, targets) losses = sum(loss for loss in loss_dict.values()) losses.backward() optimizer.step() total_loss += losses.item() print(f"Epoch {epoch+1} completed. Total Loss: {total_loss:.4f}") ``` --- ###