python opencv 直线检测

python opencv直线检测与拟合

河北一帆

2354人浏览 · 2022-07-14 18:28:16

河北一帆 · 2022-07-14 18:28:16 发布

先创建一个“点”类和一个“直线”类

class Endpoint:
    def __init__(self, x=0.0, y=0.0):
        self.x = x
        self.y = y

from endpoint import Endpoint


class Line:
    def __init__(self, p1, p2):
        self.p1 = p1
        self.p2 = p2
        self.k = (p1.y - p2.y) / (p1.x - p2.x)
        self.b = p1.y - self.k * p1.x

        self.catagory = 0

    def set_line_segment_end(self, p1, p2):
        self.segment_p1 = p1
        self.segment_p2 = p2

对图像依次进行：转灰度图、二值化、锐化、霍夫变换、直线拟合几项操作。

转灰度图

gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

采用大津自适应二值化算法

retval, dst = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)

锐化，为了让直线边界更清楚

# 锐化
kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]], np.float32)  # 定义一个核
dst2 = cv2.filter2D(dst, -1, kernel=kernel)

创建一个三通道的黑白图，方便把拟合的直线画在上面，查看结果用，二值化的图为单通道上面不能画彩色的线

dst3 = cv2.cvtColor(dst2, cv2.COLOR_GRAY2BGR)

做霍夫变换

lines = cv2.HoughLinesP(dst2, 1, np.pi / 180, 100, 100, 10)

对霍夫变换检测出来的线进行处理，创建成我们自己创建的“直线”类。

for i in range(lines.shape[0]):
    for x1, y1, x2, y2 in lines[i]:
        p1 = Endpoint(x1, y1)
        p2 = Endpoint(x2, y2)
        line = Line(p1, p2)
        if line.k > 0.1:
            line.catagory = 1
            lineList_1.append(line)
        elif line.k < -0.1:
            line.catagory = 3
            lineList_3.append(line)
        else:
            if line.b < img.shape[0] / 2:
                line.catagory = 2
                lineList_2.append(line)
            elif line.b > img.shape[0] / 2:
                line.catagory = 4
                lineList_4.append(line)

        kList.append(line.k)
        bList.append(line.b)

霍夫变换检测出来是一堆小短线，我们需要的结果是一根完整的长线，对这些短线进行拟合，拟合的方法是，取所有线段的端点，对端点进行加权最小二乘拟合，得到我们想要的线。

加权最小二乘拟合公式如下

import numpy as np


def weighted_least_square_method(X, Y, W):
    Y = Y[:, np.newaxis]
    ones = np.ones((len(X), 1))
    X = np.hstack((X[:, np.newaxis], ones))
    temp = np.matmul(W, X)
    temp2 = np.linalg.inv(np.matmul(X.T, temp))

    temp3 = np.matmul(W, Y)
    temp4 = np.matmul(X.T, temp3)

    theta = np.matmul(temp2, temp4)

    return theta

拟合之前还有一步剔除离群点

import numpy as np


def detect_outliers(data, threshold=2):
    """离群点检测"""
    mean_d = np.mean(data)
    std_d = np.std(data)
    outliers = []

    for y in data:
        z_score = (y - mean_d) / std_d
        if np.abs(z_score) > threshold:
            outliers.append(y)

    idx = []
    for i in range(len(data)):
        for j in range(len(outliers)):
            if data[i] == outliers[j]:
                idx.append(i)

    return outliers, idx

超过标准差一定倍数（threshold）的点，认为是离群点剔除。

完整代码：

import cv2
import numpy as np
import matplotlib.pyplot as plt
from endpoint import Endpoint
from line import Line
from find_isolation import detect_outliers
from weighted_lsm import weighted_least_square_method


global img

def mouse(event, x, y, flags, param):
    global img
    if event == cv2.EVENT_LBUTTONDOWN:
        xy = "%d,%d" % (x, y)
        cv2.circle(img, (x, y), 1, (255, 255, 255), thickness = -1)
        cv2.putText(img, xy, (x, y), cv2.FONT_HERSHEY_PLAIN,
                    1.0, (255, 255, 255), thickness = 1)
        cv2.imshow("image", img)


