src/utils/loss.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
from math import exp

class MultiScaleDerivativeLoss(nn.Module):
    def __init__(self, operator='scharr', p=1, reduction='mean', normalize_input=False, num_scales=4):
        """
        operator: 'scharr' (一阶) or 'laplace' (二阶)
        p: 1 for L1, 2 for L2
        reduction: 'mean' or 'sum'
        normalize_input: whether to normalize input vectors (for normals)
        num_scales: number of scales in the pyramid (e.g., 4 = 原图, 1/2, 1/4, 1/8)
        """
        super().__init__()
        assert operator in ['scharr', 'laplace']
        assert p in [1, 2]
        assert reduction in ['mean', 'sum']
        assert num_scales >= 1

        self.operator = operator
        self.p = p
        self.reduction = reduction
        self.normalize_input = normalize_input
        self.num_scales = num_scales

    def forward(self, pred, gt):
        """
        pred, gt: [B, C, H, W] tensors
        """
        pred_pyramid = self._build_pyramid(pred)
        gt_pyramid = self._build_pyramid(gt)

        total_loss = 0.0

        for pred_i, gt_i in zip(pred_pyramid, gt_pyramid):
            if self.normalize_input:
                pred_i = F.normalize(pred_i, dim=1)
                gt_i = F.normalize(gt_i, dim=1)

            grad_pred = self._compute_gradient(pred_i)
            grad_gt = self._compute_gradient(gt_i)

            diff = grad_pred - grad_gt
            if self.p == 1:
                diff = torch.abs(diff)
            else:
                diff = diff ** 2

            if self.reduction == 'mean':
                total_loss += diff.mean()
            else:
                total_loss += diff.sum()

        return total_loss / self.num_scales

    def _build_pyramid(self, img):
        """Construct a multi-scale pyramid from input image"""
        pyramid = [img]
        for i in range(1, self.num_scales):
            scale = 0.5 ** i
            img = F.interpolate(img, scale_factor=scale, mode='bicubic', align_corners=False, recompute_scale_factor=True,antialias=True)
            pyramid.append(img)
        return pyramid

    def _compute_gradient(self, img):
        B, C, H, W = img.shape
        device = img.device

        if self.operator == 'scharr':
            kernel_x = torch.tensor([[[-3., 0., 3.],
                                      [-10., 0., 10.],
                                      [-3., 0., 3.]]], device=device) / 16.0
            kernel_y = torch.tensor([[[-3., -10., -3.],
                                      [0., 0., 0.],
                                      [3., 10., 3.]]], device=device) / 16.0
            kernel_x = kernel_x.unsqueeze(0).expand(C, 1, 3, 3)
            kernel_y = kernel_y.unsqueeze(0).expand(C, 1, 3, 3)

            grad_x = F.conv2d(img, kernel_x, padding=1, groups=C)
            grad_y = F.conv2d(img, kernel_y, padding=1, groups=C)
            return torch.cat([grad_x, grad_y], dim=1)  # [B, 2C, H, W]

        elif self.operator == 'laplace':
            kernel = torch.tensor([[[0., 1., 0.],
                                    [1., -4., 1.],
                                    [0., 1., 0.]]], device=device)
            kernel = kernel.unsqueeze(0).expand(C, 1, 3, 3)
            return F.conv2d(img, kernel, padding=1, groups=C)  # [B, C, H, W]

class CosineLoss(torch.nn.Module):
    def __init__(self):
        super(CosineLoss, self).__init__()

    def forward(self, N, N_hat):
        """
        N: 真实法向量, 形状 (B, C, H, W) 
        N_hat: 预测法向量, 形状应与 N 相同
        """
        # 创建非零 mask（按像素维度求L2范数）
        _,_,H,W = N.shape
        mask = (N.norm(p=2, dim=1, keepdim=True) > 0)  # shape: (B, 1, H, W)，True表示N非零
        mse = F.mse_loss(N, N_hat, reduction='mean') * H * W /2048 
        dot_product = torch.sum(N * N_hat, dim=1, keepdim=True)  # shape: (B, 1, H, W)
    
        # 仅在非零区域计算 loss
        loss = 1 - dot_product
        loss = loss[mask]  # 只取非零像素位置
        return loss.mean(), mse
    





def gaussian(window_size, sigma):
    gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
    return gauss/gauss.sum()

def create_window(window_size, channel):
    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
    return window

