diff --git a/monai/metrics/panoptic_quality.py b/monai/metrics/panoptic_quality.py index 03cef9d566..0309bda478 100644 --- a/monai/metrics/panoptic_quality.py +++ b/monai/metrics/panoptic_quality.py @@ -250,32 +250,26 @@ def _get_pairwise_iou( pred_id_list = _get_id_list(pred) true_id_list = _get_id_list(gt) - pairwise_iou = torch.zeros([len(true_id_list) - 1, len(pred_id_list) - 1], dtype=torch.float, device=device) - true_masks: list[torch.Tensor] = [] - pred_masks: list[torch.Tensor] = [] - - for t in true_id_list[1:]: - t_mask = torch.as_tensor(gt == t, device=device).int() - true_masks.append(t_mask) - - for p in pred_id_list[1:]: - p_mask = torch.as_tensor(pred == p, device=device).int() - pred_masks.append(p_mask) - - for true_id in range(1, len(true_id_list)): - t_mask = true_masks[true_id - 1] - pred_true_overlap = pred[t_mask > 0] - pred_true_overlap_id = list(pred_true_overlap.unique()) - for pred_id in pred_true_overlap_id: - if pred_id == 0: - continue - p_mask = pred_masks[pred_id - 1] - total = (t_mask + p_mask).sum() - inter = (t_mask * p_mask).sum() - iou = inter / (total - inter) - pairwise_iou[true_id - 1, pred_id - 1] = iou - - return pairwise_iou, true_id_list, pred_id_list + num_true = len(true_id_list) - 1 + num_pred = len(pred_id_list) - 1 + pairwise_iou = torch.zeros([num_true, num_pred], dtype=torch.float, device=device) + if num_true == 0 or num_pred == 0: + return pairwise_iou, true_id_list, pred_id_list + + # ids are contiguous after `remap_instance_id`, so count all pairwise intersections in one bincount + gt_flat = gt.reshape(-1).long() + pred_flat = pred.reshape(-1).long() + stride = num_pred + 1 + joint = gt_flat * stride + pred_flat + intersection = torch.bincount(joint, minlength=(num_true + 1) * stride).reshape(num_true + 1, stride).float() + true_area = torch.bincount(gt_flat, minlength=num_true + 1).float() + pred_area = torch.bincount(pred_flat, minlength=num_pred + 1).float() + + inter = intersection[1:, 1:] # drop background row/column + union = true_area[1:, None] + pred_area[None, 1:] - inter + pairwise_iou = torch.where(inter > 0, inter / union, pairwise_iou) + + return pairwise_iou.to(device), true_id_list, pred_id_list def _get_paired_iou( diff --git a/tests/metrics/test_compute_panoptic_quality.py b/tests/metrics/test_compute_panoptic_quality.py index 3f8c21debb..2f3dd6b8cb 100644 --- a/tests/metrics/test_compute_panoptic_quality.py +++ b/tests/metrics/test_compute_panoptic_quality.py @@ -18,7 +18,7 @@ from parameterized import parameterized from monai.metrics import PanopticQualityMetric, compute_panoptic_quality -from monai.metrics.panoptic_quality import compute_mean_iou +from monai.metrics.panoptic_quality import _get_pairwise_iou, compute_mean_iou from tests.test_utils import SkipIfNoModule _device = "cuda:0" if torch.cuda.is_available() else "cpu" @@ -215,6 +215,24 @@ def test_invalid_3d_shape(self): with self.assertRaises(ValueError): metric(invalid_pred, invalid_gt) + def test_pairwise_iou(self): + """`_get_pairwise_iou` returns the expected IoU matrix, including no-overlap and empty cases.""" + gt = torch.as_tensor([[1, 1, 0], [0, 2, 2], [0, 0, 0]], device=_device) + pred = torch.as_tensor([[1, 0, 0], [0, 2, 2], [0, 0, 0]], device=_device) + pairwise, _, _ = _get_pairwise_iou(pred, gt, device=_device) + np.testing.assert_allclose(pairwise.cpu().numpy(), [[0.5, 0.0], [0.0, 1.0]], atol=1e-6) + + # true instance with no overlapping prediction -> all-zero row, shape preserved + disjoint_pred = torch.as_tensor([[0, 0, 1], [0, 0, 0], [0, 0, 0]], device=_device) + pairwise, _, _ = _get_pairwise_iou(disjoint_pred, gt, device=_device) + self.assertEqual(pairwise.shape, torch.Size([2, 1])) + self.assertTrue(torch.all(pairwise == 0)) + + # empty gt / empty pred -> degenerate matrices, no error + empty = torch.zeros((3, 3), dtype=torch.int, device=_device) + self.assertEqual(_get_pairwise_iou(pred, empty, device=_device)[0].shape, torch.Size([0, 2])) + self.assertEqual(_get_pairwise_iou(empty, gt, device=_device)[0].shape, torch.Size([2, 0])) + def test_compute_mean_iou_invalid_shape(self): """Test that compute_mean_iou raises ValueError for invalid shapes.""" from monai.metrics.panoptic_quality import compute_mean_iou