Edit2Perceive/utils/cycle_loss.py

import torch
import numpy as np
import torch.nn.functional as F
import math
import torch.nn as nn

def _compute_sobel_gradients_chunked(image, chunk_size=4):
    """
    Computes Sobel gradients for a batch of images in chunks to save memory.
    
    Args:
        image: Tensor of shape [B, H, W]
        chunk_size: The size of each chunk to process.
    
    Returns:
        tuple: (grad_x, grad_y), each of shape [B, H, W]
    """
    # Sobel kernels
    sobel_x_kernel = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], 
                                  dtype=image.dtype, device=image.device).view(1, 1, 3, 3)
    sobel_y_kernel = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], 
                                  dtype=image.dtype, device=image.device).view(1, 1, 3, 3)

    batch_size = image.shape[0]
    # 如果 batch_size 小于等于 chunk_size，就不需要分块
    if batch_size <= chunk_size:
        image_unsqueezed = image.unsqueeze(1)
        grad_x = F.conv2d(image_unsqueezed, sobel_x_kernel, padding=1).squeeze(1)
        grad_y = F.conv2d(image_unsqueezed, sobel_y_kernel, padding=1).squeeze(1)
        return grad_x, grad_y

    # --- 分块计算 ---
    grads_x_list = []
    grads_y_list = []
    # 使用 torch.no_grad() 上下文可以进一步节省显存，因为我们不需要为卷积操作本身计算梯度
    # Sobel核是固定的，不需要梯度。输入image的梯度会通过torch自动追踪。
    with torch.no_grad():
        for i in range(0, batch_size, chunk_size):
            chunk = image[i:i + chunk_size]
            chunk_unsqueezed = chunk.unsqueeze(1)
            
            # 对每个块进行卷积
            chunk_grad_x = F.conv2d(chunk_unsqueezed, sobel_x_kernel, padding=1)
            chunk_grad_y = F.conv2d(chunk_unsqueezed, sobel_y_kernel, padding=1)
            
            grads_x_list.append(chunk_grad_x)
            grads_y_list.append(chunk_grad_y)
    
    # 合并结果
    grad_x = torch.cat(grads_x_list, dim=0).squeeze(1)
    grad_y = torch.cat(grads_y_list, dim=0).squeeze(1)
    
    return grad_x, grad_y

# def get_cycle_consistency_depth_loss(pred_depth, gt_depth, gt_mask=None, depth_normalization="log",eps=1e-6):
#     """
#     Cycle consistency depth loss based on AbsRel metric (Optimized batch processing version)
    
#     Args:
#         pred_depth: [B,C,H,W] or [B,1,H,W] - predicted depth maps
#         gt_depth: [B,C,H,W] or [B,1,H,W] - ground truth depth maps
#         gt_mask: [B,H,W] or [B,1,H,W] - mask for valid GT positions (optional)
#         eps: float - small value to avoid division by zero
#         depth_normalization: str - depth normalization method
    
#     Returns:
#         torch.Tensor - scalar loss value for backpropagation
#     """
    
#     # Ensure consistent shape -> [B, H, W]
#     if pred_depth.dim() == 4 and pred_depth.shape[1] != 1:
#         pred_depth = torch.mean(pred_depth, dim=1, keepdim=False)
#     elif pred_depth.dim() == 4:
#         pred_depth = pred_depth.squeeze(1)
    
#     if gt_depth.dim() == 4 and gt_depth.shape[1] != 1:
#         gt_depth = torch.mean(gt_depth, dim=1, keepdim=False)
#     elif gt_depth.dim() == 4:
#         gt_depth = gt_depth.squeeze(1)
    
#     # Handle shape mismatch by resizing prediction to match GT
#     if pred_depth.shape != gt_depth.shape:
#         pred_depth = F.interpolate(pred_depth.unsqueeze(1), size=gt_depth.shape[-2:], 
#                                  mode='bilinear', align_corners=False).squeeze(1)
    
#     # Handle invalid values in GT and prediction (batch processing)
#     gt_depth = torch.where(torch.isinf(gt_depth), torch.zeros_like(gt_depth), gt_depth)
#     gt_depth = torch.where(torch.isnan(gt_depth), torch.zeros_like(gt_depth), gt_depth)
    
