import cv2
import numpy as np 
from tqdm import tqdm
from skimage import img_as_float,img_as_ubyte
from skimage.metrics import mean_squared_error

# 该函数使用预训练的 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):
    
    h,w = img_rgb.shape[:2]

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

    # 指定模型文件的路径
    model_path ='model/hed_pretrained_bsds.caffemodel'
    prototxt_path ='model/deploy.prototxt'
    # 使用OpenCV加载预训练模型
    net = cv2.dnn.readNetFromCaffe(prototxt_path,model_path)

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

    ######### 整体推理(predict)的过程 ##########
    # 将 blob 设置为网络的输入
    net.setInput(blob)
    # 执行前向传递以获得边缘检测结果
    hed_output = net.forward()
    # 从输出中提取边缘图
    edge_map = hed_output[0, 0]

    # 参数 (5,5)用来指定kernel大小，为了生成核实的高斯曲线，请确定其大小为奇数。
    # 参数0用来指定高斯曲线在x轴的标准差，如果定义为0，则表示让代码自动根据kernel高度和宽度自动来计算。
    edge_blur=cv2.GaussianBlur(edge_map,(blur_window, blur_window),0) # opencv 高斯滤波
    edge_blur = cv2.convertScaleAbs(edge_map)
    _, hed = cv2.threshold(edge_blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    # TODO

    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_blur = cv2.GaussianBlur(img_gray, (5,5), 0)

    # 初始化最佳sigma值和min_mse
    best_sigma = None
    min_mse = float('inf')

    # 初始化一个列表，用于存储所有的mse值
    mse_values = []
    # 生成从0.01到0.8的sigma值，步长为0.01
    sigma_values = np.arange(0.01, 1.01, 0.01)

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

    # 遍历sigma值范围，计算每个去噪图像的mse
    for sigma in tqdm(sigma_values, desc=" sigma "):

        lower1 = int(max(0, (1.0 - sigma) * median))
        upper1 = int(min(255, (1.0 + sigma) * median))
        # 应用Canny算法进行边缘检测
        auto_canny = cv2.Canny(img_blur, lower1, upper1)

        # 计算均方误差（MSE）
        mse = mean_squared_error(auto_canny,  hed)
       
        mse_values.append(mse)
        if mse < min_mse:
            min_mse = mse
            best_sigma = sigma

    # 使用最佳sigma值进行去噪
    lower = int(max(0, (1.0 - best_sigma) * median))
    upper = int(min(255, (1.0 + best_sigma) * median))
    edge_img = cv2.Canny(img_blur, lower, upper)


    # TODO

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


########################## 注册新的裁剪层方法 ##############################
# 定义一个名为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.ystart=0 
        # self.xstart=0
        
        # 计算裁剪的结束位置
        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]]