-
Notifications
You must be signed in to change notification settings - Fork 2
/
DisparityMap.py
138 lines (98 loc) · 4.03 KB
/
DisparityMap.py
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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
import numpy as np
import cv2
import glob
import matplotlib.pyplot as plt
from tqdm import tqdm
def get_blocks(img,window_size):
height,width = img.shape
window_set=[]
window_loc = []
for y in range(window_size, height - window_size):
for x in range(window_size, width - window_size):
img_window = img[y:y + window_size, x:x + window_size]
window_set.append(img_window)
window_loc.append((y,x))
return window_set,window_loc
def get_boundary_condition(x,y,search_range,width,height):
range_min = max(0, x-search_range)
range_max = min(width,x+search_range )
return range_max,range_min
def get_sdd(imgl_window,imgr_window):
max = 1e10
if imgl_window.shape == imgr_window.shape:
diff = np.abs(imgl_window - imgr_window)
SSD = np.sum(np.square(imgl_window - imgr_window))
return SSD
else:
return max
def get_disparity(img1,img2,search_range,type):
if type==1:
print("method 1")
imgl = img1.copy()
imgr = img2.copy()
disparityImg = np.zeros(img2.shape)
window_size = 21
#get a patch in left image
height,width = img1.shape
for y in tqdm(range(window_size,height-window_size)):
for x in range(window_size,width-window_size):
#get window boundary conditions
# range_max,range_min = get_boundary_condition(x,y,50,width,height)
range_min = max(0, x-100)
range_max = min(width,x+100)
imgl_window = imgl[y:y+window_size, x:x+window_size]
first = True
min_val = None
index = None
#search for corresponding patch in right image
for i in range(range_min,range_max):
imgr_window = imgr[y:y+window_size, i:i+window_size]
SSD = get_sdd(imgl_window,imgr_window)
if first:
min_val = SSD
index = i
first = False
else:
if SSD < min_val:
min_val = SSD
index = i
d = np.abs(index - x)
disparityImg[y,x] = d
return disparityImg
if type == 2:
print("method 2")
# img1 = cv2.resize(img1, (int(img1.shape[1]*0.5), int(img1.shape[0] *0.5)))
# img2 = cv2.resize(img2, (int(img2.shape[1]*0.5), int(img2.shape[0] *0.5)))
imgl = img1.copy()
imgr = img2.copy()
disparityImg = np.zeros(img2.shape)
window_size = 15
#get a patch in left image
height,width = img1.shape
imgl = imgl.astype(np.int32)
imgr = imgr.astype(np.int32)
for y in tqdm(range(0,height-window_size)):
imgl_window_set = []
imgr_window_set = []
for x in range(0,width-window_size):
left_window = imgl[y:y+window_size, x:x+window_size]
right_window = imgr[y:y+window_size, x:x+window_size]
imgl_window_set.append(left_window.flatten())
imgr_window_set.append(right_window.flatten())
imgl_window_set = np.array(imgl_window_set)
imgr_window_set = np.array(imgr_window_set)
print(len(imgl_window_set))
for i in range(len(imgl_window_set)):
SSD = np.sum(np.square(imgr_window_set- imgl_window_set[i]),axis=1)
index = np.argmin(SSD)
disparity = np.abs(i - index)
# disparity = ((disparity+1) / 2) * 255
disparityImg[y,i] = np.uint8(disparity)
return disparityImg
else:
return np.zeros(img2.shape)
def get_depth(disparityImg,b,f):
depth = (b * f) / (disparityImg + 1e-10)
depth[depth > 10000] = 10000
depth_map = np.uint8(depth * 255 / np.max(depth))
return depth_map