ssd.py

#!/usr/bin/env python
# --------------------------------------------------------------------
# Simple sum of squared differences (SSD) stereo-matching using Numpy
# --------------------------------------------------------------------

# Copyright (c) 2016 David Christian
# Licensed under the MIT License
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from PIL import Image

def match(x,y,left, right, max_offset, kernel_half):
    disparity = 0
    best_offset = 0
    prev_ssd = 65534
    pathl = left[y-kernel_half:y+kernel_half, x-kernel_half:x+kernel_half]
    w = left.shape[1]    
    for offset in range(max_offset):
        
        if (x-offset)-kernel_half < 0 or (x- offset)+kernel_half >=  w-kernel_half:
            continue           
        ssd = 0                         
        pathr = right[y-kernel_half:y+kernel_half, (x-offset)-kernel_half:(x- offset)+kernel_half]
        # v and u are the x,y of our local window search, used to ensure a good match
        # because the squared differences of two pixels alone is not enough ot go on
        ssd = np.sum((pathl - pathr)**2)
        # for v in range(-kernel_half, kernel_half):
        #     for u in range(-kernel_half, kernel_half):
        #         # iteratively sum the sum of squared differences value for this block
        #         # left[] and right[] are arrays of uint8, so converting them to int saves
        #         # potential overflow
        #         ssd_temp = int(left[y+v, x+u]) - int(right[y+v, (x+u) - offset])  
        #         ssd += ssd_temp * ssd_temp              
        
        # if this value is smaller than the previous ssd at this block
        # then it's theoretically a closer match. Store this value against
        # this block..

        if ssd < prev_ssd:
            prev_ssd = ssd
            disparity = offset

    return disparity    

def matchAndPlot(x,y,left, right, max_offset, kernel_half, title):
    disparity = 0
    best_offset = 0
    prev_ssd = 65534
    pathl = left[y-kernel_half:y+kernel_half, x-kernel_half:x+kernel_half]
    w = left.shape[1]    
    costs = []
    for offset in range(max_offset):
        
        if (x-offset)-kernel_half < 0 or (x- offset)+kernel_half >=  w-kernel_half:
            continue           
        ssd = 0                         
        pathr = right[y-kernel_half:y+kernel_half, (x-offset)-kernel_half:(x- offset)+kernel_half]
        # v and u are the x,y of our local window search, used to ensure a good match
        # because the squared differences of two pixels alone is not enough ot go on
        ssd = np.sum((pathl - pathr)**2)
        # for v in range(-kernel_half, kernel_half):
        #     for u in range(-kernel_half, kernel_half):
        #         # iteratively sum the sum of squared differences value for this block
        #         # left[] and right[] are arrays of uint8, so converting them to int saves
        #         # potential overflow
        #         ssd_temp = int(left[y+v, x+u]) - int(right[y+v, (x+u) - offset])  
        #         ssd += ssd_temp * ssd_temp              
        
        # if this value is smaller than the previous ssd at this block
        # then it's theoretically a closer match. Store this value against
        # this block..
        costs.append(ssd)
        if ssd < prev_ssd:
            prev_ssd = ssd
            
            disparity = offset
    fig, ax = plt.subplots(1,2,figsize=(20, 5), gridspec_kw={'width_ratios': [0.4, 0.6]})
    fig.suptitle(title)
    ax[0].imshow(left)
    rect = patches.Rectangle((y-kernel_half, y+kernel_half), 
                            kernel_half*2, 
                            kernel_half*2, 
                            linewidth=1, edgecolor='r', facecolor='none')
    ax[0].imshow(left.astype(np.uint8), cmap='gray')
    ax[0].title.set_text("Left Image")
    ax[0].add_patch(rect)
    ax[0].set_xticks([])
    ax[0].set_yticks([])
    ax[1].title.set_text("Costs")
    ax[1].set_xlabel("Disparity")
    ax[1].set_ylabel("SSD Cost")
    ax[1].plot(costs)
    ax[1].set_ylim(0, 5000)
    ax[1].grid()
    plt.tight_layout()
    plt.savefig(title + ".png")
    plt.show()
    
    return disparity          
             
def stereo_match(left_img, right_img, kernel, max_offset):
    # Load in both images, assumed to be RGBA 8bit per channel images
    left_img = Image.open(left_img).convert('L')
    left = np.asarray(left_img)
    right_img = Image.open(right_img).convert('L')
    right = np.asarray(right_img)    
    w, h = left_img.size  # assume that both images are same size   
    
    # Depth (or disparity) map
    depth = np.zeros((w, h), np.uint8)
    depth.shape = h, w
       
    kernel_half = int(kernel / 2)
    offset_adjust = 255 / max_offset  # this is used to map depth map output to 0-255 range
      
    matchAndPlot(100,214, left, right, max_offset, kernel_half, "Strong feature")
    matchAndPlot(750,450, left, right, max_offset, kernel_half, "Strong feature, 2")
    matchAndPlot(600,800, left, right, max_offset, kernel_half, "Weak feature")
    
    for y in range(kernel_half+1, h - kernel_half-1):              
        print("\rProcessing.. %d%% complete"%(y / (h - kernel_half) * 100), end="", flush=True)        
        for x in range(kernel_half+1, w - kernel_half-1):
            disparity = match(x,y,left, right, max_offset, kernel_half)
            
            depth[y, x] = disparity * offset_adjust
                                
    # Convert to PIL and save it
    Image.fromarray(depth).save('depth_no_pattern.png')

if __name__ == '__main__':
    stereo_match("imgL.png", "imgR.png", 6, 30)  # 6x6 local search kernel, 30 pixel search range