Class OccupancyMetric

embodiedscan/eval/metrics/occupancy_metric.py:18–178

Indoor scene evaluation metric. Args: iou_thr (list[float]): List of iou threshold when calculate the metric. Defaults to [0.25, 0.5]. collect_device (str, optional): Device name used for collecting results from different ranks during distributed trainin

@METRICS.register_module()
class OccupancyMetric(BaseMetric):
    """Indoor scene evaluation metric.

    Args:
        iou_thr (list[float]): List of iou threshold when calculate the
            metric. Defaults to  [0.25, 0.5].
        collect_device (str, optional): Device name used for collecting
            results from different ranks during distributed training.
            Must be 'cpu' or 'gpu'. Defaults to 'cpu'.
        prefix (str): The prefix that will be added in the metric
            names to disambiguate homonymous metrics of different evaluators.
            If prefix is not provided in the argument, self.default_prefix
            will be used instead. Default: None
    """

    def __init__(self,
                 collect_device: str = 'cpu',
                 prefix: Optional[str] = None,
                 batchwise_anns: bool = False,
                 **kwargs):
        super(OccupancyMetric, self).__init__(prefix=prefix,
                                              collect_device=collect_device)
        self.batchwise_anns = batchwise_anns

    def process(self, data_batch: dict, data_samples: Sequence[dict]) -> None:
        """Process one batch of data samples and predictions.

        The processed results should be stored in ``self.results``,
        which will be used to compute the metrics when all batches
        have been processed.

        Args:
            data_batch (dict): A batch of data from the dataloader.
            data_samples (Sequence[dict]): A batch of outputs from
                the model.
        """
        for data_sample in data_samples:
            pred_occ = data_sample['pred_occupancy']
            gt_4 = data_sample['gt_occupancy']
            gt_occ = torch.zeros_like(pred_occ)
            gt_occ[gt_4[:, 0], gt_4[:, 1], gt_4[:, 2]] = gt_4[:, 3]
            if 'gt_occupancy_masks' in data_sample:
                gt_occ_mask = data_sample['gt_occupancy_masks']
                gt_occ[~gt_occ_mask] = 255
            self.results.append((gt_occ, pred_occ))

    def compute_metrics(self, results: list) -> Dict[str, float]:
        """Compute the metrics from processed results.

        Args:
            results (list): The processed results of each batch.

        Returns:
            Dict[str, float]: The computed metrics. The keys are the names of
            the metrics, and the values are corresponding results.
        """
        logger: MMLogger = MMLogger.get_current_instance()
        num_class = len(self.dataset_meta['classes']) + 1