#     pred_depth = torch.where(torch.isinf(pred_depth), torch.zeros_like(pred_depth), pred_depth)
#     pred_depth = torch.where(torch.isnan(pred_depth), torch.zeros_like(pred_depth), pred_depth)
    
#     # Create valid mask
#     if gt_mask is not None:
#         if gt_mask.dim() == 4:
#             gt_mask = gt_mask.squeeze(1)  # [B, H, W]
#         valid_mask = gt_mask.bool()
#     else:
#         valid_mask = torch.ones_like(gt_depth, dtype=torch.bool, device=gt_depth.device)
    
#     # Clamp depths to valid range
#     pred_depth = torch.clamp(pred_depth, -1.0, 1.0)
#     gt_depth = torch.clamp(gt_depth, -1.0, 1.0)
    
#     # Apply depth normalization (batch processing)
#     if depth_normalization == "log":
#         pred_depth = torch.exp(pred_depth)
#         gt_depth = torch.exp(gt_depth)
#     elif depth_normalization == "sqrt":
#         pred_depth = pred_depth**2
#         gt_depth = gt_depth**2
#     elif depth_normalization == "disp":
#         pred_depth = 1.0/(pred_depth + eps)
#         gt_depth = 1.0/(gt_depth + eps)
#     elif depth_normalization == "sqrt_disp":
#         pred_depth = 1.0/(pred_depth**2 + eps)
#         gt_depth = 1.0/(gt_depth**2 + eps)
#     elif depth_normalization == "uniform":
#         pass # already in the same scale as real relative depth
#     else:
#         raise ValueError(f"Unknown depth normalization: {depth_normalization}")
    
#     # Compute AbsRel metric for each sample in the batch
#     pred_depth = torch.log(pred_depth + eps)  # [B, H, W]
#     gt_depth = torch.log(gt_depth + eps)      # [B, H, W
#     abs_diff = torch.abs(pred_depth - gt_depth)  # [B, H, W]
#     rel_diff = abs_diff
#     # rel_diff = abs_diff / (gt_depth + eps)       # [B, H, W], add eps to avoid division by zero
    
#     # Apply mask and compute mean for each sample
#     masked_rel_diff = rel_diff * valid_mask.float()  # [B, H, W]
    
#     # Sum over spatial dimensions and divide by number of valid pixels per sample
#     valid_pixel_count = valid_mask.sum(dim=[1, 2]).float()  # [B]
#     sample_losses = masked_rel_diff.sum(dim=[1, 2])         # [B]
    
#     # Avoid division by zero for samples with no valid pixels
#     valid_samples_mask = valid_pixel_count > 0
#     sample_losses = sample_losses / torch.clamp(valid_pixel_count, min=1.0)
    
#     # Only consider samples with valid pixels
#     if valid_samples_mask.any():
#         final_loss = sample_losses[valid_samples_mask].mean()
#     else:
#         # If no valid samples, return zero loss
#         final_loss = torch.tensor(0.0, device=pred_depth.device, requires_grad=True)
    
#     return final_loss

class ScaleAndShiftInvariantLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.name = "SSILoss"
    def forward(self, prediction, target, mask,depth_normalization="log"):
        # if mask.ndim ==3:
        #     mask = mask.unsqueeze(1)  # [B,1,H,W]
        if prediction.ndim == 4:
            prediction = torch.mean(prediction, dim=1, keepdim=False) # [B,H,W]
        if target.ndim ==4:
            target = torch.mean(target, dim=1, keepdim=False) # [B,H,W]
        mask = mask.bool()
        with torch.autocast(device_type='cuda', enabled=False):
            prediction = prediction.float()
            target     = target.float()
            if depth_normalization == "log":
                target_cp = torch.log(target).requires_grad_(False)
            elif depth_normalization == "sqrt":
                target_cp = torch.sqrt(target).requires_grad_(False)
            scale, shift = compute_scale_and_shift_masked(prediction, target_cp, mask)
            del target_cp
            scaled_prediction = scale.view(-1, 1, 1) * prediction + shift.view(-1, 1, 1)
            loss = nn.functional.l1_loss(scaled_prediction[mask], target[mask])
        return loss
    
