1. import numpy as np
2. import cv2
3. import matplotlib.pyplot as plt
4.  
5. # load images
6. left = cv2.imread('sawtooth/im0.ppm')
7. right = cv2.imread('sawtooth/im2.ppm')
8.  
9. # convert to grayscale
10. left_gray = cv2.cvtColor(left, cv2.COLOR_BGR2GRAY)
11. right_gray = cv2.cvtColor(right, cv2.COLOR_BGR2GRAY)
12.  
13. # show grayscale images
14. cv2.imshow('Left', left_gray)
15. cv2.imshow('Right', right_gray)
16. cv2.waitKey(0)
17. cv2.destroyAllWindows()
18.  
19. # define the image patch
20. patch_loc = [100, 120]
21. patch_size = [100, 100]
22.  
23. # extract the patch from the left image
24. patch_left = left_gray[patch_loc[0]:(patch_loc[0] + patch_size[0]),
25.                        patch_loc[1]:(patch_loc[1] + patch_size[1])]
26.  
27. # extract strip from the rigth image
28. strip_right = right_gray[patch_loc[0]:(patch_loc[0] + patch_size[0]), :]
29.  
30. # show left patch and right strip
31. cv2.imshow('Image Patch Left', patch_left)
32. cv2.imshow('Image Strip Right', strip_right)
33. cv2.waitKey(0)
34. cv2.destroyAllWindows()
35.  
36.  
37. ################################### FIND BEST X USING SUM OF SQUARED DIFFERENCES
38. # create a function to find the best x match location in the right strip
39. def best_x(patch, strip):
40.     """ Return the best x location for the patch in the strip using the sum of
41.    squared differences.
42.    Params:
43.    patch: The input image patch to be compared
44.    strip: The input strip image in which the patch is compared against
45.    Returns:
46.    best_x: The best x fit location of the patch in the strip
47.    """
48.     min_diff = float('inf')
49.     best_x = 0
50.     for x in range(0, strip.shape[1] - patch.shape[1] + 1):
51.         extracted_patch = strip[:, x:(x + patch.shape[1])]
52.         ssd = np.sum((patch - extracted_patch)**2)
53.         if ssd < min_diff:
54.             best_x = x
55.             min_diff = ssd
56.  
57.     return best_x
58.  
59. # get best_x and corresponding patch
60. x = best_x(patch_left, strip_right)
61. patch_right = strip_right[:, x:(x + patch_left.shape[1])]
62.  
63. # plot patch left, strip right and patch right using matplotlib
64. fig = plt.figure(figsize=(8, 6))
65. fig.canvas.set_window_title('Best X Patch Detection')
66. plt.subplot(311), plt.imshow(patch_left, 'gray')
67. plt.title('Original Patch from the Left Image'), plt.xticks([]), plt.yticks([])
68. plt.subplot(312), plt.imshow(strip_right, 'gray')
69. plt.title('Original Strip from the Right Image'), plt.xticks([]), plt.yticks([])
70. plt.subplot(313), plt.imshow(patch_right, 'gray')
71. plt.title('Extracted Patch from the Right Image'), plt.xticks([]), plt.yticks([])
72.  
73. plt.show()
74.  
75.  
76. ############################################# FIND DISPARITY VECTOR FOR 2 STRIPS
77. # define strip row (y) and the square block size (b)
78. y = 75
79. b = 100
80.  
81. # extract strip from the images
82. left_strip = left_gray[y:(y + b), :]
83. right_strip = right_gray[y:(y + b), :]
84.  
85. # show strip images
86. cv2.imshow('Left Disparity Image Strip', left_strip)
87. cv2.imshow('Right Disparity Image Strip', right_strip)
88. cv2.waitKey(0)
89. cv2.destroyAllWindows()
90.  
91. # define function to find disparity
92. def match_strips(left_strip, right_strip, block_size):
93.     ''' A function for finding the pixel disparity vector between each block in
94.        right and left strips with the given block size. '''
95.     num_blocks = left_strip.shape[1] / block_size # python 2 floors integer div
96.     # create an empty disparity vector
97.     disparity = np.zeros(num_blocks)
98.     # how many times does the block fit in the strip
99.     for i in range(num_blocks):
100.         x_left = i*block_size
101.         patch_left = left_strip[:, x_left:(x_left + block_size)] # get patch
102.         x_right = best_x(patch_left, right_strip)
103.         disparity[i] = x_left - x_right
104.  
105.     return disparity
106.  
107. # call the match strips function
108. disparity = match_strips(left_strip, right_strip, b)
109. print disparity
110.  
111. # plot strips and disparity
112. fig = plt.figure(figsize=(8, 6))
113. fig.canvas.set_window_title('Stereo Disparity')
114. plt.subplot(311), plt.imshow(left_strip, 'gray')
115. plt.title('Strip from the Left Image'), plt.xticks([]), plt.yticks([])
116. plt.subplot(312), plt.imshow(right_strip, 'gray')
117. plt.title('Strip from the Right Image'), plt.xticks([]), plt.yticks([])
118. plt.subplot(313), plt.plot(disparity)
119. plt.title('Disparity Between the Two Images')
120.  
121. plt.show()