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

def probabilistic_hough_transform(image_path, threshold=50, min_line_length=50, max_line_gap=10):
    """
    使用OpenCV实现概率霍夫变换
    :param image_path: 输入图像路径
    :param threshold: 累加器阈值
    :param min_line_length: 最小直线长度
    :param max_line_gap: 最大线段间隙
    :return: 原图和检测到直线的图像
    """
    # 读取图像
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 边缘检测
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)

    # 使用概率霍夫变换检测直线
    lines = cv2.HoughLinesP(edges, 1, np.pi/180, threshold, minLineLength=min_line_length, maxLineGap=max_line_gap)

    # 在原图上绘制检测到的直线
    img_with_lines = img.copy()
    if lines is not None:
        for x1, y1, x2, y2 in lines[:, 0]:
            cv2.line(img_with_lines, (x1, y1), (x2, y2), (0, 255, 0), 2)

    return img, img_with_lines

def manual_probabilistic_hough(image_path, threshold=50, min_line_length=50, max_line_gap=10):
    """
    手动实现概率霍夫变换
    :param image_path: 输入图像路径
    :param threshold: 累加器阈值
    :param min_line_length: 最小直线长度
    :param max_line_gap: 最大线段间隙
    :return: 概率霍夫变换检测结果
    """
    # 读取图像
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 边缘检测
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)

    # 获取边缘点坐标
    edge_points = np.column_stack(np.where(edges > 0))
    edge_points = edge_points[:, ::-1]  # 交换x和y坐标

    # 参数空间
    height, width = gray.shape
    rho_max = int(np.sqrt(height**2 + width**2))
    theta_range = np.linspace(0, np.pi, 180)

    # 随机选择边缘点
    num_samples = min(len(edge_points), 1000)  # 限制采样数量
    selected_indices = np.random.choice(len(edge_points), size=num_samples, replace=False)
    selected_points = edge_points[selected_indices]

    # 累加器
    accumulator = np.zeros((2 * rho_max, len(theta_range)), dtype=np.uint32)

    # 对选中的点进行投票
    for x, y in selected_points:
        for theta_idx, theta in enumerate(theta_range):
            rho = int(x * np.cos(theta) + y * np.sin(theta))
            if -rho_max <= rho < rho_max:
                accumulator[rho + rho_max, theta_idx] += 1

    # 找到峰值
    peaks = np.where(accumulator >= threshold)
    detected_lines = []

    for rho_idx, theta_idx in zip(peaks[0], peaks[1]):
        rho = rho_idx - rho_max
        theta = theta_range[theta_idx]
        
        # 将极坐标转换为直线的两点表示
        a = np.cos(theta)
        b = np.sin(theta)
        x0 = a * rho
        y0 = b * rho
        x1 = int(x0 + 1000 * (-b))
        y1 = int(y0 + 1000 * (a))
        x2 = int(x0 - 1000 * (-b))
        y2 = int(y0 - 1000 * (a))
        
        detected_lines.append([x1, y1, x2, y2])

    # 在原图上绘制检测到的直线
    img_with_lines = img.copy()
    for x1, y1, x2, y2 in detected_lines:
        cv2.line(img_with_lines, (x1, y1), (x2, y2), (0, 255, 0), 2)

    return img, img_with_lines

def adaptive_probabilistic_hough(image_path):
    """
    自适应概率霍夫变换
    :param image_path: 输入图像路径
    :return: 原图和自适应检测到直线的图像
    """
    # 读取图像
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 边缘检测
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)

    # 使用自适应参数的霍夫变换
    lines = cv2.HoughLinesP(
        edges, 
        rho=1,           # 距离精度
        theta=np.pi/180, # 角度精度
        threshold=100,   # 阈值
        minLineLength=50, # 最小线段长度
        maxLineGap=10    # 最大线段间隔
    )

    # 在原图上绘制检测到的直线
    img_with_lines = img.copy()
    if lines is not None:
        for x1, y1, x2, y2 in lines[:, 0]:
            cv2.line(img_with_lines, (x1, y1), (x2, y2), (0, 255, 0), 2)

    return img, img_with_lines

