# 该函数使用预训练的 HED (Holistically-Nested Edge Detection) 模型来检测图像中的边缘。
# 通过使用 OpenCV 库加载模型并传递输入图像，该函数生成并返回二值化的边缘检测结果。
# 参数:
#   img_rgb (numpy.ndarray): 输入的 RGB 图像，图像尺寸应为 (height, width, 3)。
#   blur_window (int, 可选): 用于对输入图像和边缘检测结果进行高斯模糊的窗口大小。默认值为 5。
#   scalefactor (float, 可选): 输入图像在创建 blob 时的缩放系数。默认值为 1.0。
# 返回:
#   hed (numpy.ndarray): 二值化的边缘检测结果，图像尺寸与输入图像相同。
# 示例:
#   >>> img = cv2.imread("example.jpg")
#   >>> hed = get_hed(img)
#   >>> plt.imshow(hed)
def get_hed(img_rgb,blur_window=5, scalefactor=1.0):

   # 加载 HED 模型
    model_path='model/hed_pretrained_model.caffemodel'
    prototxt_path='model/deploy.prototxt'
    hed_model = cv2.dnn.readNet(model_path, prototxt_path)
    # 将图像转换为 blob
    blob = cv2.dnn.blobFromImage(img_rgb, scalefactor=scalefactor, size=(256, 256), mean=(104.00698793, 116.66876762, 122.67891434), swapRB=False, crop=False)
    # 输入 blob 到 HED 模型中进行边缘检测
    hed_model.setInput(blob)
    hed = hed_model.forward()
    # 获取输出结果中的第一个通道作为边缘检测结果
    hed = cv2.resize(hed[0, 0], (img_rgb.shape[1], img_rgb.shape[0]))
    # 对结果进行二值化处理
    _, hed = cv2.threshold(hed, 0.9 * hed.max(), 255, cv2.THRESH_BINARY)
    hed = np.uint8(hed)
    # 对结果进行高斯模糊处理
    hed = cv2.GaussianBlur(hed, (blur_window, blur_window), 0)
    return hed




# 该函数通过比较输入的灰度图像与 HED 边缘检测结果之间的均方误差 (MSE)，自动选择最佳的 Canny 边缘检测阈值参数。使用最佳阈值参数计算最终的边缘检测结果。
# 参数:
#   img_gray (numpy.ndarray): 输入的灰度图像，图像尺寸应为 (height, width)。
#   hed (numpy.ndarray): 使用 HED 模型生成的边缘检测结果，图像尺寸应与输入灰度图像相同。
# 返回:
#   edge_img (numpy.ndarray): 使用最佳阈值参数生成的 Canny 边缘检测结果。
#   best_params (tuple): 包含最佳阈值参数的元组，格式为 (best_sigma, min_mse, lower, upper)。
#       best_sigma: 为具有最低MSE的sigma值
#       min_mse: 最低的mse值
#       lower, upper: 用best_sigma进行canny时的上下限
#   plot_data (tuple): 包含用于绘制 MSE 与 sigma 值关系图的数据，格式为 (sigma_values, mse_values)。
#       sigma_values: 所有测试过的sigma_value,建议范围从0.001到1.0实验100个值
#       mse_values: 对于每一个sigma的mse值
# 示例:
# >>> img_gray = cv2.cvtColor(cv2.imread("example.jpg"), cv2.COLOR_BGR2GRAY)
# >>> hed = get_hed(cv2.imread("example.jpg"))
# >>> edge_img, best_params, plot_data = auto_canny(img_gray, hed)
# >>> plt.imshow(edge_img)
def auto_canny(img_gray,hed):
    sigma_values = np.linspace(0.001, 1.0, 100)  # 生成 100 个 sigma 值
    for element in tqdm(sigma_values,desc="提示词:",unit="每秒单位"):
        time.sleep(0.1)
    mse_values = []
    for sigma in sigma_values:
        # 使用给定的 sigma 值计算 Canny 边缘检测
        canny_edges = cv2.Canny(img_gray, int(sigma * 50), int(sigma * 100))
        # 计算均方误差
        mse = np.mean((canny_edges.astype("float") - hed.astype("float")) ** 2)
        mse_values.append(mse)
    # 找到最小 MSE 对应的索引
    min_mse_index = np.argmin(mse_values)
    best_sigma = sigma_values[min_mse_index]
    # 根据最佳 sigma 值计算 Canny 边缘检测
    lower = int(best_sigma * 50)
    upper = int(best_sigma * 100)
    edge_img = cv2.Canny(img_gray, lower, upper)
    best_params = (best_sigma, mse_values[min_mse_index], lower, upper)
    plot_data = (sigma_values, mse_values)
    return edge_img, best_params, plot_data