Skip to content

Commit 45aa9a6

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

File tree

1 file changed

+24
-1
lines changed

1 file changed

+24
-1
lines changed

monai/utils/module.py

+24-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,15 @@
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+
}
6372

6473
def look_up_option(
6574
opt_str: Hashable,
@@ -335,6 +344,20 @@ class OptionalImportError(ImportError):
335344
Could not import APIs from an optional dependency.
336345
"""
337346

347+
def __init__(self, msg: str = "") -> None:
348+
super().__init__(msg)
349+
self.msg = msg
350+
351+
def __str__(self) -> str:
352+
original_msg = super().__str__()
353+
if "ImageWriter" in original_msg and "backend found for" in original_msg:
354+
ext = original_msg.split("for ")[-1].strip(".")
355+
suggested_packages = WRITER_PACKAGE_MAP.get(ext, [])
356+
if suggested_packages:
357+
package_list = ", ".join(suggested_packages)
358+
return f"{original_msg} Please install one of the following packages: {package_list}"
359+
return original_msg
360+
338361

339362
def optional_import(
340363
module: str,

0 commit comments

Comments
 (0)