stereo-algorithms-evolution/python/stereo3_SGM.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
# Stereo Matching using Semi-Global Matching - smoothness cost function approach
# Computes a disparity map from a rectified stereo pair using Semi-Global Matching

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

# 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
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 minimum cost paths 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)

# Compute minimum cost paths for left to right direction
for x in range(cols-1):
    currentCosts = (matchingCosts[:,x,:] + fromLeft[:,x,:])[:,np.newaxis,:]
    fromLeft[:,x+1,:] = computeCosts_BpLike(currentCosts,smoothnessCosts)[:,0,:]

# Compute minimum cost paths for right to left direction
for x in range(cols-1,0,-1):
    currentCosts = (matchingCosts[:,x,:] + fromRight[:,x,:])[:,np.newaxis,:]
    fromRight[:,x-1,:] = computeCosts_BpLike(currentCosts,smoothnessCosts)[:,0,:]

# Compute minimum cost paths for up to down direction
for y in range(rows-1):
    currentCosts = (matchingCosts[y,:,:] + fromUp[y,:,:])[np.newaxis,:,:]
    fromUp[y+1,:,:] = computeCosts_BpLike(currentCosts,smoothnessCosts)[0,:,:]

# Compute minimum cost paths for down to up direction
for y in range(rows-1,0,-1):
    currentCosts = (matchingCosts[y,:,:] + fromDown[y,:,:])[np.newaxis,:,:]
    fromDown[y-1,:,:] = computeCosts_BpLike(currentCosts,smoothnessCosts)[0,:,:]

# Compute total costs
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.imshow(dispImg,cmap="gray")
plt.show(block=False)
plt.pause(0.01)

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

plt.show()