Skip to content

Commit 08523ce

Browse files
committed
Print what package should be installed when suitable writer is missing
Fixes #7980 Signed-off-by: ytl0623 <[email protected]>
1 parent 6c23fd0 commit 08523ce

File tree

1 file changed

+25
-1
lines changed

1 file changed

+25
-1
lines changed

Diff for: monai/utils/module.py

+25-1
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
from pydoc import locate
2727
from re import match
2828
from types import FunctionType, ModuleType
29-
from typing import Any, Iterable, cast
29+
from typing import Any, Iterable, cast, Dict, List
3030

3131
import torch
3232

@@ -60,6 +60,16 @@
6060
"pytorch_after",
6161
]
6262

63+
WRITER_PACKAGE_MAP: Dict[str, List[str]] = {
64+
"png": ["pillow"],
65+
"jpg": ["pillow"],
66+
"jpeg": ["pillow"],
67+
"nii": ["nibabel"],
68+
"nii.gz": ["nibabel"],
69+
"dcm": ["pydicom"],
70+
# Add more mappings as needed
71+
}
72+
6373

6474
def look_up_option(
6575
opt_str: Hashable,
@@ -334,6 +344,20 @@ class OptionalImportError(ImportError):
334344
"""
335345
Could not import APIs from an optional dependency.
336346
"""
347+
348+
def __init__(self, msg: str = "") -> None:
349+
super().__init__(msg)
350+
self.msg = msg
351+
352+
def __str__(self) -> str:
353+
original_msg = super().__str__()
354+
if "ImageWriter" in original_msg and "backend found for" in original_msg:
355+
ext = original_msg.split("for ")[-1].strip(".")
356+
suggested_packages = WRITER_PACKAGE_MAP.get(ext, [])
357+
if suggested_packages:
358+
package_list = ", ".join(suggested_packages)
359+
return f"{original_msg} Please install one of the following packages: {package_list}"
360+
return original_msg
337361

338362

339363
def optional_import(

0 commit comments

Comments
 (0)