计算机视觉（6）双目相机基于SGM（semi-global matching）算法获取深度信息

摘要：
最近在做双目视差估计算法，在OpenCV里有一些算法，其中半全局块匹配（Semi-Global Block Matching，SGBM）算法具有视差效果好速度快的特点，因此常常被广泛应用。本文主要讨论的就是SGBM算法。
SGM聚合步骤示意图（视差图呈现）：
  1 import argparse
  2 import sys
  3 import time as t
  4 
  5 import cv2
  6 import numpy as np
  7 
  8 
  9 class Direction:
 10     def __init__(self, direction=(0, 0), name='invalid'):
 11         """
 12         represent a cardinal direction in image coordinates (top left = (0, 0) and bottom right = (1, 1)).
 13         :param direction: (x, y) for cardinal direction.
 14         :param name: common name of said direction.
 15         """
 16         self.direction = direction
 17         self.name = name
 18 
 19 
 20 # 8 defined directions for sgm
 21 N = Direction(direction=(0, -1), name='north')
 22 NE = Direction(direction=(1, -1), name='north-east')
 23 E = Direction(direction=(1, 0), name='east')
 24 SE = Direction(direction=(1, 1), name='south-east')
 25 S = Direction(direction=(0, 1), name='south')
 26 SW = Direction(direction=(-1, 1), name='south-west')
 27 W = Direction(direction=(-1, 0), name='west')
 28 NW = Direction(direction=(-1, -1), name='north-west')
 29 
 30 
 31 class Paths:
 32     def __init__(self):
 33         """
 34         represent the relation between the directions.
 35         """
 36         self.paths = [N, NE, E, SE, S, SW, W, NW]
 37         self.size = len(self.paths)
 38         self.effective_paths = [(E,  W), (SE, NW), (S, N), (SW, NE)]
 39 
 40 
 41 class Parameters:
 42     def __init__(self, max_disparity=128, P1=10, P2=120, csize=(7, 7), bsize=(3, 3)):
 43         """
 44         represent all parameters used in the sgm algorithm.
 45         :param max_disparity: maximum distance between the same pixel in both images.
 46         :param P1: penalty for disparity difference = 1
 47         :param P2: penalty for disparity difference > 1
 48         :param csize: size of the kernel for the census transform.
 49         :param bsize: size of the kernel for blurring the images and median filtering.
 50         """
 51         self.max_disparity = max_disparity
 52         self.P1 = P1
 53         self.P2 = P2
 54         self.csize = csize
 55         self.bsize = bsize
 56 
 57 
 58 def load_images(left_name, right_name, parameters):
 59     """
 60     read and blur stereo image pair.
 61     :param left_name: name of the left image.
 62     :param right_name: name of the right image.
 63     :param parameters: structure containing parameters of the algorithm.
 64     :return: blurred left and right images.
 65     """
 66     left = cv2.imread(left_name, 0)
 67     left = cv2.GaussianBlur(left, parameters.bsize, 0, 0)
 68     right = cv2.imread(right_name, 0)
 69     right = cv2.GaussianBlur(right, parameters.bsize, 0, 0)
 70     return left, right
 71 
 72 
 73 def get_indices(offset, dim, direction, height):
 74     """
 75     for the diagonal directions (SE, SW, NW, NE), return the array of indices for the current slice.
 76     :param offset: difference with the main diagonal of the cost volume.
 77     :param dim: number of elements along the path.
 78     :param direction: current aggregation direction.
 79     :param height: H of the cost volume.
 80     :return: arrays for the y (H dimension) and x (W dimension) indices.
 81     """
 82     y_indices = []
 83     x_indices = []
 84 
 85     for i in range(0, dim):
 86         if direction == SE.direction:
 87             if offset < 0:
 88                 y_indices.append(-offset + i)
 89                 x_indices.append(0 + i)
 90             else:
 91                 y_indices.append(0 + i)
 92                 x_indices.append(offset + i)
 93 
 94         if direction == SW.direction:
 95             if offset < 0:
 96                 y_indices.append(height + offset - i)
 97                 x_indices.append(0 + i)
 98             else:
 99                 y_indices.append(height - i)
