diff --git a/monai/auto3dseg/analyzer.py b/monai/auto3dseg/analyzer.py index a731546a9e..88e36f9d13 100644 --- a/monai/auto3dseg/analyzer.py +++ b/monai/auto3dseg/analyzer.py @@ -217,22 +217,6 @@ def __init__(self, image_key: str, stats_name: str = DataStatsKeys.IMAGE_STATS) self.update_ops(ImageStatsKeys.INTENSITY, SampleOperations()) def __call__(self, data): - # Input Validation Addition - if not isinstance(data, dict): - raise TypeError(f"Input data must be a dict, but got {type(data).__name__}.") - if self.image_key not in data: - raise KeyError(f"Key '{self.image_key}' not found in input data.") - image = data[self.image_key] - if not isinstance(image, (np.ndarray, torch.Tensor, MetaTensor)): - raise TypeError( - f"Value for '{self.image_key}' must be a numpy array, torch.Tensor, or MetaTensor, " - f"but got {type(image).__name__}." - ) - if image.ndim < 3: - raise ValueError( - f"Image data under '{self.image_key}' must have at least 3 dimensions, but got shape {image.shape}." - ) - # --- End of validation --- """ Callable to execute the pre-defined functions @@ -242,7 +226,9 @@ def __call__(self, data): has stats pre-defined by SampleOperations (max, min, ....). Raises: - RuntimeError if the stats report generated is not consistent with the pre- + ValueError: if ``nda_croppeds`` is present in the input dict but is not a + list/tuple or has a different length than the number of image channels. + RuntimeError: if the stats report generated is not consistent with the pre- defined report_format. Note: @@ -255,36 +241,46 @@ def __call__(self, data): restore_grad_state = torch.is_grad_enabled() torch.set_grad_enabled(False) - ndas = [d[self.image_key][i] for i in range(d[self.image_key].shape[0])] - if "nda_croppeds" not in d: - nda_croppeds = [get_foreground_image(nda) for nda in ndas] - - # perform calculation - report = deepcopy(self.get_report_format()) - - report[ImageStatsKeys.SHAPE] = [list(nda.shape) for nda in ndas] - report[ImageStatsKeys.CHANNELS] = len(ndas) - report[ImageStatsKeys.CROPPED_SHAPE] = [list(nda_c.shape) for nda_c in nda_croppeds] - report[ImageStatsKeys.SPACING] = ( - affine_to_spacing(data[self.image_key].affine).tolist() - if isinstance(data[self.image_key], MetaTensor) - else [1.0] * min(3, data[self.image_key].ndim) - ) + try: + ndas = [d[self.image_key][i] for i in range(d[self.image_key].shape[0])] + if "nda_croppeds" not in d: + nda_croppeds = [get_foreground_image(nda) for nda in ndas] + else: + nda_croppeds = d["nda_croppeds"] + if not isinstance(nda_croppeds, (list, tuple)) or len(nda_croppeds) != len(ndas): + raise ValueError( + f"Pre-computed 'nda_croppeds' must be a list with one entry per image channel " + f"(expected {len(ndas)}, got " + f"{len(nda_croppeds) if isinstance(nda_croppeds, (list, tuple)) else type(nda_croppeds).__name__})." + ) + + # perform calculation + report = deepcopy(self.get_report_format()) + + report[ImageStatsKeys.SHAPE] = [list(nda.shape) for nda in ndas] + report[ImageStatsKeys.CHANNELS] = len(ndas) + report[ImageStatsKeys.CROPPED_SHAPE] = [list(nda_c.shape) for nda_c in nda_croppeds] + report[ImageStatsKeys.SPACING] = ( + affine_to_spacing(data[self.image_key].affine).tolist() + if isinstance(data[self.image_key], MetaTensor) + else [1.0] * min(3, data[self.image_key].ndim) + ) - report[ImageStatsKeys.SIZEMM] = [ - a * b for a, b in zip(report[ImageStatsKeys.SHAPE][0], report[ImageStatsKeys.SPACING]) - ] + report[ImageStatsKeys.SIZEMM] = [ + a * b for a, b in zip(report[ImageStatsKeys.SHAPE][0], report[ImageStatsKeys.SPACING]) + ] - report[ImageStatsKeys.INTENSITY] = [ - self.ops[ImageStatsKeys.INTENSITY].evaluate(nda_c) for nda_c in nda_croppeds - ] + report[ImageStatsKeys.INTENSITY] = [ + self.ops[ImageStatsKeys.INTENSITY].evaluate(nda_c) for nda_c in nda_croppeds + ] - if not verify_report_format(report, self.get_report_format()): - raise RuntimeError(f"report generated by {self.__class__} differs from the report format.") + if not verify_report_format(report, self.get_report_format()): + raise RuntimeError(f"report generated by {self.__class__} differs from the report format.") - d[self.stats_name] = report + d[self.stats_name] = report + finally: + torch.set_grad_enabled(restore_grad_state) - torch.set_grad_enabled(restore_grad_state) logger.debug(f"Get image stats spent {time.time() - start}") return d diff --git a/tests/apps/test_auto3dseg.py b/tests/apps/test_auto3dseg.py index 6c840e944d..8021e07ba7 100644 --- a/tests/apps/test_auto3dseg.py +++ b/tests/apps/test_auto3dseg.py @@ -178,7 +178,10 @@ def __call__(self, data): class TestDataAnalyzer(unittest.TestCase): + """Integration tests for the auto3dseg analyzer pipeline.""" + def setUp(self): + """Create temporary directory and write simulated datalist JSON file.""" self.test_dir = tempfile.TemporaryDirectory() work_dir = self.test_dir.name self.dataroot_dir = os.path.join(work_dir, "sim_dataroot") @@ -188,6 +191,7 @@ def setUp(self): @parameterized.expand(SIM_CPU_TEST_CASES) def test_data_analyzer_cpu(self, input_params): + """Verify DataAnalyzer produces per-case stats on CPU across dim/label combinations.""" sim_dim = input_params["sim_dim"] label_key = input_params["label_key"] image_only = not bool(label_key) @@ -204,6 +208,7 @@ def test_data_analyzer_cpu(self, input_params): assert len(datastat["stats_by_cases"]) == len(sim_datalist["training"]) def test_data_analyzer_histogram(self): + """Verify DataAnalyzer runs in histogram_only mode with no label key.""" create_sim_data( self.dataroot_dir, sim_datalist, [32] * 3, image_only=True, rad_max=8, rad_min=1, num_seg_classes=1 ) @@ -221,6 +226,7 @@ def test_data_analyzer_histogram(self): @parameterized.expand(SIM_GPU_TEST_CASES) @skip_if_no_cuda def test_data_analyzer_gpu(self, input_params): + """Verify DataAnalyzer produces per-case stats on GPU (skipped if CUDA unavailable).""" sim_dim = input_params["sim_dim"] label_key = input_params["label_key"] image_only = not bool(label_key) @@ -236,6 +242,7 @@ def test_data_analyzer_gpu(self, input_params): assert len(datastat["stats_by_cases"]) == len(sim_datalist["training"]) def test_basic_operation_class(self): + """Verify Operations.evaluate returns correct stat keys and shapes with and without axis.""" op = TestOperations() test_data = np.random.rand(10, 10).astype(np.float64) test_ret_1 = op.evaluate(test_data) @@ -250,6 +257,7 @@ def test_basic_operation_class(self): assert test_ret_2["max"].ndim == 1 def test_sample_operations(self): + """Verify SampleOperations works with both numpy arrays and MetaTensors.""" op = SampleOperations() test_data_np = np.random.rand(10, 10).astype(np.float64) test_data_mt = MetaTensor(test_data_np, device=device) @@ -265,6 +273,7 @@ def test_sample_operations(self): assert "sum" in test_ret_np def test_summary_operations(self): + """Verify SummaryOperations reduces a stat dict to scalar summary values.""" op = SummaryOperations() test_dict = {"min": [0, 1, 2, 3], "max": [2, 3, 4, 5], "mean": [1, 2, 3, 4], "sum": [2, 4, 6, 8]} test_ret = op.evaluate(test_dict) @@ -277,6 +286,7 @@ def test_summary_operations(self): assert isinstance(test_ret["sum"], Number) def test_basic_analyzer_class(self): + """Verify a custom Analyzer subclass computes and stores stats in the output dict.""" test_data = {} test_data["image_test"] = np.random.rand(10, 10) report_format = {"stats": None} @@ -288,6 +298,7 @@ def test_basic_analyzer_class(self): assert result["test"]["stats"]["mean"] == np.mean(test_data["image_test"]) def test_transform_analyzer_class(self): + """Verify a custom Analyzer integrates correctly as a step in a Compose transform.""" transform = Compose([LoadImaged(keys=["image"]), TestImageAnalyzer(image_key="image")]) create_sim_data(self.dataroot_dir, sim_datalist, (32, 32, 32), rad_max=8, rad_min=1, num_seg_classes=1) files, _ = datafold_read(sim_datalist, self.dataroot_dir, fold=-1) @@ -302,6 +313,7 @@ def test_transform_analyzer_class(self): assert "mean" in d["test_image"]["test_stats"] def test_image_stats_case_analyzer(self): + """Verify ImageStats produces a report matching the expected format for 3-D images.""" analyzer = ImageStats(image_key="image") transform = Compose( [ @@ -323,6 +335,7 @@ def test_image_stats_case_analyzer(self): assert verify_report_format(d["image_stats"], report_format) def test_foreground_image_stats_cases_analyzer(self): + """Verify FgImageStats produces a valid foreground stats report.""" analyzer = FgImageStats(image_key="image", label_key="label") transform_list = [ LoadImaged(keys=["image", "label"]), @@ -345,6 +358,7 @@ def test_foreground_image_stats_cases_analyzer(self): assert verify_report_format(d["image_foreground_stats"], report_format) def test_label_stats_case_analyzer(self): + """Verify LabelStats produces a valid report including per-label statistics.""" analyzer = LabelStats(image_key="image", label_key="label") transform = Compose( [ @@ -369,6 +383,7 @@ def test_label_stats_case_analyzer(self): @parameterized.expand(LABEL_STATS_DEVICE_TEST_CASES) def test_label_stats_mixed_device_analyzer(self, input_params): + """Verify LabelStats handles tensors split across CPU and CUDA devices.""" image_device = torch.device(input_params["image_device"]) label_device = torch.device(input_params["label_device"]) @@ -413,6 +428,7 @@ def test_label_stats_mixed_device_analyzer(self, input_params): self.assertAlmostEqual(foreground_stats[1]["mean"], 14.75) def test_filename_case_analyzer(self): + """Verify FilenameStats records both image and label paths in the output dict.""" analyzer_image = FilenameStats("image", DataStatsKeys.BY_CASE_IMAGE_PATH) analyzer_label = FilenameStats("label", DataStatsKeys.BY_CASE_IMAGE_PATH) transform_list = [LoadImaged(keys=["image", "label"]), analyzer_image, analyzer_label] @@ -426,6 +442,7 @@ def test_filename_case_analyzer(self): assert DataStatsKeys.BY_CASE_IMAGE_PATH in d def test_filename_case_analyzer_image_only(self): + """Verify FilenameStats handles image-only input and stores 'None' for the label path.""" analyzer_image = FilenameStats("image", DataStatsKeys.BY_CASE_IMAGE_PATH) analyzer_label = FilenameStats(None, DataStatsKeys.BY_CASE_IMAGE_PATH) transform_list = [LoadImaged(keys=["image"]), analyzer_image, analyzer_label] @@ -440,6 +457,7 @@ def test_filename_case_analyzer_image_only(self): assert d[DataStatsKeys.BY_CASE_IMAGE_PATH] == "None" def test_image_stats_summary_analyzer(self): + """Verify ImageStatsSumm correctly aggregates per-case image stats.""" summary_analyzer = ImageStatsSumm("image_stats") transform_list = [ @@ -463,6 +481,7 @@ def test_image_stats_summary_analyzer(self): assert verify_report_format(summary_report, report_format) def test_fg_image_stats_summary_analyzer(self): + """Verify FgImageStatsSumm correctly aggregates per-case foreground stats.""" summary_analyzer = FgImageStatsSumm("image_foreground_stats") transform_list = [ @@ -488,6 +507,7 @@ def test_fg_image_stats_summary_analyzer(self): assert verify_report_format(summary_report, report_format) def test_label_stats_summary_analyzer(self): + """Verify LabelStatsSumm correctly aggregates per-case label stats.""" summary_analyzer = LabelStatsSumm("label_stats") transform_list = [ @@ -513,6 +533,7 @@ def test_label_stats_summary_analyzer(self): assert verify_report_format(summary_report, report_format) def test_seg_summarizer(self): + """Verify SegSummarizer produces a summary with image, foreground, and label stat keys.""" summarizer = SegSummarizer("image", "label") keys = ["image", "label"] transform_list = [ @@ -539,7 +560,45 @@ def test_seg_summarizer(self): assert str(DataStatsKeys.FG_IMAGE_STATS) in report assert str(DataStatsKeys.LABEL_STATS) in report + def test_image_stats_precomputed_nda_croppeds(self): + """Verify ImageStats handles pre-populated nda_croppeds without crashing. + + Previously raised UnboundLocalError because nda_croppeds was only assigned + inside the ``if "nda_croppeds" not in d`` branch but used unconditionally. + """ + analyzer = ImageStats(image_key="image") + image = torch.rand(1, 10, 10, 10) + precomputed = [np.random.rand(8, 8, 8)] # simulated pre-cropped foreground + data = {"image": MetaTensor(image), "nda_croppeds": precomputed} + result = analyzer(data) + assert "image_stats" in result + assert verify_report_format(result["image_stats"], analyzer.get_report_format()) + + def test_analyzer_grad_state_restored_after_call(self): + """Verify ImageStats restores torch grad-enabled state on both normal and disabled entry. + + Checks that the try/finally guard correctly restores the state regardless of + whether grad was enabled or disabled before the call. + """ + analyzer = ImageStats(image_key="image") + image = torch.rand(1, 10, 10, 10) + data = {"image": MetaTensor(image)} + + # grad enabled before call → must still be enabled after + torch.set_grad_enabled(True) + analyzer(data) + assert torch.is_grad_enabled(), "grad state was not restored after ImageStats call" + + # grad disabled before call → must still be disabled after + torch.set_grad_enabled(False) + try: + analyzer(data) + assert not torch.is_grad_enabled(), "grad state was not restored after ImageStats call" + finally: + torch.set_grad_enabled(True) # always restore for subsequent tests + def tearDown(self) -> None: + """Remove the temporary test directory.""" self.test_dir.cleanup()