Permalink
Cannot retrieve contributors at this time
import numpy as np | |
import cv2 | |
def gradient_abs_value_mask(image, sobel_kernel=3, axis='x', threshold=(0, 255)): | |
""" | |
Masks the image based on gradient absolute value. | |
Parameters | |
---------- | |
image : Image to mask. | |
sobel_kernel : Kernel of the Sobel gradient operation. | |
axis : Axis of the gradient, 'x' or 'y'. | |
threshold : Value threshold for it to make it to appear in the mask. | |
Returns | |
------- | |
Image mask with 1s in activations and 0 in other pixels. | |
""" | |
# Take the absolute value of derivative in x or y given orient = 'x' or 'y' | |
if axis == 'x': | |
sobel = np.absolute(cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=sobel_kernel)) | |
if axis == 'y': | |
sobel = np.absolute(cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=sobel_kernel)) | |
# Scale to 8-bit (0 - 255) then convert to type = np.uint8 | |
sobel = np.uint8(255 * sobel / np.max(sobel)) | |
# Create a mask of 1's where the scaled gradient magnitude is > thresh_min and < thresh_max | |
mask = np.zeros_like(sobel) | |
# Return this mask as your binary_output image | |
mask[(sobel >= threshold[0]) & (sobel <= threshold[1])] = 1 | |
return mask | |
def gradient_magnitude_mask(image, sobel_kernel=3, threshold=(0, 255)): | |
""" | |
Masks the image based on gradient magnitude. | |
Parameters | |
---------- | |
image : Image to mask. | |
sobel_kernel : Kernel of the Sobel gradient operation. | |
threshold : Magnitude threshold for it to make it to appear in the mask. | |
Returns | |
------- | |
Image mask with 1s in activations and 0 in other pixels. | |
""" | |
# Take the gradient in x and y separately | |
sobel_x = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=sobel_kernel) | |
sobel_y = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=sobel_kernel) | |
# Calculate the magnitude | |
magnitude = np.sqrt(sobel_x ** 2 + sobel_y ** 2) | |
# Scale to 8-bit (0 - 255) and convert to type = np.uint8 | |
magnitude = (magnitude * 255 / np.max(magnitude)).astype(np.uint8) | |
# Create a binary mask where mag thresholds are met | |
mask = np.zeros_like(magnitude) | |
mask[(magnitude >= threshold[0]) & (magnitude <= threshold[1])] = 1 | |
# Return this mask as your binary_output image | |
return mask | |
def gradient_direction_mask(image, sobel_kernel=3, threshold=(0, np.pi / 2)): | |
""" | |
Masks the image based on gradient direction. | |
Parameters | |
---------- | |
image : Image to mask. | |
sobel_kernel : Kernel of the Sobel gradient operation. | |
threshold : Direction threshold for it to make it to appear in the mask. | |
Returns | |
------- | |
Image mask with 1s in activations and 0 in other pixels. | |
""" | |
# Take the gradient in x and y separately | |
sobel_x = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=sobel_kernel) | |
sobel_y = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=sobel_kernel) | |
# Take the absolute value of the x and y gradients and calculate the direction of the gradient | |
direction = np.arctan2(np.absolute(sobel_y), np.absolute(sobel_x)) | |
# Create a binary mask where direction thresholds are met | |
mask = np.zeros_like(direction) | |
# Return this mask as your binary_output image | |
mask[(direction >= threshold[0]) & (direction <= threshold[1])] = 1 | |
return mask | |
def color_threshold_mask(image, threshold=(0, 255)): | |
""" | |
Masks the image based on color intensity. | |
Parameters | |
---------- | |
image : Image to mask. | |
threshold : Color intensity threshold. | |
Returns | |
------- | |
Image mask with 1s in activations and 0 in other pixels. | |
""" | |
mask = np.zeros_like(image) | |
mask[(image > threshold[0]) & (image <= threshold[1])] = 1 | |
return mask | |
def get_edges(image, separate_channels=False): | |
""" | |
Masks the image based on a composition of edge detectors: gradient value, | |
gradient magnitude, gradient direction and color. | |
Parameters | |
---------- | |
image : Image to mask. | |
separate_channels : Flag indicating if we need to put masks in different color channels. | |
Returns | |
------- | |
Image mask with 1s in activations and 0 in other pixels. | |
""" | |
# Convert to HLS color space and separate required channel | |
hls = cv2.cvtColor(np.copy(image), cv2.COLOR_RGB2HLS).astype(np.float) | |
s_channel = hls[:, :, 2] | |
# Get a combination of all gradient thresholding masks | |
gradient_x = gradient_abs_value_mask(s_channel, axis='x', sobel_kernel=3, threshold=(20, 100)) | |
gradient_y = gradient_abs_value_mask(s_channel, axis='y', sobel_kernel=3, threshold=(20, 100)) | |
magnitude = gradient_magnitude_mask(s_channel, sobel_kernel=3, threshold=(20, 100)) | |
direction = gradient_direction_mask(s_channel, sobel_kernel=3, threshold=(0.7, 1.3)) | |
gradient_mask = np.zeros_like(s_channel) | |
gradient_mask[((gradient_x == 1) & (gradient_y == 1)) | ((magnitude == 1) & (direction == 1))] = 1 | |
# Get a color thresholding mask | |
color_mask = color_threshold_mask(s_channel, threshold=(170, 255)) | |
if separate_channels: | |
return np.dstack((np.zeros_like(s_channel), gradient_mask, color_mask)) | |
else: | |
mask = np.zeros_like(gradient_mask) | |
mask[(gradient_mask == 1) | (color_mask == 1)] = 1 | |
return mask |