Function ssitrim_loss

model/midas_loss.py:4–97

pred_depth shape: [B, V, 1, H, W] target_depth shape: [B, V, 1, H, W] impl according to MiDaS paper Eq.6 and Splatter-a-Video paper Eq.4 Args: pred_depth: predicted depth target_depth: ground truth depth mask: validity mask ignore_large_loss: threshol

(pred_depth, target_depth, mask=None, ignore_large_loss=0.0, smooth_transition=True, range_utilization=0.0, eps=1e-8, inverse_depth=True)

from einops import rearrange

def ssitrim_loss(pred_depth, target_depth, mask=None, ignore_large_loss=0.0, smooth_transition=True, range_utilization=0.0, eps=1e-8, inverse_depth=True):
    """
    pred_depth shape: [B, V, 1, H, W]
    target_depth shape: [B, V, 1, H, W]
    impl according to MiDaS paper Eq.6 and Splatter-a-Video paper Eq.4
    Args:
        pred_depth: predicted depth
        target_depth: ground truth depth
        mask: validity mask
        ignore_large_loss: threshold for filtering large losses
        smooth_transition: whether to use smooth transition between small and large losses
        range_utilization: weight for encouraging better utilization of depth range
        eps: small constant for numerical stability
    """
    # (B, H * W)
    if mask is None:
        mask = torch.ones_like(target_depth)

    if pred_depth.dim() == 5:
        # seperate depth
        # pred_depth = rearrange(pred_depth, 'B V C H W -> (B V) C H W')
        # target_depth = rearrange(target_depth, 'B V C H W -> (B V) C H W')
        # mask = rearrange(mask, 'B V C H W -> (B V) C H W')
        # unify depth across all views
        pred_depth = rearrange(pred_depth, 'B V C H W -> B (V C) H W')
        target_depth = rearrange(target_depth, 'B V C H W -> B (V C) H W')
        mask = rearrange(mask, 'B V C H W -> B (V C) H W')
    
    # Add range utilization encouragement
    if range_utilization > 0:
        # Calculate min and max depth values
        min_depth = pred_depth.min(dim=1, keepdim=True)[0]
        max_depth = pred_depth.max(dim=1, keepdim=True)[0]
        
        # Encourage min to be close to 0 and max to be close to 1
        min_penalty = torch.relu(min_depth)  # Penalize if min > 0
        max_penalty = torch.relu(1 - max_depth)  # Penalize if max < 1
        
        # Encourage spread by penalizing small ranges
        range_size = max_depth - min_depth
        range_penalty = torch.relu(0.5 - range_size)  # Penalize if range < 0.5
        
        # Combine penalties
        utilization_loss = (min_penalty + max_penalty + range_penalty).mean() * range_utilization
    else:
        utilization_loss = 0.0
    
    if inverse_depth:
        # Convert to inverse depth
        pred_depth = 1.0 / (pred_depth + eps)
    
    pred_depth, target_depth = pred_depth * mask.float(), target_depth * mask.float()

    pred_depth, target_depth = pred_depth.flatten(1), target_depth.flatten(1)

    pix_num = pred_depth.shape[1]
    
    pred_t = torch.median(pred_depth.float(), dim=1).values