def compute_scale_and_shift_masked(prediction, target, mask):
    # system matrix: A = [[a_00, a_01], [a_10, a_11]]
    a_00 = torch.sum(mask * prediction * prediction, (1, 2))
    a_01 = torch.sum(mask * prediction, (1, 2))
    a_11 = torch.sum(mask, (1, 2))
    # right hand side: b = [b_0, b_1]
    b_0 = torch.sum(mask * prediction * target, (1, 2))
    b_1 = torch.sum(mask * target, (1, 2))
    # solution: x = A^-1 . b = [[a_11, -a_01], [-a_10, a_00]] / (a_00 * a_11 - a_01 * a_10) . b
    x_0 = torch.zeros_like(b_0)
    x_1 = torch.zeros_like(b_1)
    det = a_00 * a_11 - a_01 * a_01
    # A needs to be a positive definite matrix.
    valid = det > 0 #1e-3
    x_0[valid] = (a_11[valid] * b_0[valid] - a_01[valid] * b_1[valid]) / det[valid]
    x_1[valid] = (-a_01[valid] * b_0[valid] + a_00[valid] * b_1[valid]) / det[valid]
    return x_0, x_1

def get_cycle_consistency_depth_loss(pred, gt, mask=None, depth_normalization="log",eps=1e-8,
                   alpha=0.85, # SILog 损失的权重
                   beta=1.0,   # 梯度损失的权重
                   gamma=0.0,
                   per_sample_weights=None): # 可选的法线一致性损失的权重
    """
    一个专注于几何结构的、稳定版的深度损失函数。
    它结合了尺度不变损失(SILog)和局部结构损失(梯度匹配)。
    Args:
        pred: [B,3,H,W] 或 [B,H,W,3], 范围 [-1, 1]
        gt:   [B,3,H,W] 或 [B,H,W,3], 范围 [-1, 1]
        mask: None 或 [B,H,W] / [B,1,H,W] / [B,H,W,1]
        eps:  数值稳定项
        alpha: SILog 损失中的方差项权重 (通常 0.5 或 0.85)
        beta: 梯度损失的权重
        gamma: 法线一致性损失的权重 (如果>0，会计算法线损失)
    Returns:
        单标量 tensor（可反传）
    """
    # --- 形状、device、dtype 规范化 ---
    if pred.dim() == 4 and pred.shape[1] != 3 and pred.shape[-1] == 3:
        pred = pred.permute(0, 3, 1, 2)
    if gt.dim() == 4 and gt.shape[1] != 3 and gt.shape[-1] == 3:
        gt = gt.permute(0, 3, 1, 2)

    B, C, H, W = pred.shape
    device = pred.device
    dtype = pred.dtype
    # if per_sample_weights is None:
    #     per_sample_weights = torch.ones((B,), device=device, dtype=dtype)
    # --- 1. 预处理 ---
    # 取3通道均值得到单通道深度图
    pred = pred.mean(dim=1,keepdim=True) # [B,1,H,W]
    gt   = gt.mean(dim=1,keepdim=True)   # [B,1,H,W]

    # 将 [-1, 1] 范围映射到正数范围 (e.g., [eps, 1+eps])，为 log 做准备
    # 这一步非常重要，因为它将归一化的值转回到了一个类似深度的正数空间
    pred = (pred.clamp(min=-1, max=1) + 1.0) / 2.0 + eps
    gt   = (gt.clamp(min=-1, max=1) + 1.0) / 2.0 + eps
    # if depth_normalization == "log":
    #     pred = torch.exp(pred)
    #     gt = torch.exp(gt)
    # elif depth_normalization == "sqrt":
    #     pred = pred**2
    #     gt = gt**2
    # elif depth_normalization == "disp":
    #     pred = 1.0/(pred + eps)
    #     gt = 1.0/(gt + eps)
    # elif depth_normalization == "sqrt_disp":
    #     pred = 1.0/(pred**2 + eps)
    #     gt = 1.0/(gt**2 + eps)
    # else:
    #     # uniform 
    #     pass # already in the same scale as real relative depth
    # --- mask 处理 ---
    if mask is None:
        # 如果没有mask，创建一个全为True的mask
        mask = torch.ones_like(pred, dtype=torch.bool, device=device)
    else:
        # 确保mask是 [B,H,W] 的 bool tensor
        if mask.ndim == 4:
            mask = torch.mean(mask, dim=1,keepdim=True).bool()  # [B,1,H,W]
        elif mask.ndim ==3:
            mask = mask.unsqueeze(1).bool()  # [B,1,H,W]
    mask = mask.bool().to(device)
    # 检查有效像素数量
    total_valid = mask.sum()
    if total_valid.item() == 0:
        return torch.zeros((), device=device, dtype=dtype, requires_grad=True)

    # --- 2. 计算复合损失 ---

    # --- Component A: Scale-Invariant Logarithmic (SILog) Loss ---
    # 这是点对点损失的鲁棒版本，对尺度不敏感
    log_diff = torch.log(1+pred[mask]) - torch.log(1+gt[mask]) # B,N

    # # SILog 公式: E[d^2] - alpha * (E[d])^2
    silog_term1 = torch.mean(log_diff ** 2)
    silog_term2 = (torch.mean(log_diff)) ** 2
    loss_silog = silog_term1 - alpha * silog_term2
    loss_silog = (loss_silog).mean()
    # --- Component B: Gradient Matching Loss ---
    # 这是结构损失，确保表面的局部坡度一致
    # 使用简单的 sobel 算子来计算梯度
    sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], device=device, dtype=dtype).view(1, 1, 3, 3)
    sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], device=device, dtype=dtype).view(1, 1, 3, 3)

    pred_grad_x = F.conv2d(pred, sobel_x, padding=1)
    pred_grad_y = F.conv2d(pred, sobel_y, padding=1)
    gt_grad_x = F.conv2d(gt, sobel_x, padding=1)
    gt_grad_y = F.conv2d(gt, sobel_y, padding=1)

    # 计算梯度差异的 L1 损失
    grad_loss_x = torch.abs(pred_grad_x - gt_grad_x)
    grad_loss_y = torch.abs(pred_grad_y - gt_grad_y)

    # 只在 mask 区域计算损失
    # 稍微腐蚀一下mask，因为边界处的梯度可能不准确
    mask_eroded = F.max_pool2d(mask.float(), kernel_size=3, stride=1, padding=1).bool()
    loss_grad = (grad_loss_x[mask_eroded].mean() + grad_loss_y[mask_eroded].mean())
    # loss_grad = ((grad_loss_x[mask_eroded].view(B, -1).mean(dim=1) + grad_loss_y[mask_eroded].view(B, -1).mean(dim=1)) * per_sample_weights).mean()

    # --- Component C (Optional): Normal Consistency Loss ---
    # 如果 gamma > 0，则计算从深度图重建的法线之间的一致性损失
    loss_normal = torch.tensor(0.0, device=device, dtype=dtype)
    if gamma > 0:
        # 从深度图计算法线 (使用梯度)
        # 法线 n = [-gx, -gy, 1]，然后归一化
        # 注意：这里假设相机内参 f_x, f_y = 1
        ones = torch.ones_like(pred)
        pred_normal = torch.cat([-pred_grad_x, -pred_grad_y, ones], dim=1)
        gt_normal   = torch.cat([-gt_grad_x,   -gt_grad_y,   ones], dim=1)

        # 使用你成功的法线损失函数
        loss_normal = get_cycle_consistency_normal_loss(
            pred_normal, gt_normal, mask=mask_eroded.squeeze(1),
        )

    # --- 最终损失 ---
    # 加权求和
    # 权重可以根据你的任务进行调整，beta=1.0, gamma=0.0 是一个很好的起点
    total_loss = loss_silog + beta * loss_grad + gamma * loss_normal

    return total_loss


