Function ssimse_loss

model/midas_loss.py:99–152

Scale and translation invariant MSE loss Args: pred_depth: predicted depth [B, 1, H, W] target_depth: ground truth depth [B, 1, H, W] mask: validity mask [B, 1, H, W] ignore_large_loss: threshold for filtering large losses eps: small constant for nume

(pred_depth, target_depth, mask=None, ignore_large_loss=0.0, eps=1e-8)

    return loss_val + utilization_loss

def ssimse_loss(pred_depth, target_depth, mask=None, ignore_large_loss=0.0, eps=1e-8):
    """
    Scale and translation invariant MSE loss
    Args:
        pred_depth: predicted depth [B, 1, H, W]
        target_depth: ground truth depth [B, 1, H, W]
        mask: validity mask [B, 1, H, W]
        ignore_large_loss: threshold for filtering large losses
        eps: small constant for numerical stability
    """
    if mask is None:
        mask = torch.ones_like(target_depth)
    
    # Apply mask and flatten
    mask = mask.float()
    pred_depth = pred_depth * mask
    target_depth = target_depth * mask
    
    pred_depth = pred_depth.flatten(1)  # [B, H*W]
    target_depth = target_depth.flatten(1)
    mask_flat = mask.flatten(1)
    
    # Compute means on valid pixels only
    valid_pixels = mask_flat.sum(1) + eps
    gt_mean = (target_depth * mask_flat).sum(1) / valid_pixels
    pred_mean = (pred_depth * mask_flat).sum(1) / valid_pixels
    
    # Center the depths
    pred_centered = pred_depth - pred_mean[:, None] 
    target_centered = target_depth - gt_mean[:, None]
    
    # Compute scale factor (least squares)
    numerator = (pred_centered * target_centered * mask_flat).sum(1)
    denominator = (target_centered**2 * mask_flat).sum(1) + eps
    
    # Clamp scale factor for stability
    s = (numerator / denominator).clamp(-10, 10)
    t = pred_mean - s * gt_mean
    
    # Apply scale and translation
    pred_aligned = s[:, None] * pred_depth + t[:, None]
    
    # Compute loss on valid pixels only
    delta = (pred_aligned - target_depth) * mask_flat
    
    if ignore_large_loss > 0:
        valid_mask = ((delta ** 2) < ignore_large_loss) & (mask_flat > 0)
        if valid_mask.any():
            delta = delta[valid_mask]
        else:
            return torch.tensor(0.0, device=pred_depth.device)
    
    loss_val = (delta ** 2).mean()
    return loss_val