# 该函数使用预训练的 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):

    model_path = './model/hed_pretrained_bsds.caffemodel'
    prototxt_path = './model/deploy.prototxt'

    # 使用OpenCV加载预训练模型
    net = cv2.dnn.readNetFromCaffe(prototxt_path, model_path)

    H, W = img_rgb.shape[:2]

    img = cv2.GaussianBlur(img_rgb, (blur_window, blur_window), 0)  # opencv 高斯滤波

    # 从图像创建blob(一大块数据)
    blob = cv2.dnn.blobFromImage(img, scalefactor, size=(W, H), mean=(105, 117, 123), swapRB=False, crop=False)

    # 将 blob 设置为网络的输入
    net.setInput(blob)
    # 执行前向传递以获得边缘检测结果
    hed_output = net.forward()
    # 从输出中提取边缘图
    hed = hed_output[0, 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):

    img = cv2.GaussianBlur(img_gray, (5, 5), 0)  #plot_data (tuple): 包含用于绘制 MSE 与 sigma 值关系图的数据，格式为 (sigma_values, mse_values)。
#       sigma_values: 所有测试过的sigma_value,建议范围从0.001到1.0实验100个值
#       mse_values: 对于每一个sigma的mse值
# 示例: opencv 高斯滤波

    # 计算图像的灰度中值
    median = np.median(img)

    # 定义函数get_bound，计算并返回自动确定的低阈值和高阈值
    def get_bound(sigma):
        lower = int(max(0, (1.0 - sigma) * median))
        upper = int(min(255, (1.0 + sigma) * median))
        return lower, upper

    # 生成从0.01到1的sigma值，步长为0.01
    sigma_values = np.arange(0.01, 1, 0.01)
    # 初始化最佳sigma值和最高PSNR
    best_sigma = None
    min_mse = 10000000

    # 初始化一个列表，用于存储所有的MSE值
    mse_values = []

    # 遍历sigma值范围，计算每个去噪图像的MSE
    for sigma in tqdm(sigma_values, desc="搜索最佳 sigma", unit="sigma"):
        lower_1, upper_1 = get_bound(sigma)
        img_autocanny = cv2.Canny(img, lower_1, upper_1)
        current_mse = mean_squared_error(hed, img_autocanny)
        mse_values.append(current_mse)
        if current_mse < min_mse:
            min_mse = current_mse
            best_sigma = sigma

    # 使用最佳sigma值进行去噪
    print(best_sigma)
    lower, upper = get_bound(best_sigma)
    edge_img = cv2.Canny(img_autocanny, lower, upper)

    return edge_img,(best_sigma,min_mse,lower,upper),(sigma_values,mse_values)


class CropLayer(cv2.dnn.Layer):
    # 构造函数，用于初始化裁剪的起始和结束位置
    def __init__(self, params, blobs):
        super(CropLayer, self).__init__()
        self.xstart = 0
        self.xend = 0
        self.ystart = 0
        self.yend = 0

    # 该层接收两个输入。
    # 我们需要将第一个输入blob裁剪成与第二个输入blob相同的形状（保持批大小和通道数不变）
    def getMemoryShapes(self, inputs):
        # 获取输入形状和目标形状
        inputShape, targetShape = inputs[0], inputs[1]
        # 获取批大小和通道数
        batchSize, numChannels = inputShape[0], inputShape[1]

        height_in, width_in = inputShape[2], inputShape[3]  # 获取输入的高度和宽度
        height_out, width_out = targetShape[2], targetShape[3]  # 获取目标形状的高度和宽度

        # 居中裁剪方案
        self.ystart = (height_in - height_out) // 2
        self.xstart = (width_in - width_out) // 2

        # 计算裁剪的结束位置
        self.yend = self.ystart + height_out
        self.xend = self.xstart + width_out
        # 返回裁剪后的形状
        return [[batchSize, numChannels, height_out, width_out]]

    # 前向传播函数，进行裁剪操作
    def forward(self, inputs):
        return [inputs[0][:, :, self.ystart:self.yend, self.xstart:self.xend]]


# 使用cv2.dnn_registerLayer函数将自定义裁剪层注册到网络中
cv2.dnn_registerLayer('Crop', CropLayer)