# ===================================================================
# 核心辅助函数：从深度图计算法线
# ===================================================================
# def depth_to_normals(depth, camera_intrinsics=None, eps=1e-8):
#     """
#     从深度图计算表面法线。
#     Args:
#         depth: [B, 1, H, W] 深度图
#         camera_intrinsics: [B, 3, 3] 相机内参矩阵 K。如果为 None，则在像素空间计算。
#     Returns:
#         normals: [B, 3, H, W] 法线图
#     """
#     B, _, H, W = depth.shape
#     device = depth.device
#     dtype = depth.dtype

#     # 创建像素坐标网格
#     y_coords, x_coords = torch.meshgrid(torch.arange(H, device=device, dtype=dtype),
#                                         torch.arange(W, device=device, dtype=dtype),
#                                         indexing='ij')
#     coords = torch.stack([x_coords, y_coords, torch.ones_like(x_coords)], dim=0).unsqueeze(0).repeat(B, 1, 1, 1)

#     if camera_intrinsics is not None:
#         # 反投影到相机坐标系
#         K_inv = torch.inverse(camera_intrinsics).view(B, 3, 3)
#         cam_coords = K_inv @ coords.to(torch.float32).view(B, 3, -1)
#         cam_coords = cam_coords.view(B, 3, H, W)
#         point_cloud = cam_coords * depth
#     else:
#         # 在像素空间计算（一个合理的近似）
#         point_cloud = torch.cat([coords[:, :2, ...], depth], dim=1)

