帮我写一段python的LSWMS直线检测算法

LSWMS（Least Square Weighted Mean Subtraction）直线检测算法，是一种基于最小二乘法的直线检测算法。它可以在图像中自动识别直线和曲线，并准确地提取它们的参数，常用于图像处理、计算机视觉和机器人等领域。

下面是一个基于Python的LSWMS直线检测算法的例子：

import numpy as np
import cv2

def LSWMS(img, threshold):
    # 预处理图像，将其转换为灰度图像
    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # 计算图像的梯度
    sobelx = cv2.Sobel(gray_img, cv2.CV_64F, 1, 0, ksize=3)
    sobely = cv2.Sobel(gray_img, cv2.CV_64F, 0, 1, ksize=3)
    # 计算梯度方向
    theta = np.arctan2(sobely, sobelx)
    # 初始化图像空间
    height, width = img.shape[:2]
    accumulator = np.zeros((height, width), dtype=np.uint64)
    # 遍历图像像素，累加直线
    for y in range(height):
        for x in range(width):
            if np.abs(theta[y, x]) < np.pi/4 or np.abs(theta[y, x]) > 3*np.pi/4:
                r = x
            else:
                r = y
            # 计算权重
            weight = np.exp(-np.power((sobelx[y, x]*np.cos(theta[y, x]) + sobely[y, x]*np.sin(theta[y, x])), 2)/threshold**2)
            accumulator[r, int(y*np.sin(theta[y, x]) + x*np.cos(theta[y, x]))] += weight
    # 选出最优直线
    lines = []
    for r in range(height):
        for c in range(width):
            if accumulator[r, c] > threshold:
                lines.append((r, c, accumulator[r, c]))
    lines = sorted(lines, key=lambda x:x[2], reverse=True)
    return lines

# 测试
img = cv2.imread('test.jpg')
lines = LSWMS(img, 100)
for line in lines:
    r, c, weight = line
    theta = np.arctan2(sobely[r, c], sobelx[r, c])
    a = np.cos(theta)
    b = np.sin(theta)
    x0 = a*c
    y0 = b*c
    x1 = int(x0 + 1000*(-b))
    y1 = int(y0 + 1000*(a))
    x2 = int(x0 - 1000*(-b))
    y2 = int(y0 - 1000*(a))
    cv2.line(img,(x1,y1),(x2,y2),(0,0,255),2)
cv2.imshow('LSWMS', img)
cv2.waitKey(0)
cv2.destroyAllWindows()

这段代码首先通过Sobel算子计算了图像的梯度，然后根据梯度方向和权重对图像中的每个像素进行了累加，从而检测出了直线。最后，选出了累加值最高的直线，并在原图上画出了此直线。