def compare_standard_vs_probabilistic(image_path):
    """
    比较标准霍夫变换和概率霍夫变换
    :param image_path: 输入图像路径
    :return: 原图、标准霍夫变换结果和概率霍夫变换结果
    """
    # 读取图像
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 边缘检测
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)

    # 标准霍夫变换
    standard_lines = cv2.HoughLines(edges, 1, np.pi/180, 150)

    # 绘制标准霍夫变换结果
    standard_result = img.copy()
    if standard_lines is not None:
        for rho, theta in standard_lines[:, 0]:
            a = np.cos(theta)
            b = np.sin(theta)
            x0 = a * rho
            y0 = b * rho
            x1 = int(x0 + 1000 * (-b))
            y1 = int(y0 + 1000 * (a))
            x2 = int(x0 - 1000 * (-b))
            y2 = int(y0 - 1000 * (a))
            cv2.line(standard_result, (x1, y1), (x2, y2), (0, 0, 255), 2)

    # 概率霍夫变换
    prob_lines = cv2.HoughLinesP(edges, 1, np.pi/180, 100, minLineLength=50, maxLineGap=10)

    # 绘制概率霍夫变换结果
    prob_result = img.copy()
    if prob_lines is not None:
        for x1, y1, x2, y2 in prob_lines[:, 0]:
            cv2.line(prob_result, (x1, y1), (x2, y2), (0, 255, 0), 2)

    return img, standard_result, prob_result

def multi_scale_probabilistic_hough(image_path, scales=[0.5, 1.0, 1.5]):
    """
    多尺度概率霍夫变换
    :param image_path: 输入图像路径
    :param scales: 不同的尺度因子
    :return: 原图和不同尺度下的检测结果
    """
    # 读取图像
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    results = []
    for scale in scales:
        # 根据尺度调整参数
        scaled_img = cv2.resize(gray, None, fx=scale, fy=scale)
        
        # 边缘检测
        edges = cv2.Canny(scaled_img, 50, 150, apertureSize=3)

        # 概率霍夫变换
        lines = cv2.HoughLinesP(
            edges, 
            rho=1, 
            theta=np.pi/180, 
            threshold=int(100 * scale),  # 根据尺度调整阈值
            minLineLength=int(50 * scale), 
            maxLineGap=int(10 * scale)
        )

        # 绘制结果
        result = cv2.cvtColor(scaled_img, cv2.COLOR_GRAY2BGR)
        if lines is not None:
            for x1, y1, x2, y2 in lines[:, 0]:
                cv2.line(result, (x1, y1), (x2, y2), (0, 255, 0), 2)
        
        results.append(result)

    return img, results

def robust_line_detection(image_path):
    """
    鲁棒直线检测（结合多种技术）
    :param image_path: 输入图像路径
    :return: 原图和鲁棒直线检测结果
    """
    # 读取图像
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 使用不同的预处理方法
    # 方法1：高斯模糊
    blurred1 = cv2.GaussianBlur(gray, (5, 5), 0)
    edges1 = cv2.Canny(blurred1, 50, 150)
    
    # 方法2：双边滤波
    blurred2 = cv2.bilateralFilter(gray, 9, 75, 75)
    edges2 = cv2.Canny(blurred2, 50, 150)
    
    # 组合边缘图
    combined_edges = cv2.addWeighted(edges1, 0.5, edges2, 0.5, 0)

    # 概率霍夫变换
    lines = cv2.HoughLinesP(
        combined_edges, 
        rho=1, 
        theta=np.pi/180, 
        threshold=80, 
        minLineLength=30, 
        maxLineGap=15
    )

    # 绘制结果
    result = img.copy()
    if lines is not None:
        for x1, y1, x2, y2 in lines[:, 0]:
            cv2.line(result, (x1, y1), (x2, y2), (0, 255, 0), 2)

    return img, result

# 使用示例
if __name__ == "__main__":
    # 注意：需要提供实际的图像路径
    # img, result = probabilistic_hough_transform('image.jpg', 50, 50, 10)
    # manual_img, manual_result = manual_probabilistic_hough('image.jpg', 50, 50, 10)
    # adaptive_img, adaptive_result = adaptive_probabilistic_hough('image.jpg')
    # comp_img, std_result, prob_result = compare_standard_vs_probabilistic('image.jpg')
    # multi_img, multi_results = multi_scale_probabilistic_hough('image.jpg', [0.5, 1.0, 1.5])
    # robust_img, robust_result = robust_line_detection('image.jpg')
    pass