#     # 使用 padding + 卷积计算梯度，更稳定
#     kernel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], device=device, dtype=dtype).view(1, 1, 3, 3) / 8.0
#     kernel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], device=device, dtype=dtype).view(1, 1, 3, 3) / 8.0
    
#     # 对点云的每个分量求梯度
#     point_cloud = point_cloud.to(torch.bfloat16)
#     p_x, p_y, p_z = point_cloud[:, 0:1], point_cloud[:, 1:2], point_cloud[:, 2:3]
    
#     tg_x = torch.cat([F.conv2d(p, kernel_x, padding='same') for p in (p_x, p_y, p_z)], dim=1)
#     tg_y = torch.cat([F.conv2d(p, kernel_y, padding='same') for p in (p_x, p_y, p_z)], dim=1)

#     # 计算法线 (cross product)
#     # tg_y x tg_x 的方向通常指向相机
#     normals = torch.cross(tg_y, tg_x, dim=1)

#     # 单位化
#     normals = normals / normals.norm(dim=1, keepdim=True).clamp_min(eps)
#     return normals


# # ===================================================================
# # 你的全新损失函数！
# # ===================================================================
# def get_cycle_consistency_depth_loss(pred_depth, gt_depth, mask=None, camera_intrinsics=None, depth_normalization="log",eps=1e-8,per_sample_weights=None):
#     """
#     将深度估计问题转化为法线几何一致性问题的损失函数。
#     1. 使用最小二乘法对齐 pred 和 gt_depth。
#     2. 将对齐后的 pred 和 gt_depth 转换为法线图。
#     3. 使用 get_cycle_consistency_normal_loss 计算法线间的损失。

#     Args:
#         pred: [B,3,H,W] or [B,H,W,3], 模型预测的深度图，值域 [-1, 1]
#         input_depth: [B,H,W] or [B,1,H,W], 原始的、未归一化的真实深度图
#         mask: None or [B,H,W], 有效像素区域
#         camera_intrinsics: 可选, [B, 3, 3] 相机内参, 用于更精确的法线计算
#         eps: 数值稳定项

#     Returns:
#         单标量 tensor (可反传)
#     """
#     # --- 形状、device、dtype 规范化 ---
#     device = pred_depth.device
#     dtype = pred_depth.dtype
#     if pred_depth.ndim ==4:
#         pred_depth = torch.mean(pred_depth, dim=1)  # [B, H, W]

#     if gt_depth.ndim == 4:
#         gt_depth = torch.mean(gt_depth, dim=1)  # [B, H, W]
        
#     B, H, W = pred_depth.shape
#     if depth_normalization == "log":
#         pred_depth = (pred_depth + 1.0) / 2.0 + eps
#         pred_depth = torch.exp(pred_depth)
#     elif depth_normalization == "sqrt":
#         pred_depth = (pred_depth + 1.0) / 2.0 + eps
#         pred_depth = pred_depth**2
#     elif depth_normalization == "disp":
#         pred_depth = (pred_depth + 1.0) / 2.0 + eps
#         pred_depth = 1.0/(pred_depth + eps)
#     elif depth_normalization == "sqrt_disp":
#         pred_depth = (pred_depth + 1.0) / 2.0 + eps
#         pred_depth = 1.0/(pred_depth**2 + eps)
        
#     # --- Mask 处理 ---
#     if mask is None:
#         mask = (gt_depth > eps).detach()
#     else:
#         # if mask.dim() == 3:
#             # mask = mask.unsqueeze(1)
#         mask = mask.bool().to(device)