def _ssim(img1, img2, window, window_size, channel, size_average = True, stride=None):
    mu1 = F.conv2d(img1, window, padding = (window_size-1)//2, groups = channel, stride=stride)
    mu2 = F.conv2d(img2, window, padding = (window_size-1)//2, groups = channel, stride=stride)

    mu1_sq = mu1.pow(2)
    mu2_sq = mu2.pow(2)
    mu1_mu2 = mu1*mu2

    sigma1_sq = F.conv2d(img1*img1, window, padding = (window_size-1)//2, groups = channel, stride=stride) - mu1_sq
    sigma2_sq = F.conv2d(img2*img2, window, padding = (window_size-1)//2, groups = channel, stride=stride) - mu2_sq
    sigma12 = F.conv2d(img1*img2, window, padding = (window_size-1)//2, groups = channel, stride=stride) - mu1_mu2

    C1 = 0.01**2
    C2 = 0.03**2

    ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))

    if size_average:
        return ssim_map.mean()
    else:
        return ssim_map.mean(1).mean(1).mean(1)

class SSIM(torch.nn.Module):
    def __init__(self, window_size = 3, size_average = True, stride=3):
        super(SSIM, self).__init__()
        self.window_size = window_size
        self.size_average = size_average
        self.channel = 1
        self.stride = stride
        self.window = create_window(window_size, self.channel)

    def forward(self, img1, img2):
        """
        img1, img2: torch.Tensor([b,c,h,w])
        """
        (_, channel, _, _) = img1.size()

        if channel == self.channel and self.window.data.type() == img1.data.type():
            window = self.window
        else:
            window = create_window(self.window_size, channel)

            if img1.is_cuda:
                window = window.cuda(img1.get_device())
            window = window.type_as(img1)

            self.window = window
            self.channel = channel


        return _ssim(img1, img2, window, self.window_size, channel, self.size_average, stride=self.stride)


def ssim(img1, img2, window_size = 11, size_average = True):
    (_, channel, _, _) = img1.size()
    window = create_window(window_size, channel)

    if img1.is_cuda:
        window = window.cuda(img1.get_device())
    window = window.type_as(img1)

    return _ssim(img1, img2, window, window_size, channel, size_average)




class S3IM(torch.nn.Module):
    def __init__(self, kernel_size=4, stride=4, repeat_time=10, patch_height=64, patch_width=32):
        super(S3IM, self).__init__()
        self.kernel_size = kernel_size
        self.stride = stride
        self.repeat_time = repeat_time
        self.patch_height = patch_height
        self.patch_width = patch_width
        self.ssim_loss = SSIM(window_size=self.kernel_size, stride=self.stride)

    def forward(self, src_vec, tar_vec):
        """
        Args:
            src_vec: [B, N, C] e.g., [batch, pixels, channels]
            tar_vec: [B, N, C]
        Returns:
            loss: scalar tensor
        """
        B, N, C = src_vec.shape
        device = src_vec.device
        patch_list_src, patch_list_tar = [], []

        for b in range(B):
            index_list = []
            for i in range(self.repeat_time):
                if i == 0:
                    tmp_index = torch.arange(N, device=device)
                else:
                    tmp_index = torch.randperm(N, device=device)
                index_list.append(tmp_index)

            res_index = torch.cat(index_list)  # [M * N]
            tar_all = tar_vec[b][res_index]    # [M*N, C]
            src_all = src_vec[b][res_index]    # [M*N, C]

            # reshape into [1, C, H, W]
            tar_patch = tar_all.permute(1, 0).reshape(1, C, self.patch_height, self.patch_width * self.repeat_time)
            src_patch = src_all.permute(1, 0).reshape(1, C, self.patch_height, self.patch_width * self.repeat_time)

            patch_list_tar.append(tar_patch)
            patch_list_src.append(src_patch)

        # Stack all batches: [B, C, H, W]
        tar_tensor = torch.cat(patch_list_tar, dim=0)
        src_tensor = torch.cat(patch_list_src, dim=0)

        # 计算 batch-wise SSIM，输出为 [B]
        ssim_scores = self.ssim_loss(src_tensor, tar_tensor)

        # 损失为 1 - mean SSIM
        loss = 1.0 - ssim_scores
        return loss



torch.manual_seed(0)

# 假设每张图片提取出 64 x 64 个像素，每个像素 3 通道
# H, W, C = 64, 32, 3
# N = H * W
# B = 4
# # 随机生成两个图像特征向量：[N, C]
# src_vec = torch.rand(B, N, C)  # 模拟重建图像
# tar_vec = torch.rand(B, N, C)  # 模拟 ground truth 图像

# # 初始化 S3IM 模块
# s3im_loss_fn = S3IM(kernel_size=4, stride=4, repeat_time=10, patch_height=64, patch_width=32)

# # 计算损失
# loss = s3im_loss_fn(src_vec, tar_vec)

def weighted_huber_loss(
    input: torch.Tensor,
    target: torch.Tensor,
    weight: torch.Tensor,          # 新增的置信度权重张量
    reduction: str = 'mean',
    delta: float = 1.0,
) -> torch.Tensor:    
    # 广播对齐所有张量
    expanded_input, expanded_target = torch.broadcast_tensors(input, target)
    expanded_weight, _ = torch.broadcast_tensors(weight, input)  # 确保权重可广播
    
    # 计算逐元素误差
    diff = expanded_input - expanded_target
    abs_diff = torch.abs(diff)
    
    # Huber损失分段计算
    loss = torch.where(
        abs_diff <= delta,
        0.5 * (diff ** 2),
        delta * (abs_diff - 0.5 * delta)
    )
    
    # 应用权重
    weighted_loss = expanded_weight * loss
    
    # 汇总方式
    if reduction == 'mean':
        return torch.mean(weighted_loss)
    elif reduction == 'sum':
        return torch.sum(weighted_loss)
    elif reduction == 'none':
        return weighted_loss
    else:
        raise ValueError(f"Unsupported reduction: {reduction}")