3. 半全局匹配(SGM)算法实现

任务目标

实现半全局匹配(SGM)算法的核心流程:代价计算、代价聚合(多路径扫描线优化)与视差优化(WTA + 左右一致性检查),对标准立体图像对生成高质量视差图。

环境依赖

Python 3.8+,NumPy,OpenCV(仅用于图像读写与基础处理)。无需GPU,纯CPU可运行。

核心实现

import numpy as np
import cv2

def census_transform(img, win_size=7):
    """中心像素与邻域比较,生成比特串代价"""
    h, w = img.shape
    r = win_size // 2
    cost = np.zeros((h, w), dtype=np.uint64)
    center = cv2.copyMakeBorder(img, r, r, r, r, cv2.BORDER_REPLICATE)
    for dy in range(-r, r+1):
        for dx in range(-r, r+1):
            if dy == 0 and dx == 0:
                continue
            shift = (dy + r) * win_size + (dx + r)
            # 邻域像素大于中心像素则置1
            bit = (center[dy+r:dy+r+h, dx+r:dx+r+w] > img).astype(np.uint64)
            cost |= (bit << shift)
    return cost

def hamming_distance(costL, costR, max_disp):
    """计算每个像素在各视差下的汉明距离代价"""
    h, w = costL.shape
    disp_cost = np.zeros((h, w, max_disp), dtype=np.float32)
    for d in range(max_disp):
        if d == 0:
            disp_cost[:, :, d] = (costL != costR).sum(axis=2) if costL.ndim == 3 else (costL != costR).astype(np.float32)
        else:
            shifted = np.zeros_like(costR)
            shifted[:, d:] = costR[:, :-d]
            disp_cost[:, :, d] = (costL != shifted).astype(np.float32)
    return disp_cost

def cost_aggregation(disp_cost, P1=10, P2=120, num_paths=4):
    """多路径扫描线代价聚合(SGM核心)"""
    h, w, max_disp = disp_cost.shape
    agg_cost = np.zeros_like(disp_cost)
    
    # 定义4条路径方向: 左→右, 右→左, 上→下, 下→上
    paths = [(0, 1), (0, -1), (1, 0), (-1, 0)]
    
    for dy, dx in paths:
        # 初始化路径代价
        path_cost = np.full_like(disp_cost, np.inf)
        # 按路径方向遍历
        if dx == 1:   # 左→右
            for x in range(w):
                _aggregate_step(path_cost, disp_cost, x, 0, dy, dx, P1, P2, h, w, max_disp)
        elif dx == -1: # 右→左
            for x in range(w-1, -1, -1):
                _aggregate_step(path_cost, disp_cost, x, 0, dy, dx, P1, P2, h, w, max_disp)
        elif dy == 1: # 上→下
            for y in range(h):
                _aggregate_step(path_cost, disp_cost, 0, y, dy, dx, P1, P2, h, w, max_disp)
        else:         # 下→上
            for y in range(h-1, -1, -1):
                _aggregate_step(path_cost, disp_cost, 0, y, dy, dx, P1, P2, h, w, max_disp)
        
        agg_cost += path_cost
    
    return agg_cost

def _aggregate_step(path_cost, disp_cost, start_x, start_y, dy, dx, P1, P2, h, w, max_disp):
    """单步聚合:沿路径方向递推"""
    y, x = start_y, start_x
    while 0 <= y < h and 0 <= x < w:
        if y == start_y and x == start_x:
            path_cost[y, x, :] = disp_cost[y, x, :]
        else:
            prev_y, prev_x = y - dy, x - dx
            if 0 <= prev_y < h and 0 <= prev_x < w:
                prev_cost = path_cost[prev_y, prev_x, :]
                # 递推公式: L_r(p,d) = C(p,d) + min(L_r(p-r,d), L_r(p-r,d±1)+P1, min_k L_r(p-r,k)+P2) - min_k L_r(p-r,k)
                min_prev = np.min(prev_cost)
                term1 = prev_cost
                term2 = np.full(max_disp, np.inf)
                term2[1:] = prev_cost[:-1] + P1
                term2[:-1] = np.minimum(term2[:-1], prev_cost[1:] + P1)
                term3 = min_prev + P2
                path_cost[y, x, :] = disp_cost[y, x, :] + np.minimum(np.minimum(term1, term2), term3) - min_prev
        y += dy
        x += dx

