From 6bbc5348ab7678f9f77ae16b4ecd20b82bcf14e5 Mon Sep 17 00:00:00 2001 From: songhahaha66 Date: Wed, 5 Nov 2025 22:25:36 +0800 Subject: [PATCH] Add dependency hints to image writer errors Signed-off-by: songhahaha66 --- monai/data/image_writer.py | 29 ++++++++++++++++++++++++++++- tests/data/test_image_rw.py | 28 +++++++++++++++++++++++++++- 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/monai/data/image_writer.py b/monai/data/image_writer.py index b9e8b9e68e..38235abb9e 100644 --- a/monai/data/image_writer.py +++ b/monai/data/image_writer.py @@ -62,6 +62,7 @@ ] SUPPORTED_WRITERS: dict = {} +WRITER_DEPENDENCY_HINTS: dict[type, tuple[str, str]] = {} def register_writer(ext_name, *im_writers): @@ -106,17 +107,34 @@ def resolve_writer(ext_name, error_if_not_found=True) -> Sequence: if fmt.startswith("."): fmt = fmt[1:] avail_writers = [] + dependency_hints: set[tuple[str, str]] = set() default_writers = SUPPORTED_WRITERS.get(EXT_WILDCARD, ()) for _writer in look_up_option(fmt, SUPPORTED_WRITERS, default=default_writers): try: _writer() # this triggers `monai.utils.module.require_pkg` to check the system availability avail_writers.append(_writer) except OptionalImportError: + hint = WRITER_DEPENDENCY_HINTS.get(_writer) + if hint: + dependency_hints.add(hint) continue 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_msg = "" + if dependency_hints: + sorted_hints = sorted(dependency_hints, key=lambda item: item[0].lower()) + if len(sorted_hints) == 1: + pkg, cmd = sorted_hints[0] + hint_msg = f" Install `{pkg}` (e.g. `{cmd}`) to enable writing {fmt} images." + else: + pkg_names = ", ".join(f"`{pkg}`" for pkg, _ in sorted_hints) + commands = ", ".join(f"`{cmd}`" for _, cmd in sorted_hints) + hint_msg = ( + f" Install one of the supported dependencies {pkg_names} " + f"(for example: {commands}) to enable writing {fmt} images." + ) + raise OptionalImportError(f"No ImageWriter backend found for {fmt}.{hint_msg}") writer_tuple = ensure_tuple(avail_writers) SUPPORTED_WRITERS[fmt] = writer_tuple return writer_tuple @@ -862,6 +880,15 @@ def create_backend_obj( return PILImage.fromarray(data, mode=kwargs.pop("image_mode", None)) +WRITER_DEPENDENCY_HINTS.update( + { + ITKWriter: ("ITK", "pip install itk"), + NibabelWriter: ("Nibabel", "pip install nibabel"), + PILWriter: ("Pillow", "pip install pillow"), + } +) + + def init(): """ Initialize the image writer modules according to the filename extension. diff --git a/tests/data/test_image_rw.py b/tests/data/test_image_rw.py index d90c1c8571..fdb30270b4 100644 --- a/tests/data/test_image_rw.py +++ b/tests/data/test_image_rw.py @@ -21,11 +21,12 @@ import torch from parameterized import parameterized +import monai.data.image_writer as image_writer 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.meta_tensor import MetaTensor from monai.transforms import LoadImage, SaveImage, moveaxis -from monai.utils import MetaKeys, OptionalImportError, optional_import +from monai.utils import MetaKeys, OptionalImportError, optional_import, require_pkg from tests.test_utils import TEST_NDARRAYS, assert_allclose _, has_itk = optional_import("itk", allow_namespace_pkg=True) @@ -150,6 +151,31 @@ def test_1_new(self): register_writer("new2", lambda x: x + 1) self.assertEqual(resolve_writer("new")[0](0), 1) + def test_missing_dependency_hint(self): + ext = ".needshint" + fmt_key = ext.lstrip(".").lower() + previous = image_writer.SUPPORTED_WRITERS.get(fmt_key) + + @require_pkg(pkg_name="__monai_missing_test_pkg__") + class MissingHintWriter(image_writer.ImageWriter): + pass + + image_writer.WRITER_DEPENDENCY_HINTS[MissingHintWriter] = ("FakePkg", "pip install fakepkg") + + try: + register_writer(ext, MissingHintWriter) + with self.assertRaises(OptionalImportError) as ctx: + resolve_writer(ext) + err_msg = str(ctx.exception) + self.assertIn("FakePkg", err_msg) + self.assertIn("pip install fakepkg", err_msg) + finally: + image_writer.WRITER_DEPENDENCY_HINTS.pop(MissingHintWriter, None) + if previous is None: + image_writer.SUPPORTED_WRITERS.pop(fmt_key, None) + else: + image_writer.SUPPORTED_WRITERS[fmt_key] = previous + @unittest.skipUnless(has_itk, "itk not installed") class TestLoadSaveNrrd(unittest.TestCase):