Threshold an image for a given range and Sobel kernel

Question

The class defines functions to threshold an image for a given range and Sobel kernel. Each function has similarities, i.e. they share the same following tasks:

• each function converts an image to grayscale
• take the gradient
• rescale to 8 bit
• apply a threshold, and create a binary image result

class GRADIENT_THRESHOLD(object):
    '''
    Define functions to threshold an image for a given range and Sobel kernel
    '''

    def __init__(self, args):
        self.args = args

    def gradient_abs_sobel(self, img, orient='x', sobel_kernel=3, thresh=(0, 255)):
        # convert to grayscale
        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        # take the absolute value of the gradient in given orient = 'x' or 'y'
        if orient == 'x':
            abs_sobel = np.absolute(cv2.Sobel(gray, cv2.CV_64F, 1, 0))
        if orient == 'y':
            abs_sobel = np.absolute(cv2.Sobel(gray, cv2.CV_64F, 0, 1))
        # scale to 8-bit (0 - 255) then convert to type = np.uint8
        scaled_sobel = np.uint8(255*abs_sobel/np.max(abs_sobel))
        # create a mask of 1's where the scaled gradient magnitude
        # is > thresh_min and < thresh_max
        binary_output = np.zeros_like(scaled_sobel)
        # return this mask as your binary_output image
        binary_output[(scaled_sobel >= thresh[0]) & (scaled_sobel <= thresh[1])] = 1
        # return the binary image
        return binary_output 

    def gradient_magnitude(self, img, sobel_kernel=3, thresh=(0, 255)):
        # convert to grayscale
        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        # take both Sobel x and y gradients
        sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=sobel_kernel)
        sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=sobel_kernel)
        # calculate the gradient magnitude
        gradmag = np.sqrt(sobelx**2 + sobely**2)
        # rescale to 8 bit
        scale_factor = np.max(gradmag)/255 
        gradmag = (gradmag/scale_factor).astype(np.uint8) 
        # create a binary image of ones where threshold is met, zeros otherwise
        binary_output = np.zeros_like(gradmag)
        binary_output[(gradmag >= thresh[0]) & (gradmag <= thresh[1])] = 1
        # return the binary image
        return binary_output


    # Define a function to threshold an image for a given range and Sobel kernel
    def gradient_direction(self, img, sobel_kernel=3, thresh=(0, np.pi/2)):
        # convert to grayscale
        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        # take both Sobel x and y gradients
        sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=sobel_kernel)
        sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=sobel_kernel)
        # take the absolute value of the gradient direction
        absgraddir = np.arctan2(np.absolute(sobely), np.absolute(sobelx))
        # apply a threshold, and create a binary image result
        binary_output =  np.zeros_like(absgraddir)
        binary_output[(absgraddir >= thresh[0]) & (absgraddir <= thresh[1])] = 1
        # return the binary image
        return binary_output

Answer 1

as things are now, there is little use for the class in this case. the self argument is never used in the functions.

This is easily factored into different functions

def grayscale(img):
    return cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

Whether this deserves a separate function or can be inlined is your choice. If you need to type this a lot, or if the method of grayscaling can change in the future, it will be easier to put it in a separate method.. You could also move the colorspace to the arguments to make this more versatile and also accept HSV or Luv encoded images.

def sobel_xy(gray, sobel_kernel=3, absolute=True):
    sobel_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=sobel_kernel)
    sobel_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=sobel_kernel)
    if absolute:
        sobel_x, sobel_y = np.absolute(sobel_x), np.absolute(sobel_y)
    return sobel_x, sobel_y

This calculation is done in all methods, so refactoring it is simple and effective

def gradient_magnitude(gray, sobel_kernel=3):
    return np.hypot(*sobel_xy(gray, sobel_kernel))

def gradient_direction(gray, sobel_kernel=3):
    return np.arctan2(*sobel_xy(gray, sobel_kernel))

This part is pretty straightforward. All I changed was use np.hypot to calculate the magnitude

def gradient_abs_sobel(gray, orient='x'):
    dx, dy = (1, 0) if orient == 'x' else (0, 1)
    sobel = cv2.Sobel(gray, cv2.CV_64F, dx, dy)
    return np.absolute(sobel)

def gradient_abs_sobel2(gray, orient='x'):
    orientation = {'x': 0,'y': 1}[orient]
    sobel = sobel_xy(gray)
    return sobel[orientation]