100                 x_indices.append(offset + i)
101 
102     return np.array(y_indices), np.array(x_indices)
103 
104 
105 def get_path_cost(slice, offset, parameters):
106     """
107     part of the aggregation step, finds the minimum costs in a D x M slice (where M = the number of pixels in the
108     given direction)
109     :param slice: M x D array from the cost volume.
110     :param offset: ignore the pixels on the border.
111     :param parameters: structure containing parameters of the algorithm.
112     :return: M x D array of the minimum costs for a given slice in a given direction.
113     """
114     other_dim = slice.shape[0]
115     disparity_dim = slice.shape[1]
116 
117     disparities = [d for d in range(disparity_dim)] * disparity_dim
118     disparities = np.array(disparities).reshape(disparity_dim, disparity_dim)
119 
120     penalties = np.zeros(shape=(disparity_dim, disparity_dim), dtype=slice.dtype)
121     penalties[np.abs(disparities - disparities.T) == 1] = parameters.P1
122     penalties[np.abs(disparities - disparities.T) > 1] = parameters.P2
123 
124     minimum_cost_path = np.zeros(shape=(other_dim, disparity_dim), dtype=slice.dtype)
125     minimum_cost_path[offset - 1, :] = slice[offset - 1, :]
126 
127     for i in range(offset, other_dim):
128         previous_cost = minimum_cost_path[i - 1, :]
129         current_cost = slice[i, :]
130         costs = np.repeat(previous_cost, repeats=disparity_dim, axis=0).reshape(disparity_dim, disparity_dim)
131         costs = np.amin(costs + penalties, axis=0)
132         minimum_cost_path[i, :] = current_cost + costs - np.amin(previous_cost)
133     return minimum_cost_path
134 
135 
136 def aggregate_costs(cost_volume, parameters, paths):
137     """
138     sgm第二步，在N个可能方向匹配
139     second step of the sgm algorithm, aggregates matching costs for N possible directions (8 in this case).
140     :param cost_volume: array containing the matching costs.
141     :param parameters: structure containing parameters of the algorithm.
142     :param paths: structure containing all directions in which to aggregate costs.
143     :return: H x W x D x N array of matching cost for all defined directions.
144     """
145     height = cost_volume.shape[0]
146     width = cost_volume.shape[1]
147     disparities = cost_volume.shape[2]
148     start = -(height - 1)
149     end = width - 1
150 
151     aggregation_volume = np.zeros(shape=(height, width, disparities, paths.size), dtype=cost_volume.dtype)
152 
153     path_id = 0
154     for path in paths.effective_paths:
155         print('\tProcessing paths {} and {}...'.format(path[0].name, path[1].name), end='')
156         sys.stdout.flush()
157         dawn = t.time()
158 
159         main_aggregation = np.zeros(shape=(height, width, disparities), dtype=cost_volume.dtype)
160         opposite_aggregation = np.copy(main_aggregation)
161 
162         main = path[0]
163         if main.direction == S.direction:
164             for x in range(0, width):
165                 south = cost_volume[0:height, x, :]
166                 north = np.flip(south, axis=0)
167                 main_aggregation[:, x, :] = get_path_cost(south, 1, parameters)
168                 opposite_aggregation[:, x, :] = np.flip(get_path_cost(north, 1, parameters), axis=0)
169 
170         if main.direction == E.direction:
171             for y in range(0, height):
172                 east = cost_volume[y, 0:width, :]
173                 west = np.flip(east, axis=0)
174                 main_aggregation[y, :, :] = get_path_cost(east, 1, parameters)
175                 opposite_aggregation[y, :, :] = np.flip(get_path_cost(west, 1, parameters), axis=0)
176 
177         if main.direction == SE.direction:
178             for offset in range(start, end):
179                 south_east = cost_volume.diagonal(offset=offset).T
180                 north_west = np.flip(south_east, axis=0)
181                 dim = south_east.shape[0]
182                 y_se_idx, x_se_idx = get_indices(offset, dim, SE.direction, None)
183                 y_nw_idx = np.flip(y_se_idx, axis=0)
184                 x_nw_idx = np.flip(x_se_idx, axis=0)
185                 main_aggregation[y_se_idx, x_se_idx, :] = get_path_cost(south_east, 1, parameters)
186                 opposite_aggregation[y_nw_idx, x_nw_idx, :] = get_path_cost(north_west, 1, parameters)
187 
188         if main.direction == SW.direction:
189             for offset in range(start, end):
190                 south_west = np.flipud(cost_volume).diagonal(offset=offset).T
191                 north_east = np.flip(south_west, axis=0)
192                 dim = south_west.shape[0]
193                 y_sw_idx, x_sw_idx = get_indices(offset, dim, SW.direction, height - 1)
194                 y_ne_idx = np.flip(y_sw_idx, axis=0)
195                 x_ne_idx = np.flip(x_sw_idx, axis=0)
196                 main_aggregation[y_sw_idx, x_sw_idx, :] = get_path_cost(south_west, 1, parameters)
197                 opposite_aggregation[y_ne_idx, x_ne_idx, :] = get_path_cost(north_east, 1, parameters)
198 
199         aggregation_volume[:, :, :, path_id] = main_aggregation
200         aggregation_volume[:, :, :, path_id + 1] = opposite_aggregation
201         path_id = path_id + 2
202 
203         dusk = t.time()
204         print('\t(done in {:.2f}s)'.format(dusk - dawn))
205 
206     return aggregation_volume
207 
208 
209 def compute_costs(left, right, parameters, save_images):
210     """
211     SGM第一步，基于人口普查算法和汉明距离
212     first step of the sgm algorithm, matching cost based on census transform and hamming distance.
213     :param left: left image.
214     :param right: right image.
215     :param parameters: structure containing parameters of the algorithm.
216     :param save_images: whether to save census images or not.
217     :return: H x W x D array with the matching costs.
218     """
219     assert left.shape[0] == right.shape[0] and left.shape[1] == right.shape[1], 'left & right must have the same shape.'
220     assert parameters.max_disparity > 0, 'maximum disparity must be greater than 0.'
221 
222     height = left.shape[0]
223     width = left.shape[1]
224     cheight = parameters.csize[0]
225     cwidth = parameters.csize[1]
226     y_offset = int(cheight / 2)
227     x_offset = int(cwidth / 2)
228     disparity = parameters.max_disparity
229 
230     left_img_census = np.zeros(shape=(height, width), dtype=np.uint8)
231     right_img_census = np.zeros(shape=(height, width), dtype=np.uint8)
232     left_census_values = np.zeros(shape=(height, width), dtype=np.uint64)
233     right_census_values = np.zeros(shape=(height, width), dtype=np.uint64)
234 
235     print('\tComputing left and right census...', end='')
236     sys.stdout.flush()
237     dawn = t.time()
238     # pixels on the border will have no census values
239     for y in range(y_offset, height - y_offset):
240         for x in range(x_offset, width - x_offset):
241             left_census = np.int64(0)
242             center_pixel = left[y, x]
243             reference = np.full(shape=(cheight, cwidth), fill_value=center_pixel, dtype=np.int64)
244             image = left[(y - y_offset):(y + y_offset + 1), (x - x_offset):(x + x_offset + 1)]
245             comparison = image - reference
246             for j in range(comparison.shape[0]):
247                 for i in range(comparison.shape[1]):
248                     if (i, j) != (y_offset, x_offset):
249                         left_census = left_census << 1
250                         if comparison[j, i] < 0:
251                             bit = 1
252                         else:
253                             bit = 0
254                         left_census = left_census | bit
255             left_img_census[y, x] = np.uint8(left_census)
256             left_census_values[y, x] = left_census
257 
258             right_census = np.int64(0)
259             center_pixel = right[y, x]
260             reference = np.full(shape=(cheight, cwidth), fill_value=center_pixel, dtype=np.int64)
261             image = right[(y - y_offset):(y + y_offset + 1), (x - x_offset):(x + x_offset + 1)]
262             comparison = image - reference
263             for j in range(comparison.shape[0]):
264                 for i in range(comparison.shape[1]):
265                     if (i, j) != (y_offset, x_offset):
266                         right_census = right_census << 1
267                         if comparison[j, i] < 0:
268                             bit = 1
269                         else:
270                             bit = 0
271                         right_census = right_census | bit
272             right_img_census[y, x] = np.uint8(right_census)
273             right_census_values[y, x] = right_census
274 
275     dusk = t.time()
276     print('\t(done in {:.2f}s)'.format(dusk - dawn))
277 
278     if save_images:
279         cv2.imwrite('left_census.png', left_img_census)
280         cv2.imwrite('right_census.png', right_img_census)
281 
282     print('\tComputing cost volumes...', end='')
283     sys.stdout.flush()
284     dawn = t.time()
285     left_cost_volume = np.zeros(shape=(height, width, disparity), dtype=np.uint32)
286     right_cost_volume = np.zeros(shape=(height, width, disparity), dtype=np.uint32)
287     lcensus = np.zeros(shape=(height, width), dtype=np.int64)
288     rcensus = np.zeros(shape=(height, width), dtype=np.int64)
289     for d in range(0, disparity):
290         rcensus[:, (x_offset + d):(width - x_offset)] = right_census_values[:, x_offset:(width - d - x_offset)]
291         left_xor = np.int64(np.bitwise_xor(np.int64(left_census_values), rcensus))
292         left_distance = np.zeros(shape=(height, width), dtype=np.uint32)
293         while not np.all(left_xor == 0):
294             tmp = left_xor - 1
295             mask = left_xor != 0
296             left_xor[mask] = np.bitwise_and(left_xor[mask], tmp[mask])
297             left_distance[mask] = left_distance[mask] + 1
298         left_cost_volume[:, :, d] = left_distance
299 
300         lcensus[:, x_offset:(width - d - x_offset)] = left_census_values[:, (x_offset + d):(width - x_offset)]
301         right_xor = np.int64(np.bitwise_xor(np.int64(right_census_values), lcensus))
302         right_distance = np.zeros(shape=(height, width), dtype=np.uint32)
303         while not np.all(right_xor == 0):
304             tmp = right_xor - 1
305             mask = right_xor != 0
306             right_xor[mask] = np.bitwise_and(right_xor[mask], tmp[mask])
307             right_distance[mask] = right_distance[mask] + 1
308         right_cost_volume[:, :, d] = right_distance
309 
310     dusk = t.time()
311     print('\t(done in {:.2f}s)'.format(dusk - dawn))
312 
313     return left_cost_volume, right_cost_volume
314 
315 
316 def select_disparity(aggregation_volume):
317     """
318     最后一步
319     last step of the sgm algorithm, corresponding to equation 14 followed by winner-takes-all approach.
320     :param aggregation_volume: H x W x D x N array of matching cost for all defined directions.
321     :return: disparity image.
322     """
323     volume = np.sum(aggregation_volume, axis=3)
324     disparity_map = np.argmin(volume, axis=2)
325     return disparity_map
326 
327 
328 def normalize(volume, parameters):
329     """
330     transforms values from the range (0, 64) to (0, 255).
331     :param volume: n dimension array to normalize.
332     :param parameters: structure containing parameters of the algorithm.
333     :return: normalized array.
334     """
335     return 255.0 * volume / parameters.max_disparity
336     #return 255.0 * volume / 128
337 
338 
339 
340 def get_recall(disparity, gt, args):
341     """
342     computes the recall of the disparity map.
343     :param disparity: disparity image.
344     :param gt: path to ground-truth image.
345     :param args: program arguments.
346     :return: rate of correct predictions.
347     """
348     gt = np.float32(cv2.imread(gt, cv2.IMREAD_GRAYSCALE))
349     gt = np.int16(gt / 255.0 * float(args.disp))
350     disparity = np.int16(np.float32(disparity) / 255.0 * float(args.disp))
351     correct = np.count_nonzero(np.abs(disparity - gt) <= 3)
352     return float(correct) / gt.size
353 
354 
355 def sgm():
356     """
357     main function applying the semi-global matching algorithm.
358     :return: void.
359     """
360     parser = argparse.ArgumentParser()
361     parser.add_argument('--left', default='cones/im2.png', help='name (path) to the left image')
362     parser.add_argument('--right', default='cones/im6.png', help='name (path) to the right image')
363     parser.add_argument('--left_gt', default='cones/disp2.png', help='name (path) to the left ground-truth image')
364     parser.add_argument('--right_gt', default='cones/disp6.png', help='name (path) to the right ground-truth image')
365     parser.add_argument('--output', default='disparity_map.png', help='name of the output image')
366     parser.add_argument('--disp', default=33, type=int, help='maximum disparity for the stereo pair')
367     parser.add_argument('--images', default=True, type=bool, help='save intermediate representations')
368     parser.add_argument('--eval', default=True, type=bool, help='evaluate disparity map with 3 pixel error')
369     args = parser.parse_args()
370 
371     left_name = args.left
372     right_name = args.right
373     left_gt_name = args.left_gt
374     right_gt_name = args.right_gt
375     output_name = args.output
376     disparity = args.disp
377     save_images = args.images
378     evaluation = args.eval
379 
380     dawn = t.time()
381 
382     parameters = Parameters(max_disparity=disparity, P1=10, P2=120, csize=(7, 7), bsize=(3, 3))
383     paths = Paths()
384 
385     print('\nLoading images...')
386     left, right = load_images(left_name, right_name, parameters)
387 
388     #第一步
389     print('\nStarting cost computation...')
390     left_cost_volume, right_cost_volume = compute_costs(left, right, parameters, save_images)
391     if save_images:
392         left_disparity_map = np.uint8(normalize(np.argmin(left_cost_volume, axis=2), parameters))
393         cv2.imwrite('disp_map_left_cost_volume.png', left_disparity_map)
394         right_disparity_map = np.uint8(normalize(np.argmin(right_cost_volume, axis=2), parameters))
395         cv2.imwrite('disp_map_right_cost_volume.png', right_disparity_map)
396     #第二步
397     print('\nStarting left aggregation computation...')
398     left_aggregation_volume = aggregate_costs(left_cost_volume, parameters, paths)
399     print('\nStarting right aggregation computation...')
400     right_aggregation_volume = aggregate_costs(right_cost_volume, parameters, paths)
401 
402     print('\nSelecting best disparities...')
403     left_disparity_map = np.uint8(normalize(select_disparity(left_aggregation_volume), parameters))
404     right_disparity_map = np.uint8(normalize(select_disparity(right_aggregation_volume), parameters))
405     if save_images:
406         cv2.imwrite('left_disp_map_no_post_processing.png', left_disparity_map)
407         cv2.imwrite('right_disp_map_no_post_processing.png', right_disparity_map)
408 
409     print('\nApplying median filter...')
410     #中值滤波
411     left_disparity_map = cv2.medianBlur(left_disparity_map, parameters.bsize[0])
412     right_disparity_map = cv2.medianBlur(right_disparity_map, parameters.bsize[0])
413     # right_disparity_map = cv2.GaussianBlur(right_disparity_map,(5,5), 15)
414     # left_disparity_map = cv2.GaussianBlur(left_disparity_map,(5,5), 15)
415     cv2.imwrite(f'left_{output_name}', left_disparity_map)
416     cv2.imwrite(f'right_{output_name}', right_disparity_map)
417 
418     if evaluation:
419         print('\nEvaluating left disparity map...')
420         recall = get_recall(left_disparity_map, left_gt_name, args)
421         print('\tRecall = {:.2f}%'.format(recall * 100.0))
422         print('\nEvaluating right disparity map...')
423         recall = get_recall(right_disparity_map, right_gt_name, args)
424         print('\tRecall = {:.2f}%'.format(recall * 100.0))
425 
426     dusk = t.time()
427     print('\nFin.')
428     print('\nTotal execution time = {:.2f}s'.format(dusk - dawn))
429 
430 
431 if __name__ == '__main__':
432     sgm()