import cv2
import numpy as np
import math
import os


def find_and_mask_weld_seam_v3(image_path):
    """
    加载图像，使用更激进的预处理方法检测低对比度直线。
    :param image_path: 输入图像的完整路径。
    :return: None
    """
    # --- 步骤 1: 加载图像 ---
    image = cv2.imread(image_path)
    if image is None:
        print(f"错误: 无法加载图像 '{image_path}'。")
        return
    visualization_image = image.copy()
    height, width, _ = image.shape
    image_center_x = width / 2

    # --- 步骤 2: 图像预处理 (超强版) ---
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 2.1 高斯模糊以去除噪声
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)

    # 2.2 自适应阈值二值化 (关键步骤!)
    # 这个操作对于光照不均、对比度低的区域效果非常好，能强制分离出边缘
    # blocksize: 邻域大小，必须是奇数
    # C: 从均值或加权均值中减去的常数，用于微调
    binary_image = cv2.adaptiveThreshold(
        blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV, 21, 5
    )
    # 使用THRESH_BINARY_INV是因为我们感兴趣的边缘是暗色背景下的亮色线条

    # 2.3 形态学操作来清理和连接线条
    # 创建一个垂直的结构元素，因为我们的目标是垂直线
    kernel = np.ones((15, 1), np.uint8)  # 强调垂直方向
    # 使用“闭运算”来填充断开的线段中的小洞
    closed = cv2.morphologyEx(binary_image, cv2.MORPH_CLOSE, kernel)

    # --- 步骤 3: 霍夫直线变换 (直接在处理后的二值图上进行) ---
    # Canny不再是必需的，因为二值化已经完成了边缘提取
    lines = cv2.HoughLinesP(
        closed,  # 在闭运算结果上检测
        rho=1,
        theta=np.pi / 180,
        threshold=100,  # 需要较高的阈值，因为线现在更连续了
        minLineLength=int(height * 0.5),  # 线至少要有一半图像高度那么长
        maxLineGap=50  # 允许较大的间隙
    )

    # --- 步骤 4: 筛选最佳直线 (逻辑与之前类似) ---
    best_line = None
    max_length = 0

    if lines is not None:
        for line in lines:
            x1, y1, x2, y2 = line[0]
            # 角度过滤，寻找近乎垂直的线
            angle = abs(math.degrees(math.atan2(y2 - y1, x2 - x1)))
            if abs(angle - 90) < 5:  # 检查角度是否在90度左右（垂直）
                length = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
                if length > max_length:
                    max_length = length
                    best_line = line[0]

    # --- 步骤 5: 处理和显示结果 ---
    if best_line is not None:
        bx1, by1, bx2, by2 = best_line
        print(f"成功找到直线！起点:({bx1}, {by1}), 终点:({bx2}, {by2})")

        cv2.line(visualization_image, (bx1, by1), (bx2, by2), (0, 0, 255), 3)

        # 显示中间处理结果以供调试
        cv2.imshow('1. Original Image', image)
        cv2.imshow('2. Adaptive Threshold', binary_image)
        cv2.imshow('3. Morphological Closing', closed)  # 检查闭运算效果
        cv2.imshow('4. Final Detection', visualization_image)

    else:
        print("未能在图像中找到符合条件的直线。")
        cv2.imshow('1. Original Image', image)
        cv2.imshow('2. Adaptive Threshold', binary_image)
        cv2.imshow('3. Morphological Closing', closed)

    cv2.waitKey(0)
    cv2.destroyAllWindows()


# --- 主程序入口 ---
if __name__ == '__main__':
    try:
        script_dir = os.path.dirname(os.path.abspath(__file__))
    except NameError:
        script_dir = os.getcwd()

    # 请确保你的图片路径正确
    image_filename = 'data/test3/l2.jpeg'
    image_path = os.path.join(script_dir, image_filename)

    if not os.path.exists(image_path):
        # 如果找不到，尝试直接使用文件名（假设图片和脚本在同一目录）
        image_path = '13-l.bmp'
        if not os.path.exists(image_path):
            print(f"错误: 输入文件 '{image_path}' 不存在。")
        else:
            find_and_mask_weld_seam_v3(image_path)
    else:
        find_and_mask_weld_seam_v3(image_path)