Depending on taste and the expense of the sobel calculation you can pick either of the 2

def scale(sobel, MAX=255, dtype=np.uint8):
    return dtype(MAX * sobel / np.max(sobel))

def mask(img, lower, upper):
    return (img >= lower) & (img <= upper)

Are pretty straightforward too. Instead of the np.zeros_like you can directly use the boolean array of the comparison

Now instead of calling 1 function, you just call the 4 elements subsequently. If this is too much work you can still make 1 compound method for the 3 cases.

gray = grayscale(img)
abs_sobel = gradient_abs_sobel(gray)
result = mask(scale(abs_sobel), 10, 130)

Answer 2

Thank to Maarten Fabré's feedbacks, I have rewritten the code as the following:

def grayscale(image):
    return cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

def sobel_xy(gray, sobel_kernel=9, absolute=True):
    sobel_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=sobel_kernel)
    sobel_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=sobel_kernel)
    if absolute:
        sobel_x, sobel_y = np.absolute(sobel_x), np.absolute(sobel_y)
    return sobel_x, sobel_y

def sobel_abs(gray, orient='x'):
    orientation = {'x': 0,'y': 1}[orient]
    sobel       = sobel_xy(gray)
    return sobel[orientation]

def scale(sobel, MAX=255, dtype=np.uint8):
    return dtype(MAX * sobel / np.max(sobel))

def mask(image, lower, upper):
    return (image >= lower) & (image <= upper)

def gradient_abs_sobel(image, orient='x', thresh=(0, 255)):
    gray  = grayscale(image)
    sobel = sobel_abs(gray, orient=orient)
    sobel_scale = scale(sobel, MAX=thresh[1], dtype=np.uint8)
    return mask(sobel_scale, thresh[0], thresh[1])

def gradient_magnitude(image, sobel_kernel=9, thresh=(0, 255)):
    gray    = grayscale(image)
    gradmag = np.hypot(*sobel_xy(gray, sobel_kernel=sobel_kernel, absolute=False))
    gradmag = scale(gradmag, MAX=thresh[1], dtype=np.uint8)
    return mask(gradmag, thresh[0], thresh[1])

def gradient_direction(image, sobel_kernel=15, thresh=(0, np.pi/2)): # thresh=(0.7, 1.3)
    gray       = grayscale(image)
    absgraddir = np.arctan2(*sobel_xy(gray, sobel_kernel=sobel_kernel, absolute=True))
    return mask(absgraddir, thresh[0], thresh[1])


def main():
    # parameters and placeholders
    args = PARSE_ARGS()
    # Choose a Sobel kernel size
    ksize = 9
    # Read images
    image        = mpimg.imread(args.sand+'signs_vehicles_xygrad.jpg')
    img_solution = mpimg.imread(args.sand+'binary-combo-example.jpg')
    # Apply each of the thresholding functions
    gradx = gradient_sobel_abs(image, orient='x', sobel_kernel=ksize, thresh=(20, 100))
    grady = gradient_sobel_abs(image, orient='y', sobel_kernel=ksize, thresh=(20, 100))
    mag_binary = gradient_magnitude(image, sobel_kernel=ksize, thresh=(20, 100))
    dir_binary = gradient_direction(image, sobel_kernel=15, thresh=(0.7, 1.3))

    # combine thresholds
    combined1, combined2, combined3 = np.zeros_like(dir_binary), np.zeros_like(dir_binary), np.zeros_like(dir_binary)
    combined1[((gradx == 1) & (grady == 1)) | ((mag_binary == 1) & (dir_binary == 1))] = 1
    combined2[((gradx == 1) & (grady == 1))] = 1
    combined3[((mag_binary == 1) & (dir_binary == 1))] = 1

    # Plot the result
    row, column = [5, 2]
    figure, axes = plt.subplots(row, column, figsize=(15, 20))
    figure.tight_layout()
    list_title_image = [['Original Image',image],
                        ['Expected result', img_solution],
                        ['gradx', gradx],
                        ['grady', grady],
                        ['mag_binary', mag_binary],
                        ['dir_binary', dir_binary],                        
                        ['combined1', combined1],
                        ['combined2', combined2],
                        ['combined3', combined3],
                        ['Original Image', image] ]

    for ax, img in zip(axes.flatten(), list_title_image):
        ax.imshow(img[1], cmap='gray')
        ax.set_title(img[0], fontsize=15)
        ax.axis('off')


if __name__ == '__main__':
    main()