Threshold an image for a given range and Sobel kernel 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 &gt; thresh_min and &lt; thresh_max
        binary_output = np.zeros_like(scaled_sobel)
        # return this mask as your binary_output image
        binary_output[(scaled_sobel &gt;= thresh[0]) &amp; (scaled_sobel &lt;= 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 &gt;= thresh[0]) &amp; (gradmag &lt;= 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 &gt;= thresh[0]) &amp; (absgraddir &lt;= thresh[1])] = 1
        # return the binary image
        return binary_output
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 &gt;= lower) &amp; (img &lt;= 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)