stereo-algorithms-evolution/python/stereo4_BP1.py at main · aposb/stereo-algorithms-evolution · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
# Stereo Matching using Belief Propagation (Directional) - smoothness cost function approach
# Computes a disparity map from a rectified stereo pair using Belief Propagation (Directional)

import numpy as np
import cv2 as cv
import matplotlib.pyplot as plt
from genericFunctions import *

# Set parameters
dispLevels = 16 #disparity range: 0 to dispLevels-1
lambda_ = 10 #weight of smoothness cost
trunc = 2 #truncation of smoothness cost
iterations = 20

# Define matching cost function
computeMatchingCost = lambda left,right: np.absolute(left-right) #absolute differences

# Define smoothness cost function
computeSmoothnessCost = lambda d1,d2: lambda_*np.minimum(np.absolute(d1-d2),trunc)

# Load left and right images in grayscale
leftImg = cv.imread("left.png",cv.IMREAD_GRAYSCALE)
rightImg = cv.imread("right.png",cv.IMREAD_GRAYSCALE)

# Apply a Gaussian filter
leftImg = cv.GaussianBlur(leftImg,(5,5),0.6)
rightImg = cv.GaussianBlur(rightImg,(5,5),0.6)

# Get the size
(rows,cols) = leftImg.shape

# Convert to int32
leftImg = leftImg.astype(np.int32)
rightImg = rightImg.astype(np.int32)

# Compute pixel-based matching costs (data cost)
matchingCosts = np.zeros((rows,cols,dispLevels),dtype=np.int32)
for d in range(dispLevels):
    rightImgShifted = shiftRight(rightImg,d,0)
    matchingCosts[:,:,d] = computeMatchingCost(leftImg,rightImgShifted)

# Compute smoothness costs
d = np.arange(dispLevels)
smoothnessCosts = computeSmoothnessCost(d,d[np.newaxis,:].T)
smoothnessCosts = smoothnessCosts[np.newaxis,np.newaxis,:,:].astype(np.int32)

# Initialize messages for the 4 directions
fromLeft = np.zeros((rows,cols,dispLevels),dtype=np.int32)
fromRight = np.zeros((rows,cols,dispLevels),dtype=np.int32)
fromUp = np.zeros((rows,cols,dispLevels),dtype=np.int32)
fromDown = np.zeros((rows,cols,dispLevels),dtype=np.int32)

for it in range(iterations):
    # Left to right pass (horizontal forward) - Send messages right
    for x in range(cols-1):
        currentCosts = (matchingCosts[:,x,:] + fromLeft[:,x,:] + fromUp[:,x,:] + fromDown[:,x,:])[:,np.newaxis,:]
        fromLeft[:,x+1,:] = computeCosts_BpLike(currentCosts,smoothnessCosts)[:,0,:]

    # Right to left pass (horizontal backward) - Send messages left
    for x in range(cols-1,0,-1):
        currentCosts = (matchingCosts[:,x,:] + fromRight[:,x,:] + fromUp[:,x,:] + fromDown[:,x,:])[:,np.newaxis,:]
        fromRight[:,x-1,:] = computeCosts_BpLike(currentCosts,smoothnessCosts)[:,0,:]

    # Up to down pass (vertical forward) - Send messages down
    for y in range(rows-1):
        currentCosts = (matchingCosts[y,:,:] + fromUp[y,:,:] + fromLeft[y,:,:] + fromRight[y,:,:])[np.newaxis,:,:]
        fromUp[y+1,:,:] = computeCosts_BpLike(currentCosts,smoothnessCosts)[0,:,:]

    # Down to up pass (vertical backward) - Send messages up
    for y in range(rows-1,0,-1):
        currentCosts = (matchingCosts[y,:,:] + fromDown[y,:,:] + fromLeft[y,:,:] + fromRight[y,:,:])[np.newaxis,:,:]
        fromDown[y-1,:,:] = computeCosts_BpLike(currentCosts,smoothnessCosts)[0,:,:]

    # Compute total costs (belief)
    totalCosts = fromLeft + fromRight + fromUp + fromDown

    # Compute the disparity map
    dispMap = np.argmin(totalCosts,axis=2)

    # Normalize the disparity map for display
    scaleFactor = 256/dispLevels
    dispImg = (dispMap*scaleFactor).astype(np.uint8)

    # Show disparity map
    plt.cla()
    plt.imshow(dispImg,cmap="gray")
    plt.show(block=False)
    plt.pause(0.01)

    # Show iterations
    print("iteration: {0}/{1}".format(it+1,iterations))

# Save disparity map
cv.imwrite("disparity4_BP1.png",dispImg)

plt.show()