Cost Volume Visualization (for stereo depth estimation debuging and exploration)

cost_volume_viz.py

import numpy as np
import matplotlib.pyplot as plt
from volume_viz import VolumeViz


class CostVolumeViz:
    def __init__(self, left, right, disparity):
        self.left = left
        self.right = right
        self.disparity = self.StandardizeDisparity(disparity)
        self.cost_volume = None

    def Plot(self, title=''):
        assert self.cost_volume is not None
        fig = plt.figure(title)
        VolumeViz(fig, image=L, volume=self.cost_volume, marker=self.disparity)
        plt.show()

    def CalcCostVolume(self, left_features, right_features, normalize=True):
        rows, cols, _ = left_features.shape
        CV = np.zeros((rows, cols, cols))
        for r in range(rows):
            CV[r, :, :] = np.dot(left_features[r, :, :], right_features[r, :, :].transpose())

        if normalize:
            norm = CV.max(axis=2, keepdims=True)
            CV /= norm

        print('Cost Volume Shape = {}'.format(CV.shape))

        return CV

    def CalcVolumeFromRawFeatures(self, block_size=5):
        LF = self.CalcRawFeatures(self.left, block_size=block_size)
        RF = self.CalcRawFeatures(self.right, block_size=block_size)
        self.cost_volume = self.CalcCostVolume(LF, RF)

    #   left_feat/right_feat should be of shape [height, width, features]
    def CalcVolumeFromCustomFeatures(self, left_feat, right_feat):
        self.cost_volume = self.CalcCostVolume(left_feat, right_feat)

    #   volume should be of shape [height, width, width(disparity)]
    def LoadCostVolume(self, volume):
        self.cost_volume = volume

    def CalcRawFeatures(self, image, block_size = 5):
        cols, rows = image.shape
        padded = np.pad(image, block_size // 2, mode='edge')
        features = np.zeros((cols, rows, block_size**2))
        for c in range(cols):
            for r in range(rows):
                f = padded[c:c+block_size, r:r+block_size].flatten()
                norm = np.sqrt(f.dot(f) + 1e-6)
                features[c, r, :] = f / norm

        return features

    def StandardizeDisparity(self, disparity):
        std_disp = disparity // 16
        shift = np.repeat(np.expand_dims(np.arange(0, disparity.shape[1]), axis=0), axis=0, repeats=disparity.shape[0])
        std_disp = shift - std_disp
        return std_disp

main.py

import numpy as np
import matplotlib.pyplot as plt
import cv2
from cost_volume_viz import CostVolumeViz
from utils import read_disparity_file, calc_features

# Load Data
L = cv2.imread('data/Cor_00008_L.png').mean(axis=2)
L = cv2.resize(L, None, fx=0.5, fy=0.5)
R = cv2.imread('data/Cor_00008_R.png').mean(axis=2)
R = cv2.resize(R, None, fx=0.5, fy=0.5)
D, nonValid = read_disparity_file('data/LowResDisp_00008.bin', shift=False)

cost_volume_viz = CostVolumeViz(left=L, right=R, disparity=D)

# METHOD 1
#'''
method = 'CalcVolumeFromRawFeatures'
cost_volume_viz.CalcVolumeFromRawFeatures(block_size=21)
#'''

# METHOD 2
'''
method = 'CalcVolumeFromCustomFeatures'
feat_L, feat_R = calc_features(L), calc_features(R)
cost_volume_viz.CalcVolumeFromCustomFeatures(left_feat=feat_L, right_feat=feat_R)
'''

# METHOD 3
'''
method = 'LoadCostVolume'
cost_volume_viz.LoadCostVolume(volume=...)
'''

cost_volume_viz.Plot(title=method)

utils.py

import numpy as np
import cv2


def calc_features(im, radius = 7):
    FeatureXtractor = cv2.ORB_create()
    # FeatureXtractor = cv2.xfeatures2d.SURF_create()
    kp = [cv2.KeyPoint(x, y, radius) for y in range(0, im.shape[1]) for x in range(0, im.shape[0])]
    kp, feat_raw = FeatureXtractor.compute(im.astype('uint8'), kp)
    feat = np.zeros((im.shape[0], im.shape[1], feat_raw.shape[1]))
    for k, f in zip(kp, feat_raw):
        feat[int(k.pt[0] - 1), int(k.pt[1] - 1), :] = f
    return feat
  
  
def read_disparity_file(filename, shift=True, verbose=True):
  def next_int(file):
      byte = file.read(4)
      val = int.from_bytes(byte, byteorder='little', signed=True)
      return val

  with open(filename, "rb") as f:
      rows = next_int(f)
      cols = next_int(f)
      num_type = next_int(f)
      if verbose:
          print("rows: {}, cols: {}, type: {}".format(rows, cols, num_type))

      if num_type == 5:
          depth = np.reshape(np.fromfile(f, dtype='<f'), [rows, cols])
          depth[depth < 0] = -1
      elif num_type == 3:
          lowResInvLow = next_int(f)
          print('lowResInvLow: {}'.format(lowResInvLow))
          depth = np.reshape(np.fromfile(f, dtype='<i2'), [rows, cols])
          if shift:
              depth -= lowResInvLow
      else:
          assert False

      depth_range = np.min(depth), np.max(depth)
      if verbose:
          print("loaded size {}, with range {}".format(depth.shape, depth_range))

  return depth, lowResInvLow

volume_viz.py

import numpy as np
import matplotlib.pyplot as plt


class VolumeViz:
    def __init__(self, fig, image, volume, marker, verbose=True):
        self.verbose = verbose
        self.image = image
        self.volume = volume
        self.marker = marker
        assert image.shape[0] == volume.shape[0]
        assert image.shape[1] == volume.shape[1]
        self.figure = fig
        # Axis #1
        self.ax1 = fig.add_subplot(121)
        self.ax1.set_title('Select Point (CLICK)')
        self.ax1.imshow(self.image)
        # Axis #2
        self.ax2 = fig.add_subplot(122)

        self.cid = self.figure.canvas.mpl_connect('button_press_event', self)
        self.figure.canvas.draw()

    def __call__(self, event):
        if event.inaxes!=self.ax1:
            return
        if self.verbose:
            print('%s click: button=%d, x=%d, y=%d, xdata=%f, ydata=%f' % ('double' if event.dblclick else 'single', event.button, event.x, event.y, event.xdata, event.ydata))
        # Axis #2
        self.ax2.cla()
        self.ax2.plot(self.volume[np.round(event.ydata).astype('int'), np.round(event.xdata).astype('int'), :])
        self.ax2.scatter(self.marker[np.round(event.ydata).astype('int'), np.round(event.xdata).astype('int')], np.array([1]), color='g')
        self.ax2.set_title('CostVolume Projected & Expected Disparity')
        self.figure.canvas.draw()