# 引入必要的库
import cv2
from tqdm import tqdm  # tqdm库用于创建进度条
import numpy as np  # numpy库，用于数值计算，简称为np
import time  # time库，用于时间操作，例如延时
import matplotlib.pyplot as plt
from skimage.metrics import mean_squared_error

# 定义一个名为CropLayer的类
class CropLayer(object):
    # 构造函数，用于初始化裁剪的起始和结束位置
    def __init__(self, params, blobs):
        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)

# 该函数使用预训练的 HED (Holistically-Nested Edge Detection) 模型来检测图像中的边缘。
def get_hed(img_rgb, blur_window=6, 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]
    # 创建blob(一大块数据)
    blob = cv2.dnn.blobFromImage(img_rgb, scalefactor=scalefactor, size=(w,h),mean=(105, 117, 123),swapRB=False, crop=False)

    # 将blob设置为网络的输入
    net.setInput(blob)

    # 执行前向传递以获得边缘检测结果
    hed_output = net.forward()

    # 从输出中提取边缘图
    edge_map = hed_output[0, 0]

    # 输出结果
    hed = edge_map
    return hed

# 该函数通过比较输入的灰度图像与 HED 边缘检测结果之间的均方误差 (MSE)，自动选择最佳的 Canny 边缘检测阈值参数。使用最佳阈值参数计算最终的边缘检测结果。
def auto_canny(img_gray, hed):
    # 对图像进行高斯模糊处理，以降低噪声对边缘检测的影响
    img_gray = cv2.GaussianBlur(img_gray, (5, 5), 0)
    # 定义sigma值范围
    sigma_values = np.linspace(0.001, 1.0, 100)
    mse_values = []  # 用于存储每个sigma值对应的MSE

    # 初始化最佳参数
    best_sigma = 0
    min_mse = float('inf')
    lower = 0
    upper = 0

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

    # 定义函数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
    # 遍历sigma值
    for sigma in tqdm(sigma_values, desc="搜索最佳sigma: ", unit="sigma"):
        # 根据当前sigma值计算Canny边缘检测的阈值参数
        lower, upper = get_bound(sigma)
        
        # 应用Canny算法进行边缘检测
        edge_img = cv2.Canny(img_gray, lower, upper)
        
        # 计算Canny边缘检测结果与HED结果之间的均方误差
        mse = mean_squared_error(edge_img, hed)
        mse_values.append(mse)

        # 更新最佳参数
        if mse < min_mse:
            min_mse = mse
            best_sigma = sigma

        # 在每次迭代中引入延时（例如0.1秒），以便进度条不会过快完成
        time.sleep(0.1)

    # 返回结果
    lower = int(max(0, (1.0 - best_sigma) * median))
    upper = int(min(255, (1.0 + best_sigma) * median))
    best_params = (best_sigma,min_mse,lower,upper)
    plot_data = (sigma_values,mse_values) 
    return edge_img,(best_sigma,min_mse,lower,upper),(sigma_values,mse_values)