#     # --- 核心步骤 1: 最小二乘法对齐尺度和偏移 ---
#     aligned_pred_depths = []
#     for i in range(B): # 逐个样本处理
#         p = pred_depth[i][mask[i]]
#         g = gt_depth[i][mask[i]]

#         if p.numel() < 2: # 有效点太少，无法拟合
#             aligned_pred_depths.append(pred_depth[i])
#             continue

#         # 构造最小二乘问题 y = s*x + t, A*[s,t]^T = y
#         A = torch.stack([p, torch.ones_like(p)], dim=1)
#         y = g

#         # 使用 torch.linalg.lstsq 求解，稳定且高效
#         try:
#             solution = torch.linalg.lstsq(A.to(torch.float32), y.to(torch.float32)).solution
#             s, t = solution[0], solution[1]
#         except torch.linalg.LinAlgError:
#             # 如果矩阵奇异，退化为只匹配均值
#             s = torch.tensor(1.0, device=device, dtype=dtype)
#             t = g.mean() - p.mean()

#         # 对齐预测值，不让梯度流过 s 和 t
#         aligned_pred_depths.append(pred_depth[i] * s.detach() + t.detach())

#     aligned_pred_depth = torch.stack(aligned_pred_depths)

#     # --- 核心步骤 2: 从对齐后的深度图计算法线 ---
#     # pred_normals = depth_to_normals(aligned_pred_depth, camera_intrinsics, k=5, d=1, gamma=0.05, min_nghbr=4, eps=eps)
#     # gt_normals = depth_to_normals(gt_depth, camera_intrinsics, k=5, d=1, gamma=0.05, min_nghbr=4, eps=eps)
#     # pred_normals = depth_to_normals(aligned_pred_depth, camera_intrinsics,)
#     # gt_normals = depth_to_normals(gt_depth, camera_intrinsics)
    
#     # # --- 核心步骤 3: 复用成功的法线损失 ---
#     # # 注意，法线计算在边缘处可能无效，mask需要结合深度mask
#     # final_mask = mask & (gt_depth > eps) & (pred_normals.norm(dim=1, keepdim=True) > eps) & (gt_normals.norm(dim=1, keepdim=True) > eps)
    
#     # loss = get_cycle_consistency_normal_loss(pred_normals, gt_normals, final_mask)
#     if per_sample_weights is None:
#         per_sample_weights = torch.ones((B,), device=device, dtype=dtype)
#     loss = per_sample_weights.view(B,1,1)*(aligned_pred_depth-gt_depth)**2*mask  # [B,H,W] -> [B]
#     return loss.mean().clip(1e-6,1.0)

import torch
import math

