import cv2
import numpy as np
import os


def detect_cracks_in_wood(image_path, output_dir="output"):
    """
    检测木条图像中的裂缝缺陷（线状或不规则形状）

    Args:
        image_path (str): 输入图像路径
        output_dir (str): 结果保存目录
    """
    os.makedirs(output_dir, exist_ok=True)

    # 1. 读取图像
    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"无法读取图像: {image_path}")

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

    # 2. 使用高斯模糊去除噪声，并增强裂缝对比度
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)

    # 3. 使用Canny边缘检测算法查找边缘
    edges = cv2.Canny(blurred, 50, 150)

    # 4. 形态学操作：使用闭运算填充裂缝间的空隙
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    closed_edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel, iterations=2)

    # 5. 查找轮廓
    contours, _ = cv2.findContours(closed_edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 6. 筛选疑似裂缝的轮廓
    min_length = 200  # 最小长度（像素），根据实际情况调整
    max_gap = 10  # 连接断开点的最大距离（像素）

    crack_contours = []
    for cnt in contours:
        length = cv2.arcLength(cnt, False)  # 计算曲线长度
        if length > min_length:
            crack_contours.append(cnt)

    # 7. 在原图上绘制检测到的裂缝
    result_img = original.copy()
    cv2.drawContours(result_img, crack_contours, -1, (0, 0, 255), 2)  # 红色框

    # 8. 保存结果
    base_name = os.path.splitext(os.path.basename(image_path))[0]
    cv2.imwrite(os.path.join(output_dir, f"{base_name}_edges.png"), edges)
    cv2.imwrite(os.path.join(output_dir, f"{base_name}_closed_edges.png"), closed_edges)
    cv2.imwrite(os.path.join(output_dir, f"{base_name}_result.png"), result_img)

    print(f"✅ 检测完成！共找到 {len(crack_contours)} 个疑似裂缝")
    print(f"📁 结果已保存至: {output_dir}")

    # （可选）显示
    cv2.imshow("Original", cv2.resize(original, (640, 480)))
    cv2.imshow("Edges", cv2.resize(edges, (640, 480)))
    cv2.imshow("Closed Edges", cv2.resize(closed_edges, (640, 480)))
    cv2.imshow("Result", cv2.resize(result_img, (640, 480)))
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == "__main__":
    image_path = "1.jpg"  # ← 替换为你的木条图像路径
    detect_cracks_in_wood(image_path)