sobel算子python代码

def apply_sobel_operator(image_tensor):
    # 扩展Sobel算子通道数
    # 将 PyTorch 张量转换为 NumPy 数组
    if has_invalid_values(image_tensor):
        print("Tensor contains invalid values.")
    input_np = (image_tensor.detach().cpu().numpy()*255).astype(np.uint8)

    # 创建一个存储 Sobel 结果的数组
    sobel_results = np.empty_like(input_np)

    # 循环遍历每个图像
    for i in range(input_np.shape[0]):
        # 循环遍历每个通道
        for j in range(input_np.shape[1]):
            # 将彩色图像转换为灰度图像
            gray_image = input_np[i,j].reshape(1,input_np.shape[2],input_np.shape[3])

            # 使用带有 NumPy 数组的 cv2.Sobel
            gradient_x = cv2.Sobel(gray_image, cv2.CV_64F, 1, 0, ksize=3)
            gradient_y = cv2.Sobel(gray_image, cv2.CV_64F, 0, 1, ksize=3)
            # 如果需要，你可以将梯度的计算结果再转回 PyTorch 张量
            gradient_magnitude = cv2.magnitude(gradient_x, gradient_y)
            # gradient_magnitude = (gradient_magnitude - np.min(gradient_magnitude)) / (np.max(gradient_magnitude) - np.min(gradient_magnitude)).astype(np.float64)
            sobel_results[i, j] = gradient_magnitude

    # 现在，sobel_results 是一个形状为 (batch_size, channels, height, width) 的 NumPy 数组，包含 Sobel 算子的结果
    # 如果需要，你可以将其再转回 PyTorch 张量
    sobel_results_tensor = torch.from_numpy(sobel_results).cuda()
    return sobel_results_tensor/255
使用该代码时，注意输入图像像素值域,根据值域不同需要调整代码，否则容易因为精度问题出现nah警告，导致网络梯度信息错误，最好是不要在网络传递过程中使用该方法，如要使用需要自己调整代码