import torch
from torch import nn
import numpy as np
import math

def create_1D_gaussian_kernel(sigma: float, size: int) -> torch.Tensor:
    kernel = np.fromfunction(
        lambda x: 1/ (math.sqrt(2 * math.pi * sigma ** 2)) *
        math.e ** ((-(x - size // 2) ** 2) / (2 * sigma ** 2)), [size])
    kernel = torch.Tensor(kernel)
    return kernel

def separable_convolution(input: torch.Tensor, kernel: torch.Tensor) -> torch.Tensor:
    # 行卷积
    conv1 = nn.functional.conv2d(input, kernel.unsqueeze(0).unsqueeze(0), bias=None, stride=1, padding=0)
    # 列卷积
    conv2 = nn.functional.conv2d(conv1, kernel.unsqueeze(1).unsqueeze(1), bias=None, stride=1, padding=0)
    return conv2

def gaussian_blur_3d(input: torch.Tensor, sigma: float, size: int) -> torch.Tensor:
    kernel = create_1D_gaussian_kernel(sigma, size)
    return separable_convolution(input, kernel)

import torch
import torchvision
from torchvision.transforms.functional import gaussian_blur

# 创建一个3D tensor
input_tensor = torch.randn(1, 3, 64, 64, 64)

# 使用Pytorch的内置函数进行3D高斯模糊（只支持3D tensor）
blurred_tensor = gaussian_blur(input_tensor, kernel_size=3, sigma=(1.5, 1.5, 1.5))

# 使用我们自己实现的可分离2D高斯模糊函数进行3D高斯模糊
blurred_tensor = gaussian_blur_3d(input_tensor, sigma=1.5, size=3)