diff --git a/monai/data/dataset_summary.py b/monai/data/dataset_summary.py index 5b9e32afca..93d7d7043e 100644 --- a/monai/data/dataset_summary.py +++ b/monai/data/dataset_summary.py @@ -159,6 +159,8 @@ def calculate_statistics(self, foreground_threshold: int = 0): label, *_ = convert_data_type(data=label, output_type=torch.Tensor) image_foreground = image[torch.where(label > foreground_threshold)] + if image_foreground.numel() == 0: + continue voxel_max.append(image_foreground.max().item()) voxel_min.append(image_foreground.min().item()) @@ -166,6 +168,11 @@ def calculate_statistics(self, foreground_threshold: int = 0): voxel_sum += image_foreground.sum() voxel_square_sum += torch.square(image_foreground).sum() + if voxel_ct == 0: + raise ValueError( + f"no foreground voxels found in any sample with foreground_threshold={foreground_threshold}; " + "set foreground_threshold=-1 to compute statistics over whole images." + ) self.data_max, self.data_min = max(voxel_max), min(voxel_min) self.data_mean = (voxel_sum / voxel_ct).item() self.data_std = (torch.sqrt(voxel_square_sum / voxel_ct - self.data_mean**2)).item() @@ -204,11 +211,17 @@ def calculate_percentiles( label, *_ = convert_data_type(data=label, output_type=torch.Tensor) intensities = image[torch.where(label > foreground_threshold)].tolist() - if sampling_flag: - intensities = intensities[::interval] - all_intensities.append(intensities) + if intensities: + if sampling_flag: + intensities = intensities[::interval] + all_intensities.append(intensities) all_intensities = list(chain(*all_intensities)) + if not all_intensities: + raise ValueError( + f"no foreground voxels found in any sample with foreground_threshold={foreground_threshold}; " + "set foreground_threshold=-1 to compute statistics over whole images." + ) self.data_min_percentile, self.data_max_percentile = np.percentile( all_intensities, [min_percentile, max_percentile] ) diff --git a/tests/data/test_dataset_summary.py b/tests/data/test_dataset_summary.py index 21cc53de90..0735ae0a79 100644 --- a/tests/data/test_dataset_summary.py +++ b/tests/data/test_dataset_summary.py @@ -18,6 +18,7 @@ import nibabel as nib import numpy as np +import torch from monai.data import Dataset, DatasetSummary, create_test_image_3d from monai.transforms import LoadImaged @@ -99,6 +100,35 @@ def test_anisotropic_spacing(self): target_spacing = calculator.get_target_spacing(anisotropic_threshold=4.0, percentile=20.0) np.testing.assert_allclose(target_spacing, (1.0, 1.0, 1.8)) + def test_mixed_foreground_and_background(self): + data = [ + {"image": torch.rand(1, 4, 4), "label": torch.ones(1, 4, 4)}, + {"image": torch.rand(1, 4, 4), "label": torch.zeros(1, 4, 4)}, + ] + image = torch.cat([d["image"] for d in data]) + label = torch.cat([d["label"] for d in data]) + expected = image[torch.where(label > 0)] + + calculator = DatasetSummary(data, num_workers=0) + calculator.calculate_statistics() + np.testing.assert_allclose(calculator.data_mean, expected.mean().item(), rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(calculator.data_std, expected.std(correction=0).item(), rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(calculator.data_max, expected.max().item(), rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(calculator.data_min, expected.min().item(), rtol=1e-5, atol=1e-5) + + calculator.calculate_percentiles(sampling_flag=False) + np.testing.assert_allclose(calculator.data_min_percentile, np.percentile(expected, 0.5), rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(calculator.data_max_percentile, np.percentile(expected, 99.5), rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(calculator.data_median, np.median(expected), rtol=1e-5, atol=1e-5) + + def test_all_background(self): + data = [{"image": torch.rand(1, 4, 4), "label": torch.zeros(1, 4, 4)}] + calculator = DatasetSummary(data, num_workers=0) + with self.assertRaisesRegex(ValueError, "foreground_threshold"): + calculator.calculate_statistics() + with self.assertRaisesRegex(ValueError, "foreground_threshold"): + calculator.calculate_percentiles() + if __name__ == "__main__": unittest.main()