Class ChamferDistance

embodiedscan/models/losses/chamfer_distance.py:83–157

Calculate Chamfer Distance of two sets. Args: mode (str): Criterion mode to calculate distance. The valid modes are 'smooth_l1', 'l1' or 'l2'. Defaults to 'l2'. reduction (str): Method to reduce losses. The valid reduction method are 'none', 'sum' or 'mea

@MODELS.register_module()
class ChamferDistance(nn.Module):
    """Calculate Chamfer Distance of two sets.

    Args:
        mode (str): Criterion mode to calculate distance.
            The valid modes are 'smooth_l1', 'l1' or 'l2'. Defaults to 'l2'.
        reduction (str): Method to reduce losses.
            The valid reduction method are 'none', 'sum' or 'mean'.
            Defaults to 'mean'.
        loss_src_weight (float): Weight of loss_source. Defaults to l.0.
        loss_dst_weight (float): Weight of loss_target. Defaults to 1.0.
    """

    def __init__(self,
                 mode: str = 'l2',
                 reduction: str = 'mean',
                 loss_src_weight: float = 1.0,
                 loss_dst_weight: float = 1.0) -> None:
        super(ChamferDistance, self).__init__()

        assert mode in ['smooth_l1', 'l1', 'l2']
        assert reduction in ['none', 'sum', 'mean']
        self.mode = mode
        self.reduction = reduction
        self.loss_src_weight = loss_src_weight
        self.loss_dst_weight = loss_dst_weight

    def forward(
        self,
        source: Tensor,
        target: Tensor,
        src_weight: Union[Tensor, float] = 1.0,
        dst_weight: Union[Tensor, float] = 1.0,
        reduction_override: Optional[str] = None,
        return_indices: bool = False,
        **kwargs
    ) -> Union[Tuple[Tensor, Tensor, Tensor, Tensor], Tuple[Tensor, Tensor]]:
        """Forward function of loss calculation.

        Args:
            source (Tensor): Source set with shape [B, N, C] to
                calculate Chamfer Distance.
            target (Tensor): Destination set with shape [B, M, C] to
                calculate Chamfer Distance.
            src_weight (Tensor | float):
                Weight of source loss. Defaults to 1.0.
            dst_weight (Tensor | float):
                Weight of destination loss. Defaults to 1.0.
            reduction_override (str, optional): Method to reduce losses.
                The valid reduction method are 'none', 'sum' or 'mean'.
                Defaults to None.
            return_indices (bool): Whether to return indices.
                Defaults to False.

        Returns:
            tuple[Tensor]: If ``return_indices=True``, return losses of
                source and target with their corresponding indices in the
                order of ``(loss_source, loss_target, indices1, indices2)``.