Class BBoxCDLoss

embodiedscan/models/losses/chamfer_distance.py:207–285

Calculate Chamfer Distance of the corners for two bbox sets. Args: mode (str): Criterion mode to calculate distance. The valid modes are 'smooth_l1', 'l1' or 'l2'. Defaults to 'l2'. group (str): How corners are grouped. The valid groups are 'g8' or 'g4',

@MODELS.register_module()
class BBoxCDLoss(nn.Module):
    """Calculate Chamfer Distance of the corners for two bbox sets.

    Args:
        mode (str): Criterion mode to calculate distance.
            The valid modes are 'smooth_l1', 'l1' or 'l2'. Defaults to 'l2'.
        group (str): How corners are grouped.
            The valid groups are 'g8' or 'g4',
            meaning either all 8 corners are in a single group
            or the corners are divided into two groups, each containing 4.
        reduction (str): Method to reduce losses.
            The valid reduction method are 'none', 'sum' or 'mean'.
            Defaults to 'mean'.
        loss_weight (float): Weight of loss. Defaults to l.0.
    """

    def __init__(
        self,
        mode: str = 'l2',
        group: str = 'g8',
        reduction: str = 'mean',
        loss_weight: float = 1.0,
    ) -> None:
        super(BBoxCDLoss, self).__init__()

        assert mode in ['smooth_l1', 'l1', 'l2']
        assert group in ['g4', 'g8']
        assert reduction in ['none', 'sum', 'mean']
        self.mode = mode
        self.group = group
        self.reduction = reduction
        self.loss_weight = loss_weight

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

        Args:
            source (Tensor): Source bbox set with shape [N, bbox_dim] to
                calculate Chamfer Distance.
            target (Tensor): Destination bbox set with shape [M, bbox_dim] to
                calculate Chamfer Distance.
            loss_weight (Tensor | float):
                Weight of 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.

        Returns:
            Tensor: return ``loss_source``.
        """
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = (reduction_override
                     if reduction_override else self.reduction)