add hough transform

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()