laplacian_smoothing.py

import torch
import torch.nn as nn

def laplacian_smoothing_loss2d(Y_pred):
    C, H, W = Y_pred.shape[1], Y_pred.shape[2], Y_pred.shape[3]
    kernel = torch.tensor([[0, 1, 0], [1, -4, 1], [0, 1, 0]], device=device, dtype=torch.float32)
    kernel = kernel.view(1, 1, 3, 3).repeat(C, 1, 1, 1)
    Y_laplacian = nn.functional.conv2d(Y_pred, kernel, groups=C, padding=1)
    laplacian_loss = (Y_laplacian ** 2).mean()
    return laplacian_loss
 

def laplacian_smoothing_3d(Y_pred):
    C, D, H, W = Y_pred.shape[1], Y_pred.shape[2], Y_pred.shape[3], Y_pred.shape[4]
    kernel = torch.tensor([[[[0, 0, 0], [0, 1, 0], [0, 0, 0]],
                            [[0, 1, 0], [1, -6, 1], [0, 1, 0]],
                            [[0, 0, 0], [0, 1, 0], [0, 0, 0]]]],
                          device=device, dtype=torch.float32)
    kernel = kernel.view(1, 1, 3, 3, 3).repeat(C, 1, 1, 1, 1)
    Y_laplacian = nn.functional.conv3d(Y_pred, kernel, groups=C, padding=1)
    laplacian_loss = (Y_laplacian ** 2).mean()
    return laplacian_loss