def line_fit(lineList, img, type):
    X_ = []
    Y_ = []
    for line in lineList:
        X_.append(line.p1.x)
        X_.append(line.p2.x)
        Y_.append(line.p1.y)
        Y_.append(line.p2.y)

        # cv2.line(img, (int(line.p1.x), int(line.p1.y)), (int(line.p2.x), int(line.p2.y)), (255, 0, 0), 1)
        cv2.circle(img, (line.p1.x, line.p1.y), 1, (255, 0, 0), 2)
        cv2.circle(img, (line.p2.x, line.p2.y), 1, (255, 0, 0), 2)

    XX = np.array(X_)
    YY = np.array(Y_)

    outliers_X, outliers_X_indx = detect_outliers(XX)

    X = np.delete(XX, outliers_X_indx)
    Y = np.delete(YY, outliers_X_indx)

    W = np.eye(len(X))
    for i in range(len(X)):
        W[i, i] = X[i]

    if type == 2 or type == 4:
        [k, b] = np.polyfit(X, Y, deg=1)
    else:
        [k, b] = weighted_least_square_method(X, Y, W)
    p1 = Endpoint(0, b)
    p2 = Endpoint(3000, k * 3000 + b)
    line_fit = Line(p1, p2)
    p1_segment = Endpoint(min(X), k * min(X) + b)
    p2_segment = Endpoint(max(X), k * max(X) + b)
    line_fit.set_line_segment_end(p1_segment, p2_segment)

    cv2.line(img, (int(line_fit.p1.x), int(line_fit.p1.y)), (int(line_fit.p2.x), int(line_fit.p2.y)), (255, 255, 0), 1)

    return line_fit


def find_crosspoint(addr):
    print(addr)
    global img
    img = cv2.imread(addr)
    # 灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # canny = cv2.Canny(gray, 100, 200)
    # 二值化
    # dst = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 25, 5)
    retval, dst = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)

    # 腐蚀
    # kernel = np.ones((3, 3), np.uint8)
    # erosion = cv2.erode(dst, kernel, iterations=1)
    # erosion2 = cv2.cvtColor(erosion, cv2.COLOR_GRAY2BGR)

    # ##膨胀
    # kernel_d = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    # dilated = cv2.dilate(dst, kernel_d)

    # 锐化
    kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]], np.float32)  # 定义一个核
    dst2 = cv2.filter2D(dst, -1, kernel=kernel)
    dst3 = cv2.cvtColor(dst2, cv2.COLOR_GRAY2BGR)

    # 霍夫线变换
    lines = cv2.HoughLinesP(dst2, 1, np.pi / 180, 100, 100, 10)

    lineList_1 = []
    lineList_2 = []
    lineList_3 = []
    lineList_4 = []
    kList = []
    bList = []

    # 将检测的线绘制
    for i in range(lines.shape[0]):
        for x1, y1, x2, y2 in lines[i]:
            p1 = Endpoint(x1, y1)
            p2 = Endpoint(x2, y2)
            line = Line(p1, p2)
            if line.k > 0.1:
                line.catagory = 1
                lineList_1.append(line)
            elif line.k < -0.1:
                line.catagory = 3
                lineList_3.append(line)
            else:
                if line.b < img.shape[0] / 2:
                    line.catagory = 2
                    lineList_2.append(line)
                elif line.b > img.shape[0] / 2:
                    line.catagory = 4
                    lineList_4.append(line)

            kList.append(line.k)
            bList.append(line.b)

    l1 = line_fit(lineList_1, img, 1)
    l2 = line_fit(lineList_2, img, 2)
    l3 = line_fit(lineList_3, img, 3)
    l4 = line_fit(lineList_4, img, 4)

    # find corss point
    x_up = (l2.b - l1.b) / (l1.k - l2.k)
    y_up = l1.k * x_up + l1.b
    p_up = Endpoint(x_up, y_up)
    x_down = (l4.b - l3.b) / (l3.k - l4.k)
    y_down = l3.k * x_down + l3.b
    p_down = Endpoint(x_down, y_down)

    # find end point
    p_rightUP = l2.segment_p1
    p_rightDown = l4.segment_p1

    dis_parallel_real = 63
    base_dis_img = l4.b - l2.b
    dis_img_up = l2.segment_p1.x - p_up.x
    dis_real_up = (dis_img_up / base_dis_img) * dis_parallel_real

    dis_real_down = l4.segment_p1.x - p_down.x
    dis_real_down = (dis_real_down / base_dis_img) * dis_parallel_real

    dis_real = (dis_real_up + dis_real_down) / 2

    cv2.circle(img, (int(p_up.x), int(p_up.y)), 1, (0, 255, 0), 2)
    cv2.circle(img, (int(p_down.x), int(p_down.y)), 1, (0, 255, 0), 2)
    cv2.circle(img, (int(p_rightUP.x), int(p_rightUP.y)), 1, (0, 255, 0), 2)
    cv2.circle(img, (int(p_rightDown.x), int(p_rightDown.y)), 1, (0, 255, 0), 2)
    print(dis_real, dis_real_up, dis_real_down)

    cv2.imshow("img", img)
    cv2.setMouseCallback("img", mouse)
    cv2.waitKey(0)

    return p_up, p_down


name_list = [""]

for name_i in name_list:
    # addr = 注意这里路径不能有中文
    p_up, p_down = find_crosspoint(addr)