MONAILabel needed changes for DeepEdit/Deepgrow/NuClick Transform to support 0.9.1*

monai/apps/deepgrow/transforms.py

@@ -52,7 +52,7 @@ def _apply(self, label):
 
     def __call__(self, data):
         d: Dict = dict(data)
-        label = d[self.label]
+        label = d[self.label].numpy() if isinstance(data[self.label], torch.Tensor) else data[self.label]
         if label.shape[0] != 1:
             raise ValueError(f"Only supports single channel labels, got label shape {label.shape}!")
 
@@ -208,6 +208,10 @@ def _get_signal(self, image, guidance):
 
     def _apply(self, image, guidance):
         signal = self._get_signal(image, guidance)
+
+        if isinstance(image, torch.Tensor):
+            image = image.detach().cpu().numpy()
+
         image = image[0 : 0 + self.number_intensity_ch, ...]
         return np.concatenate([image, signal], axis=0)
 

monai/apps/nuclick/transforms.py

@@ -14,6 +14,7 @@
 from typing import Any, Tuple, Union
 
 import numpy as np
+import torch
 
 from monai.config import KeysCollection
 from monai.transforms import MapTransform, Randomizable, SpatialPad
@@ -60,7 +61,8 @@ def __init__(self, keys: KeysCollection, connectivity: int = 1, allow_missing_ke
     def __call__(self, data):
         d = dict(data)
         for key in self.keys:
-            d[key] = measure.label(d[key], connectivity=self.connectivity).astype(np.uint8)
+            img = d[key].numpy() if isinstance(d[key], torch.Tensor) else d[key]
+            d[key] = measure.label(img, connectivity=self.connectivity).astype(np.uint8)
         return d
 
 
@@ -159,8 +161,9 @@ def __call__(self, data):
         for key in self.keys:
             self.label = key
 
-        label = d[self.label]
-        mask_value = d[self.mask_value]
+        label = d[self.label].numpy() if isinstance(d[self.label], torch.Tensor) else d[self.label]
+        mask_value = d[self.mask_value].numpy() if isinstance(d[self.mask_value], torch.Tensor) else d[self.mask_value]
+
         mask = np.uint8(label == mask_value)
         others = (1 - mask) * label
         others = self._mask_relabeling(others[0], min_area=self.min_area)[np.newaxis]
@@ -200,7 +203,7 @@ def __init__(self, keys: KeysCollection, min_size: int = 500, allow_missing_keys
     def __call__(self, data):
         d = dict(data)
         for key in self.keys:
-            img = d[key]
+            img = d[key].numpy() if isinstance(d[key], torch.Tensor) else d[key]
             d[key] = self.filter(img)
         return d
 
@@ -284,9 +287,9 @@ def __init__(
     def __call__(self, data):
         d = dict(data)
 
-        image = d[self.image]
-        mask = d[self.label]
-        others = d[self.others]
+        image = d[self.image].numpy() if isinstance(d[self.image], torch.Tensor) else d[self.image]
+        mask = d[self.label].numpy() if isinstance(d[self.label], torch.Tensor) else d[self.label]
+        others = d[self.others].numpy() if isinstance(d[self.others], torch.Tensor) else d[self.others]
 
         inc_sig = self.inclusion_map(mask[0])
         exc_sig = self.exclusion_map(others[0], drop_rate=self.drop_rate, jitter_range=self.jitter_range)
@@ -353,7 +356,8 @@ def __call__(self, data):
         cx = [xy[0] for xy in pos]
         cy = [xy[1] for xy in pos]
 
-        img = d[self.image].astype(np.uint8)
+        img = d[self.image].numpy() if isinstance(d[self.image], torch.Tensor) else d[self.image]
+        img = img.astype(np.uint8)
         img_width = img.shape[-1]
         img_height = img.shape[-2]