def disparity_optimization(agg_cost, max_disp):
    """WTA选择最小代价视差 + 亚像素插值"""
    h, w = agg_cost.shape[:2]
    disp_map = np.zeros((h, w), dtype=np.float32)
    for y in range(h):
        for x in range(w):
            d = np.argmin(agg_cost[y, x, :])
            # 抛物线亚像素插值
            if 0 < d < max_disp-1:
                c0, c1, c2 = agg_cost[y, x, d-1], agg_cost[y, x, d], agg_cost[y, x, d+1]
                denom = (c0 - 2*c1 + c2)
                if abs(denom) > 1e-10:
                    d = d - (c2 - c0) / (2 * denom)
            disp_map[y, x] = d
    return disp_map

def left_right_check(dispL, dispR, max_disp, threshold=1.0):
    """左右一致性检查,剔除遮挡与错误匹配"""
    h, w = dispL.shape
    valid = np.ones((h, w), dtype=np.uint8)
    for y in range(h):
        for x in range(w):
            d = dispL[y, x]
            if d < 0 or d >= max_disp:
                valid[y, x] = 0
                continue
            xr = x - int(round(d))
            if 0 <= xr < w:
                dr = dispR[y, xr]
                if abs(d - dr) > threshold:
                    valid[y, x] = 0
            else:
                valid[y, x] = 0
    return valid

def sgm_pipeline(imgL, imgR, max_disp=64, P1=10, P2=120):
    """SGM完整流程"""
    # 1. 代价计算(Census变换 + 汉明距离)
    costL = census_transform(imgL)
    costR = census_transform(imgR)
    disp_cost = hamming_distance(costL, costR, max_disp)
    
    # 2. 代价聚合(4路径)
    agg_cost = cost_aggregation(disp_cost, P1, P2, num_paths=4)
    
    # 3. 视差优化(WTA)
    dispL = disparity_optimization(agg_cost, max_disp)
    
    # 4. 左右一致性检查(需右图视差图)
    # 为简化演示,此处仅对左图做WTA,右图视差图需交换输入重新计算
    # 实际使用时:dispR = sgm_pipeline(imgR, imgL, ...) 后调用本函数
    return dispL

# ========== 最小可运行示例 ==========
if __name__ == "__main__":
    # 使用OpenCV自带立体图像对(tsukuba)
    imgL = cv2.imread('tsukuba_l.png', cv2.IMREAD_GRAYSCALE)
    imgR = cv2.imread('tsukuba_r.png', cv2.IMREAD_GRAYSCALE)
    if imgL is None or imgR is None:
        print("请提供左右视图灰度图像,例如从Middlebury数据集下载tsukuba图像对")
        exit(1)
    
    # 缩放到合适尺寸加速演示
    scale = 0.5
    imgL = cv2.resize(imgL, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
    imgR = cv2.resize(imgR, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
    
    max_disp = 32
    dispL = sgm_pipeline(imgL, imgR, max_disp=max_disp)
    
    # 归一化显示
    disp_vis = cv2.normalize(dispL, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
    cv2.imwrite('sgm_disparity.png', disp_vis)
    print("视差图已保存为 sgm_disparity.png")

运行与验证

  1. 准备数据:从Middlebury数据集下载tsukuba图像对(或使用任意已校正的立体图像对),将左右灰度图命名为tsukuba_l.png和tsukuba_r.png放在脚本同目录。
  2. 运行:python sgm.py
  3. 通过条件:程序正常结束并输出sgm_disparity.png。观察视差图应满足:
    • 物体轮廓清晰可辨,视差值与深度成反比(近处物体视差大,像素更亮);
    • 平坦区域视差连续,无明显条纹状噪声;
    • 遮挡区域(如物体边缘背景)被剔除显示为黑色或低置信度。

若需定量评估,可下载Middlebury提供的真值视差图,计算误匹配率(误差>1像素视为错误),通常SGM在tsukuba上误匹配率应低于10%。

关键参数说明

  • P1/P2:平滑惩罚项,P1惩罚小视差变化(倾斜面),P2惩罚大视差跳变(深度不连续)。P2通常设为P1的8~24倍。
  • 路径数:4路径已能获得较好效果,8路径(增加对角线方向)可进一步提升精度但耗时翻倍。
  • max_disp:最大搜索视差,需大于场景中最大真实视差。