Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
134 lines (117 sloc) 5.15 KB
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