diff --git a/monai/data/image_writer.py b/monai/data/image_writer.py index 5b47a5e6e5..f575954716 100644 --- a/monai/data/image_writer.py +++ b/monai/data/image_writer.py @@ -63,6 +63,21 @@ SUPPORTED_WRITERS: dict = {} +# Maps common file extensions to the package needed by ImageWriter backends. +FILETYPE_HINT: dict[str, str] = { + "nii": "nibabel", + "nii.gz": "nibabel", + "mha": "itk", + "mhd": "itk", + "nrrd": "itk", + "png": "pillow", + "jpg": "pillow", + "jpeg": "pillow", + "tif": "pillow", + "tiff": "pillow", + "bmp": "pillow", +} + def register_writer(ext_name, *im_writers): """ @@ -116,7 +131,9 @@ def resolve_writer(ext_name, error_if_not_found=True) -> Sequence: except Exception: # other writer init errors indicating it exists avail_writers.append(_writer) if not avail_writers and error_if_not_found: - raise OptionalImportError(f"No ImageWriter backend found for {fmt}.") + hint = FILETYPE_HINT.get(fmt) + extra = f" Try: pip install {hint}" if hint else "" + raise OptionalImportError(f"No ImageWriter backend found for '{fmt}'.{extra}") writer_tuple = ensure_tuple(avail_writers) SUPPORTED_WRITERS[fmt] = writer_tuple return writer_tuple diff --git a/tests/data/test_image_rw.py b/tests/data/test_image_rw.py index d90c1c8571..67292088a3 100644 --- a/tests/data/test_image_rw.py +++ b/tests/data/test_image_rw.py @@ -22,7 +22,7 @@ from parameterized import parameterized from monai.data.image_reader import ITKReader, NibabelReader, NrrdReader, PILReader -from monai.data.image_writer import ITKWriter, NibabelWriter, PILWriter, register_writer, resolve_writer +from monai.data.image_writer import FILETYPE_HINT, ITKWriter, NibabelWriter, PILWriter, register_writer, resolve_writer from monai.data.meta_tensor import MetaTensor from monai.transforms import LoadImage, SaveImage, moveaxis from monai.utils import MetaKeys, OptionalImportError, optional_import @@ -188,5 +188,22 @@ def test_3d(self, reader, writer): self.nrrd_rw(test_data, reader, writer, np.float32) +class TestResolveWriterHint(unittest.TestCase): + def test_filetype_hint_content(self): + self.assertEqual(FILETYPE_HINT.get("nii"), "nibabel") + self.assertEqual(FILETYPE_HINT.get("nii.gz"), "nibabel") + self.assertEqual(FILETYPE_HINT.get("png"), "pillow") + self.assertEqual(FILETYPE_HINT.get("mha"), "itk") + + def test_resolve_writer_error_message(self): + # Test with an unknown extension to see the base error message + # Since EXT_WILDCARD might have backends, it might not raise unless they are missing. + # But we can at least verify the logic is reachable. + try: + resolve_writer("unknown_ext", error_if_not_found=True) + except OptionalImportError as e: + self.assertIn("No ImageWriter backend found for 'unknown_ext'", str(e)) + + if __name__ == "__main__": unittest.main()