4. 基于深度学习的立体匹配

任务目标

本案例使用PSMNet（Pyramid Stereo Matching Network）完成双目图像的立体匹配任务。你将加载预训练权重，对一对双目图像进行前向推理，输出视差图，并计算与真实视差图（若可用）之间的评估指标（EPE与D1-all）。案例聚焦于推理闭环，不包含训练流程。

环境与依赖

需要PyTorch（≥1.8）、OpenCV、NumPy。PSMNet原版实现基于Python 2与旧版PyTorch，本案例使用兼容PyTorch 2.x的社区适配版本。若需训练，还需下载Scene Flow或KITTI数据集，但本案例仅使用预训练权重完成推理。

实现代码

以下代码实现完整的PSMNet推理流程：加载模型与权重、读取双目图像、计算视差图、评估精度。模型结构定义精简自原版PSMNet，保留核心的SPP模块与3D卷积堆叠。

```python import torch import torch.nn as nn import torch.nn.functional as F import cv2 import numpy as np # ---------- PSMNet核心结构（精简版） ---------- class BasicBlock(nn.Module): def __init__(self, in_planes, planes, stride=1): super().__init__() self.conv1 = nn.Conv2d(in_planes, planes, 3, stride, 1, bias=False) self.bn1 = nn.BatchNorm2d(planes) self.conv2 = nn.Conv2d(planes, planes, 3, 1, 1, bias=False) self.bn2 = nn.BatchNorm2d(planes) self.relu = nn.ReLU(inplace=True) if stride != 1 or in_planes != planes: self.downsample = nn.Sequential( nn.Conv2d(in_planes, planes, 1, stride, bias=False), nn.BatchNorm2d(planes)) else: self.downsample = None def forward(self, x): identity = x out = self.relu(self.bn1(self.conv1(x))) out = self.bn2(self.conv2(out)) if self.downsample is not None: identity = self.downsample(identity) out += identity return self.relu(out) class FeatureExtractor(nn.Module): """CNN特征提取，输出1/4分辨率特征图""" def __init__(self): super().__init__() self.conv1 = nn.Sequential( nn.Conv2d(3, 32, 3, 2, 1, bias=False), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.Conv2d(32, 32, 3, 1, 1, bias=False), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.Conv2d(32, 32, 3, 1, 1, bias=False), nn.BatchNorm2d(32), nn.ReLU(inplace=True)) self.conv2 = nn.Sequential( nn.Conv2d(32, 64, 3, 1, 1, bias=False), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.Conv2d(64, 64, 3, 1, 1, bias=False), nn.BatchNorm2d(64), nn.ReLU(inplace=True)) self.conv3 = nn.Sequential( nn.Conv2d(64, 128, 3, 2, 1, bias=False), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.Conv2d(128, 128, 3, 1, 1, bias=False), nn.BatchNorm2d(128), nn.ReLU(inplace=True)) self.conv4 = nn.Sequential(BasicBlock(128, 128)) def forward(self, x): x = self.conv1(x) # 1/2 x = self.conv2(x) # 1/2 x = self.conv3(x) # 1/4 x = self.conv4(x) # 1/4 return x class SPP(nn.Module): """空间金字塔池化，扩大感受野""" def __init__(self): super().__init__() self.branch1 = nn.Sequential( nn.AvgPool2d(64, stride=64), nn.Conv2d(128, 32, 1, bias=False), nn.BatchNorm2d(32), nn.ReLU(inplace=True)) self.branch2 = nn.Sequential( nn.AvgPool2d(32, stride=32), nn.Conv2d(128, 32, 1, bias=False), nn.BatchNorm2d(32), nn.ReLU(inplace=True)) self.branch3 = nn.Sequential( nn.AvgPool2d(16, stride=16), nn.Conv2d(128, 32, 1, bias=False), nn.BatchNorm2d(32), nn.ReLU(inplace=True)) self.branch4 = nn.Sequential( nn.AvgPool2d(8, stride=8), nn.Conv2d(128, 32, 1, bias=False), nn.BatchNorm2d(32), nn.ReLU(inplace=True)) self.fuse = nn.Sequential( nn.Conv2d(256, 128, 3, 1, 1, bias=False), nn.BatchNorm2d(128), nn.ReLU(inplace=True)) def forward(self, x): b, c, h, w = x.shape p1 = F.interpolate(self.branch1(x), (h, w), mode='bilinear', align_corners=False) p2 = F.interpolate(self.branch2(x), (h, w), mode='bilinear', align_corners=False) p3 = F.interpolate(self.branch3(x), (h, w), mode='bilinear', align_corners=False) p4 = F.interpolate(self.branch4(x), (h, w), mode='bilinear', align_corners=False) return self.fuse(torch.cat([x, p1, p2, p3, p4], dim=1)) class CostVolume(nn.Module): """构建4D代价体并正则化""" def __init__(self, max_disp=192): super().__init__() self.max_disp = max_disp self.conv3d_1 = nn.Sequential( nn.Conv3d(128, 64, 3, padding=1, bias=False), nn.BatchNorm3d(64), nn.ReLU(inplace=True), nn.Conv3d(64, 64, 3, padding=1, bias=False), nn.BatchNorm3d(64), nn.ReLU(inplace=True)) self.conv3d_2 = nn.Sequential( nn.Conv3d(64, 64, 3, padding=1, bias=False), nn.BatchNorm3d(64), nn.ReLU(inplace=True), nn.Conv3d(64, 64, 3, padding=1, bias=False), nn.BatchNorm3d(64), nn.ReLU(inplace=True)) self.conv3d_3 = nn.Sequential( nn.Conv3d(64, 64, 3, padding=1, bias=False), nn.BatchNorm3d(64), nn.ReLU(inplace=True), nn.Conv3d(64, 64, 3, padding=1, bias=False), nn.BatchNorm3d(64), nn.ReLU(inplace=True)) self.conv3d_out = nn.Conv3d(64, 1, 3, padding=1, bias=False) def forward(self, featL, featR): b, c, h, w = featL.shape cost = torch.zeros(b, self.max_disp, h, w, device=featL.device) for d in range(self.max_disp): if d == 0: cost[:, d] = torch.sum((featL - featR) ** 2, dim=1) else: shift = torch.roll(featR, shifts=d, dims=3) shift[:, :, :, :d] = 0 cost[:, d] = torch.sum((featL - shift) ** 2, dim=1) cost = cost.unsqueeze(1) # (b,1,D,h,w) x = self.conv3d_1(cost) x = self.conv3d_2(x) x = self.conv3d_3(x) x = self.conv3d_out(x) # (b,1,D,h,w) return x.squeeze(1) # (b,D,h,w) class PSMNet(nn.Module): def __init__(self, max_disp=192): super().__init__() self.max_disp = max_disp self.feature = FeatureExtractor() self.spp = SPP() self.cost_volume = CostVolume(max_disp) def forward(self, left, right): featL = self.spp(self.feature(left)) featR = self.spp(self.feature(right)) cost = self.cost_volume(featL, featR) # (b,D,h,w) # 上采样到原分辨率并softmax cost = F.interpolate(cost, scale_factor=4, mode='trilinear', align_corners=False) prob = F.softmax(-cost, dim=1) disp = torch.sum(prob * torch.arange(self.max_disp, device=left.device).view(1, -1, 1, 1), dim=1) return disp # ---------- 推理与评估 ---------- def load_image(path, size=(320, 240)): img = cv2.imread(path) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = cv2.resize(img, size) return torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0) / 255.0 def evaluate(disp_pred, disp_gt, mask=None): """计算EPE与D1-all（误差>3px或>5%的像素比例）""" if mask is None: mask = disp_gt > 0 epe = torch.mean(torch.abs(disp_pred - disp_gt)[mask]).item() diff = torch.abs(disp_pred - disp_gt) d1 = ((diff > 3) & (diff / (disp_gt + 1e-5) > 0.05))[mask].float().mean().item() return epe, d1 if __name__ == "__main__": # 使用预训练权重（需自行下载，如KITTI2015场景流模型） model = PSMNet(max_disp=192) # 假设权重文件名为psmnet_kitti.pth，包含state_dict # model.load_state_dict(torch.load('psmnet_kitti.pth', map_location='cpu')) model.eval() # 输入图像（示例路径，需替换为实际文件） left_img = load_image('left.png') right_img = load_image('right.png') with torch.no_grad(): disp = model(left_img, right_img) # (1, H, W) # 保存视差图（可视化用） disp_np = disp.squeeze().numpy() disp_viz = (disp_np / disp_np.max() * 255).astype(np.uint8) cv2.imwrite('disp_output.png', disp_viz) # 若有真实视差图，评估精度 # gt = cv2.imread('disp_gt.png', cv2.IMREAD_UNCHANGED).astype(np.float32) / 256.0 # gt_tensor = torch.from_numpy(gt).unsqueeze(0) # epe, d1 = evaluate(disp, gt_tensor) # print(f'EPE: {epe:.3f} px, D1-all: {d1*100:.2f}%') ```

运行与验证

运行方式：将上述代码保存为psmnet_infer.py，准备一对双目图像（如KITTI数据集中的左右视图），并下载PSMNet预训练权重（如官方仓库提供的KITTI模型）。取消注释权重加载行并指定正确路径后执行：

python psmnet_infer.py

验证方法：

  1. 程序运行无报错，生成disp_output.png视差图文件；
  2. 视差图中近处物体（如车辆、行人）亮度高，远处背景亮度低，轮廓清晰可辨；
  3. 若提供真实视差图（KITTI训练集格式，单位为1/256像素），取消评估代码注释，观察输出的EPE与D1-all指标。KITTI 2015验证集上，PSMNet的D1-all通常在2%-3%左右（此为模型公开性能参考，非本案例实测值）。

常见问题：若显存不足（>6GB），可降低输入分辨率（如size=(320, 240)）或减小max_disp至128；权重加载时若键名不匹配，需检查模型定义与预训练权重是否来自同一版本。

← 上一章 📖 返回目录 下一章 →