diff --git a/computer_vision/hough_transform.py b/computer_vision/hough_transform.py new file mode 100644 index 000000000..37e4cb497 --- /dev/null +++ b/computer_vision/hough_transform.py @@ -0,0 +1,175 @@ +""" +The Hough transform can be used to detect lines, circles or +other parametric curves. It works by transforming +the image edge map (obtained using a Sobel Filter) to Polar coordinates +and then selecting local maxima in the Parametric space as lines based on +majority voting. + +References: + https://en.wikipedia.org/wiki/Hough_transform + https://www.cs.cmu.edu/~16385/s17/Slides/5.3_Hough_Transform.pdf + https://www.uio.no/studier/emner/matnat/ifi/INF4300/h09/undervisningsmateriale/hough09.pdf + +Requirements (pip): + - matplotlib + - cv2 +""" + +import cv2 +import matplotlib.pyplot as plt +import numpy as np + +from digital_image_processing.edge_detection import canny + + +def generate_accumulator(edges: np.ndarray) -> np.ndarray: + """ + - Generates an accumulator by transforming edge coordinates from Cartesian + to polar coordinates (Hough space). + + - The accumulator array can be indexed as `accumulator[p][theta]` + + Params: + ------ + edges (np.ndarray): The edge-detected binary image (single-channel). + + Returns: + ------ + np.ndarray: The accumulator array with votes for line candidates. + + Example: + ------ + >>> img = np.array([[1, 0, 0,], [1, 0, 0,], [1, 0, 0,],]) + >>> np.sum(generate_accumulator(img)) + np.float64(540.0) + """ + n, m = edges.shape + theta_min, theta_max = 0, 180 + p_min, p_max = 0, int(n * np.sqrt(2) + 1) + accumulator = np.zeros((int(theta_max - theta_min), int(p_max - p_min))) + for x in range(n): + for y in range(m): + if edges[x][y]: + for theta in range(theta_min, theta_max): + p = int( + x * np.cos(np.deg2rad(theta)) + y * np.sin(np.deg2rad(theta)) + ) + accumulator[theta][p] += 1 + return accumulator + + +def hough_transform( + img: np.ndarray, threshold: int = 30, max_num_lines: int = 5 +) -> list[tuple[int, int, np.float64]]: + """ + Performs the Hough transform to detect lines in the input image. + + Params: + ------ + img (np.ndarray): Single-channel grayscale image. + threshold (int): Minimum vote count in the accumulator to consider a line. + max_num_lines (int): Maximum number of lines to return. + + Returns: + ------ + list[tuple[int, int, int]]: List of detected lines in (theta, p, votes) format. + + Raises: + ------ + AssertionError: If the image is not square or single-channel. + + Example: + ------ + >>> img = np.vstack([np.zeros((30, 50)),np.ones((1, 50)),np.zeros((19, 50))]) + >>> hough_transform(img, 30, 1) + [(0, 28, np.float64(48.0))] + """ + assert img.shape[0] == img.shape[1], "image must have equal dimensions" + assert len(img.shape) == 2, "image should be single-channel" + + # Obtain edge map for image + edges = canny.canny(img) + + # Transform to Polar Coordinates + n, _ = img.shape + theta_min, theta_max = 0, 180 + p_min, p_max = 0, int(n * np.sqrt(2) + 1) + accumulator = generate_accumulator(edges) + + # Select maxima in Polar space + res = [] + for theta in range(theta_min, theta_max): + for p in range(p_min, p_max): + if accumulator[theta][p] > threshold: + res.append((theta, p, accumulator[theta][p])) + + res = sorted(res, key=lambda x: x[2], reverse=True)[:max_num_lines] + return res + + +def draw_hough_lines( + img: np.ndarray, + lines: list[tuple], + thickness: int = 1, + color: tuple[int, int, int] = (255, 0, 0), +) -> None: + """ + Draws detected Hough lines on the image. + + Params: + ------ + img (np.ndarray): The input image to draw lines on. + lines (list[tuple[int, int, int]]): + List of (theta, p, votes) for detected lines. + thickness (int): Line thickness. + color (tuple[int, int, int]): BGR color of the lines. + + Example: + ------ + >>> draw_hough_lines(create_dummy_img(), [(50, 0),]) + """ + for line in lines: + theta, p = line[0], line[1] + a = np.sin(np.deg2rad(theta)) + b = np.cos(np.deg2rad(theta)) + x0, y0 = a * p, b * p + x1, y1 = int(x0 + 100 * (-b)), int(y0 + 100 * (a)) + x2, y2 = int(x0 - 100 * (-b)), int(y0 - 100 * (a)) + cv2.line(img, (x1, y1), (x2, y2), color, thickness) + + +def create_dummy_img(height: int = 50, width: int = 50) -> np.ndarray: + """ + Test function to create dummy 3-channel image of specified width and height + Example: + ------ + >>> create_dummy_img(100, 120).shape + (100, 120, 3) + """ + img = np.zeros((height, width), dtype=np.uint8) + cv2.line(img, (10, 10), (int(0.6 * height), int(0.8 * width)), 255, 1) # type: ignore[call-overload] + img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) # type: ignore[assignment] + return img + + +if __name__ == "__main__": + import doctest + + # Run doctests + doctest.testmod() + + img = create_dummy_img(60, 80) + # Preprocess Image + img = cv2.resize(img, (64, 64), interpolation=cv2.INTER_AREA) + gray_image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) + plt.imshow(img) + plt.show() + # Accumulator + accumulator = generate_accumulator(canny.canny(gray_image)) + plt.imshow(accumulator) + plt.show() + # Hough Transform + res = hough_transform(gray_image, 30, 1) + draw_hough_lines(img, res) + plt.imshow(img) + plt.show()