手动实现霍夫直线检测

手动实现霍夫直线检测

参考视频：hough transform(霍夫变换直线检测)

注意：$r$和$\theta$不是极坐标，这里一根直线只有一组$r$和$\theta$

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors

np.random.seed(777)
np.set_printoptions(precision=3, suppress=True)
np.set_printoptions(formatter={'float': '{: 0.3f}'.format})

def fake_image():
    image = np.zeros([100, 100, 1], dtype=np.uint8)
    rs = []
    thetas = []
    for _ in range(3):
        theta = np.random.uniform(0, np.pi/2)
        r = np.random.randint(30, 70)
        rs.append(r)
        thetas.append(np.rad2deg(theta))
        if theta > np.pi/4:
            x = np.arange(100)
            y = np.uint8((r - x*np.cos(theta)) / np.sin(theta))
            print(list(zip(x, y)))
            valid_y = y[(y >= 0) & (y < 100)]
            valid_x = x[(y >= 0) & (y < 100)]
        else:
            y = np.arange(100)
            x = np.uint8((r - y*np.sin(theta) / np.cos(theta)))
            valid_x = x[(x >= 0) & (x < 100)]
            valid_y = y[(x >= 0) & (x < 100)]
        for i, j in zip(valid_x, valid_y):
            image[j, i, 0] = 255
    return image, np.array(rs), np.array(thetas)

def houghline(image, thetas, rhos, r_max):
    cos_t = np.cos(thetas)
    sin_t = np.sin(thetas)
    
    y_idx, x_idx, _ = np.nonzero(image)
    
    vote = np.zeros([len(rhos), len(thetas)])
    
    for i in range(len(x_idx)):
        x = x_idx[i]
        y = y_idx[i]
        for j in range(len(thetas)):
            r = x*cos_t[j] + y*sin_t[j] + r_max
            vote[int(r), j] += 1
    
    return vote

def get_line(vote, thetas, rhos):
    row, col = np.unravel_index(np.argmax(vote, axis=None), vote.shape)
    max_value = vote[row, col]
    thr = 0.5 * max_value
    ii, jj = np.nonzero(vote > thr)
    
    # NMS
    cands = list(zip(ii, jj, vote[ii, jj]))
    cands = sorted(cands, key=lambda x: x[2])
    final_cands = []
    while cands:
        final_cands.append(cands.pop())
        temp_cands = []
        for i in range(len(cands)):
            if abs(cands[i][0] - final_cands[-1][0]) > 2 or abs(cands[i][1] - final_cands[-1][1]) > 2:
                temp_cands.append(cands[i])

        cands = temp_cands
    
    pred_rs = []
    pred_thetas = []
    for cand in final_cands:
        theta = thetas[cand[1]]
        rho = rhos[cand[0]]
        pred_rs.append(rho)
        pred_thetas.append(np.rad2deg(theta))
    return np.array(pred_rs), np.array(pred_thetas)
    
    
        
    
if __name__ == "__main__":
    image, target_rs, target_thetas = fake_image()
    
    thetas = np.deg2rad(np.arange(0, 180))
    height, width = image.shape[:2]
    r_max = np.ceil(np.linalg.norm([height, width]))
    rhos = np.linspace(-r_max, r_max, int(r_max*2) + 1)
    
    vote = houghline(image, thetas, rhos, r_max)
    
    pred_rs, pred_thetas = get_line(vote, thetas, rhos)
    
    print(f"target rs: ", target_rs, "target thetas: ", target_thetas)
    print(f"pred rs: ", pred_rs, "pred thetas: ", pred_thetas)
    
    fig, axes = plt.subplots(1, 2, figsize=(8, 8))
    axes[0].imshow(image, cmap=plt.cm.gray)
    axes[0].set_title('Input image')
    pcm = axes[1].imshow(vote, extent=[0, vote.shape[1], vote.shape[0]//2, -vote.shape[0]//2])
    fig.colorbar(pcm, ax=axes[1])
    axes[1].set_title('Hough transform')
    axes[1].set_xlabel('Angle (degree)')
    axes[1].set_ylabel('Distance (pixel)')
    plt.tight_layout()
    plt.show()

"""
target rs:  [68 37 69] target thetas:  [ 13.740  5.583  22.325]
pred rs:  [ 65.000  36.000  63.000] pred thetas:  [ 14.000  6.000  22.000]
"""