Skip to content
Permalink
master
Switch branches/tags

Name already in use

A tag already exists with the provided branch name. Many Git commands accept both tag and branch names, so creating this branch may cause unexpected behavior. Are you sure you want to create this branch?
Go to file
 
 
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