def get_cycle_consistency_normal_loss(pred, gt, mask=None, eps=1e-8, per_sample_weights=None):
    """
    稳定版的法线角度损失（可直接反传）
    Args:
        pred: [B,3,H,W] 或 [B,H,W,3]
        gt:   [B,3,H,W] 或 [B,H,W,3]
        mask: None 或 [B,H,W] / [B,1,H,W] / [B,H,W,1]
        eps:  数值稳定项
        per_sample_weights: None 或 [B]，每个样本的权重
    Returns:
        单标量 tensor（可反传）
    """
    # --- 形状、device、dtype 规范化 ---
    if pred.dim() == 4 and pred.shape[1] != 3 and pred.shape[-1] == 3:
        pred = pred.permute(0, 3, 1, 2)
    if gt.dim() == 4 and gt.shape[1] != 3 and gt.shape[-1] == 3:
        gt = gt.permute(0, 3, 1, 2)

    B, C, H, W = pred.shape
    assert C == 3, "输入必须是3通道法线向量"

    device = pred.device
    dtype = pred.dtype

    # --- 单位化（防止除零） ---
    pred = pred / pred.norm(dim=1, keepdim=True).clamp_min(eps)
    gt   = gt   / gt.norm(dim=1, keepdim=True).clamp_min(eps)

    # --- mask 处理 ---
    if mask is None:
        mask = torch.ones((B, H, W), dtype=torch.bool, device=device)
    if mask.dim() == 4:
        mask = mask.squeeze(1)
    if mask.dim() == 3 and mask.shape[1] == 1:  # 兼容 [B,1,H,W]
        mask = mask.squeeze(1)
    mask = mask.bool().to(device)

    # flatten为 [B,3,N] / [B,N]
    pred = pred.view(B, 3, -1)
    gt   = gt.view(B, 3, -1)
    mask = mask.view(B, -1).bool()  # [B, N]

    # --- 使用 dot 和 cross -> atan2 更稳定 ---
    dot = (pred * gt).sum(dim=1)  # [B, N], 取代 acos 的 cos 值
    cross = torch.cross(pred, gt, dim=1)  # [B,3,N]
    sin_val = cross.norm(dim=1)  # [B, N]

    # clamp 防止极端数值（但不把 dot 强行压到 ±1）
    dot = dot.clamp(-1.0 + 1e-7, 1.0 - 1e-7)
    sin_val = sin_val.clamp_min(1e-12)

    # angle in degrees
    angle = torch.atan2(sin_val, dot) * (180.0 / math.pi)  # [B, N]

    # --- 按 mask 计算 mean / accuracy（避免空mask导致 NaN） ---
    mask_f = mask.float()  # [B, N]
    counts_per_batch = mask_f.sum(dim=1)  # [B]
    # total = counts_per_batch.sum()  # scalar

    # if total.item() == 0:
    #     # 没有有效像素：返回 0 且保留梯度（避免图断裂）
    #     return torch.zeros((), device=device, dtype=dtype, requires_grad=True)
    # # mean angular error over masked elements
    # loss = (angle * mask_f ).sum() / total
    if per_sample_weights is None:
        per_sample_weights = torch.ones((B,), device=device, dtype=dtype)
    loss = (angle * mask_f).sum(dim=1) / counts_per_batch.clamp_min(1)  # [B]
    loss = (loss * per_sample_weights).mean()
    return loss / 100.0  # scale down for稳定（按你原代码保留）


