EnsureChannelFirst: avoid re-creation of AddChannel (#3649)

* avoid re-creation of AddChannel in EnsureChannelFirst
Project-MONAI · Jan 12, 2022 · 35c2b37 · 35c2b37
1 parent d58e234
commit 35c2b37
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/monai/inferers/inferer.py b/monai/inferers/inferer.py
@@ -30,7 +30,8 @@ class Inferer(ABC):
     Example code::
 
         device = torch.device("cuda:0")
-        data = ToTensor()(LoadImage()(filename=img_path)).to(device)
+        transform = Compose([ToTensor(), LoadImage(image_only=True)])
+        data = transform(img_path).to(device)
         model = UNet(...).to(device)
         inferer = SlidingWindowInferer(...)
 

diff --git a/monai/transforms/utility/array.py b/monai/transforms/utility/array.py
@@ -202,6 +202,7 @@ def __init__(self, strict_check: bool = True):
             strict_check: whether to raise an error when the meta information is insufficient.
         """
         self.strict_check = strict_check
+        self.add_channel = AddChannel()
 
     def __call__(self, img: NdarrayOrTensor, meta_dict: Optional[Mapping] = None) -> NdarrayOrTensor:
         """
@@ -223,7 +224,7 @@ def __call__(self, img: NdarrayOrTensor, meta_dict: Optional[Mapping] = None) ->
             warnings.warn(msg)
             return img
         if channel_dim == "no_channel":
-            return AddChannel()(img)
+            return self.add_channel(img)
         return AsChannelFirst(channel_dim=channel_dim)(img)