霍夫直线检测

使用python做霍夫直线检测，直接详细代码，注释清晰

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

if __name__ == '__main__':
    pic_path = 'C:/Users/Echo/Desktop/python_file/hough_project/02.jpg'  # 路径
    img = plt.imread(pic_path)
    blurred = cv2.GaussianBlur(img, (3, 3), 0)  # 高斯线性平滑滤波，高斯核为3*3，标准差为0
    if not len(blurred.shape) == 2:
        gray = cv2.cvtColor(blurred, cv2.COLOR_RGB2GRAY)  # 这个函数的参数本可以是灰度图！
    else:
        gray = blurred
    edge = cv2.Canny(blurred, 50, 150)  # 根据两个阈值canny边缘检测得到二值图，参数必须是灰度图！

    # 霍夫直线检测
    img_size = edge.shape
    max_dist = np.sqrt(img_size[0] ** 2 + img_size[1] ** 2)  # 对角线的长度
    max_dist = int(np.ceil(max_dist))  # 投票箱的高度r，取最大值到对角线长度，向上取整
    w = int(max_dist / 50)
    theta_dim = 90  # 角度分90个区间，2°一个区间
    accumulator = np.zeros((theta_dim, max_dist))  # 投票箱零矩阵
    sin_theta = [np.sin(t * np.pi / theta_dim) for t in range(theta_dim)]  # 先把三角函数值计算好，备用d
    cos_theta = [np.cos(t * np.pi / theta_dim) for t in range(theta_dim)]
    for i in range(img_size[0]):
        for j in range(img_size[1]):  # 挨个检测是边缘的点
            if not edge[i, j] == 0:
                for k in range(theta_dim):  # 每个角度都要计算一次
                    accumulator[k][int(round(i * cos_theta[k] + j * sin_theta[k]))] += 1  # 投票箱加一

    # 阈值化，大于这个数的都是直线
    M = np.max(accumulator)
    threshold = int(M * 2.3875 / 10)  # 值是可以改的，随便设
    r = np.array(np.where(accumulator > threshold))  # 这个得到的是满足条件的索引号，因为是二维矩阵得到的是两行的二维矩阵！

    # 非极大值抑制
    temp = [[], []]
    for i in range(r.shape[1]):  # 遍历r中所有的坐标点，下面这样写法在边缘的时候是2*2矩阵，其他时候是3*3矩阵
        eight_neiborhood = accumulator[max(0, r[0, i] - 1):min(r[0, i] + 2, accumulator.shape[0]),
                                       max(0, r[1, i] - w + 1):min(r[1, i] + w, accumulator.shape[1])]
        if (accumulator[r[0, i], r[1, i]] >= eight_neiborhood).all():  # 矩阵和阈值对比，全为真才返回真
            temp[0].append(r[0, i])
            temp[1].append(r[1, i])
    temp = np.array(temp)  # 变为numpy格式的矩阵，而不是python再带的list格式
    temp = temp.astype(np.float64)  # 强制类型转换
    temp[0] = temp[0] * np.pi / theta_dim  # 化成弧度制的角度，之前的是分成的块数

    # 在原图上画直线
    if len(blurred.shape) == 2:
        img_end = np.dstack((blurred, blurred, blurred))  # 如果是灰度图要叠加成三层的图
    else:
        img_end = blurred
    color = (255, 0, 0)  # 在原图上画红色的线
    cos_ = np.cos(temp[0])
    sin_ = np.sin(temp[0])
    for i in range(blurred.shape[0]):  # 遍历整个图片像素点
        for j in range(blurred.shape[1]):
            e = np.abs(temp[1] - i * cos_ - j * sin_)  # 把元素点带入到方程，有一个成立即可，然后只画一个点
            if (e < 3).any():  # 全为否则返回否，但凡有一个<3成立就可以
                img_end[i, j] = color  # 只画一个点
    plt.imshow(img_end, cmap='gray')
    plt.axis('off')
    plt.show()