def get_cycle_consistency_matting_loss(pred_alpha, gt_alpha, trimap):
    """
    一个集成的、可微的Matting损失函数 (PyTorch, Batch version)。
    
    此函数将SAD、MSE和Gradient Loss的计算合并在一起，并对未知区域进行操作。
    它被设计为简单、直接，易于集成到训练循环中。

    Args:
        pred_alpha (torch.Tensor): 预测的alpha matte, shape [B, 1, H, W], 值在 [0, 1] 之间。
        gt_alpha (torch.Tensor): 真实的alpha matte, shape [B, 1, H, W], 值在 [0, 1] 之间。
        trimap (torch.Tensor): Trimap, shape [B, 1, H, W], 值为 0, 128, 255。

    Returns:
        torch.Tensor: 一个用于反向传播的标量损失值。
    """
    # --- 1. 预处理: 创建未知区域的掩码 ---
    # 未知区域是trimap中值为128的像素
    if pred_alpha.ndim == 4:
        pred_alpha = torch.mean(pred_alpha, dim=1, keepdim=True)  # [B,1,H,W]
    if gt_alpha.ndim == 4:
        gt_alpha = torch.mean(gt_alpha, dim=1, keepdim=True)      # [B,1,H,W]
    if trimap.ndim == 3:
        trimap = trimap.unsqueeze(1)  # [B,1,H,W]
    unknown_mask = ((trimap >= 127.5) & (trimap <= 128.5)).float()

    # --- 2. SAD Loss (Sum of Absolute Differences) ---
    error_map_sad = torch.abs(pred_alpha - gt_alpha)
    # 对每个样本的未知区域求和，并按原始实现进行缩放
    loss_sad_per_sample = (error_map_sad * unknown_mask).sum(dim=(1, 2, 3)) / 1000.0

    # --- 3. MSE Loss (Mean Squared Error) ---
    error_map_mse = (pred_alpha - gt_alpha) ** 2
    # 对每个样本的未知区域求和，然后除以该区域的像素数
    loss_mse_per_sample = (error_map_mse * unknown_mask).sum(dim=(1, 2, 3)) / (unknown_mask.sum(dim=(1, 2, 3)) + 1e-8)

    # --- 4. Gradient Loss ---
    # 在函数内部动态创建高斯梯度核，以匹配输入张量的设备
    device = pred_alpha.device
    sigma = 1.4
    
    # 创建高斯导数核
    epsilon = 1e-2
    halfsize = math.ceil(sigma * math.sqrt(-2 * math.log(math.sqrt(2 * math.pi) * sigma * epsilon)))
    size = 2 * halfsize + 1
    coords = torch.arange(-halfsize, halfsize + 1, dtype=torch.float32, device=device)
    
    gauss_vals = torch.exp(-coords**2 / (2 * sigma**2)) / (sigma * math.sqrt(2 * math.pi))
    dgauss_vals = -coords * gauss_vals / (sigma**2)
    
    # x和y方向的核
    hx = (gauss_vals.unsqueeze(1) * dgauss_vals.unsqueeze(0))
    hy = hx.t()
    
    # 归一化
    hx = hx / torch.sqrt(torch.sum(torch.abs(hx) * torch.abs(hx)))
    hy = hy / torch.sqrt(torch.sum(torch.abs(hy) * torch.abs(hy)))
    
    # 为conv2d调整形状: [out_channels, in_channels, H, W]
    kernel_x = hx.to(pred_alpha.dtype).unsqueeze(0).unsqueeze(0)
    kernel_y = hy.to(pred_alpha.dtype).unsqueeze(0).unsqueeze(0)

    # 使用卷积计算梯度
    pred_gx = F.conv2d(pred_alpha, kernel_x, padding='same')
    pred_gy = F.conv2d(pred_alpha, kernel_y, padding='same')
    gt_gx = F.conv2d(gt_alpha, kernel_x, padding='same')
    gt_gy = F.conv2d(gt_alpha, kernel_y, padding='same')

    # 计算梯度幅值
    pred_amp = torch.sqrt(pred_gx**2 + pred_gy**2)
    gt_amp = torch.sqrt(gt_gx**2 + gt_gy**2)

    # 计算梯度损失
    error_map_grad = (pred_amp - gt_amp) ** 2
    loss_grad_per_sample = (error_map_grad * unknown_mask).sum(dim=(1, 2, 3)) / 1000.0

    # --- 5. 组合损失并在Batch维度上取平均 ---
    # 你可以根据需要为各项损失添加权重, e.g., 1.0 * loss_sad + 0.5 * loss_mse + ...
    total_loss_per_sample = loss_sad_per_sample + loss_mse_per_sample + loss_grad_per_sample
    final_loss = total_loss_per_sample.mean()

    return final_loss

def get_disperse_loss(eta, tau=1.0):
    """
    计算 Dispersive Loss.
    公式: L_disp = log E_{i,j} [exp(-|| η_i - η_j ||_2^2 / τ)]
    
    参数:
        eta (torch.Tensor): 一个批次的中间特征，形状为 (batch_size, seq_len, dim) 或 (batch_size, features)。
        tau (float): 温度超参数，根据论文设置为 1.0。
    
    返回:
        torch.Tensor: 计算出的 dispersive loss (标量).
    """
    batch_size = eta.shape[0]
    # 如果批次大小小于等于1，无法计算配对距离，损失为0
    if batch_size <= 1:
        return torch.tensor(0.0, device=eta.device, dtype=eta.dtype)

    # 如果特征是序列化的 (B, N, D)，则将其展平为 (B, N*D)
    if eta.dim() > 2:
        eta = eta.view(batch_size, -1)

    # 高效计算批次内样本间的成对L2距离的平方
    # (x-y)^2 = x^2 - 2xy + y^2
    eta_norm_sq = torch.sum(eta**2, dim=1, keepdim=True)
    # dist_sq 是一个 (batch_size, batch_size) 的矩阵，存储了距离的平方
    dist_sq = eta_norm_sq - 2 * torch.mm(eta, eta.t()) + eta_norm_sq.t()
    
    # 确保数值稳定性，距离的平方应为非负数
    dist_sq = torch.clamp(dist_sq, min=0.0)

    # 计算 e^(-d^2/τ)
    loss_matrix = torch.exp(-dist_sq / tau)
    
    # 计算期望 E_{i,j}，即对所有 B*B 配对取均值
    expectation = torch.mean(loss_matrix)
    
    # 最终损失是期望的对数
    loss = torch.log(expectation)
    
    return loss