hough.py 文件源码

def hough(img, nb_lines):
    """Applies the Hough Transformation to an image.
    Args:
        img The image
        nb_lines The number of lines to search for.
    Returns:
        Accumulator image,
        Local maxima in accumulator image,
        image with detected lines"""
    height, width = img.shape
    magnitude = grad_magnitude(img)
    mag_avg = np.average(magnitude)

    max_d = math.sqrt(height**2 + width**2)
    min_d = -max_d
    alphas = np.linspace(0, np.pi, NB_QUANTIZATION_STEPS)
    distances = np.linspace(min_d, max_d, NB_QUANTIZATION_STEPS)
    accumulator = np.zeros((NB_QUANTIZATION_STEPS, NB_QUANTIZATION_STEPS))
    for y in range(1, height-1):
        for x in range(1, width-1):
            if magnitude[y, x] > mag_avg:
                for alpha_idx, alpha in enumerate(alphas):
                    distance = x * math.cos(alpha) + y * math.sin(alpha)
                    distance_idx = util.quantize_idx(distance, distances)
                    accumulator[alpha_idx, distance_idx] += 1

    img_hough = np.zeros((height, width, 3))
    img_hough[:, :, 0] = np.copy(img)
    img_hough[:, :, 1] = np.copy(img)
    img_hough[:, :, 2] = np.copy(img)

    local_maxima = find_local_maxima(accumulator)
    peaks_idx = get_peak_indices(local_maxima, nb_lines)
    for value, (alpha_idx, distance_idx) in peaks_idx:
        peak_alpha_rad = alphas[alpha_idx]
        peak_distance = distances[distance_idx]

        x0 = 0
        x1 = width - 1
        y0 = (peak_distance - 0 * np.cos(peak_alpha_rad)) / (np.sin(peak_alpha_rad) + 1e-8)
        y1 = (peak_distance - (width-1) * np.cos(peak_alpha_rad)) / (np.sin(peak_alpha_rad) + 1e-8)

        y0 = np.clip(y0, 0, height-1)
        y1 = np.clip(y1, 0, height-1)

        rr, cc = draw.line(int(y0), int(x0), int(y1), int(x1))
        img_hough[rr, cc, 0] = 1
        img_hough[rr, cc, 1] = 0
        img_hough[rr, cc, 2] = 0

    return accumulator, local_maxima, img_hough