forked from ndrplz/self-driving-car
-
Notifications
You must be signed in to change notification settings - Fork 0
/
binarization_utils.py
124 lines (88 loc) · 3.64 KB
/
binarization_utils.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
import cv2
import numpy as np
import glob
import matplotlib.pyplot as plt
# selected threshold to highlight yellow lines
yellow_HSV_th_min = np.array([0, 70, 70])
yellow_HSV_th_max = np.array([50, 255, 255])
def thresh_frame_in_HSV(frame, min_values, max_values, verbose=False):
"""
Threshold a color frame in HSV space
"""
HSV = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
min_th_ok = np.all(HSV > min_values, axis=2)
max_th_ok = np.all(HSV < max_values, axis=2)
out = np.logical_and(min_th_ok, max_th_ok)
if verbose:
plt.imshow(out, cmap='gray')
plt.show()
return out
def thresh_frame_sobel(frame, kernel_size):
"""
Apply Sobel edge detection to an input frame, then threshold the result
"""
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
sobel_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=kernel_size)
sobel_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=kernel_size)
sobel_mag = np.sqrt(sobel_x ** 2 + sobel_y ** 2)
sobel_mag = np.uint8(sobel_mag / np.max(sobel_mag) * 255)
_, sobel_mag = cv2.threshold(sobel_mag, 50, 1, cv2.THRESH_BINARY)
return sobel_mag.astype(bool)
def get_binary_from_equalized_grayscale(frame):
"""
Apply histogram equalization to an input frame, threshold it and return the (binary) result.
"""
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
eq_global = cv2.equalizeHist(gray)
_, th = cv2.threshold(eq_global, thresh=250, maxval=255, type=cv2.THRESH_BINARY)
return th
def binarize(img, verbose=False):
"""
Convert an input frame to a binary image which highlight as most as possible the lane-lines.
:param img: input color frame
:param verbose: if True, show intermediate results
:return: binarized frame
"""
h, w = img.shape[:2]
binary = np.zeros(shape=(h, w), dtype=np.uint8)
# highlight yellow lines by threshold in HSV color space
HSV_yellow_mask = thresh_frame_in_HSV(img, yellow_HSV_th_min, yellow_HSV_th_max, verbose=False)
binary = np.logical_or(binary, HSV_yellow_mask)
# highlight white lines by thresholding the equalized frame
eq_white_mask = get_binary_from_equalized_grayscale(img)
binary = np.logical_or(binary, eq_white_mask)
# get Sobel binary mask (thresholded gradients)
sobel_mask = thresh_frame_sobel(img, kernel_size=9)
binary = np.logical_or(binary, sobel_mask)
# apply a light morphology to "fill the gaps" in the binary image
kernel = np.ones((5, 5), np.uint8)
closing = cv2.morphologyEx(binary.astype(np.uint8), cv2.MORPH_CLOSE, kernel)
if verbose:
f, ax = plt.subplots(2, 3)
f.set_facecolor('white')
ax[0, 0].imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
ax[0, 0].set_title('input_frame')
ax[0, 0].set_axis_off()
ax[0, 0].set_axis_bgcolor('red')
ax[0, 1].imshow(eq_white_mask, cmap='gray')
ax[0, 1].set_title('white mask')
ax[0, 1].set_axis_off()
ax[0, 2].imshow(HSV_yellow_mask, cmap='gray')
ax[0, 2].set_title('yellow mask')
ax[0, 2].set_axis_off()
ax[1, 0].imshow(sobel_mask, cmap='gray')
ax[1, 0].set_title('sobel mask')
ax[1, 0].set_axis_off()
ax[1, 1].imshow(binary, cmap='gray')
ax[1, 1].set_title('before closure')
ax[1, 1].set_axis_off()
ax[1, 2].imshow(closing, cmap='gray')
ax[1, 2].set_title('after closure')
ax[1, 2].set_axis_off()
plt.show()
return closing
if __name__ == '__main__':
test_images = glob.glob('test_images/*.jpg')
for test_image in test_images:
img = cv2.imread(test_image)
binarize(img=img, verbose=True)