added better checks

wangkun543604 · Jun 13, 2023 · 5ff57b9 · 5ff57b9
1 parent d2b39b1
commit 5ff57b9
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/GANDLF/utils/tensor.py b/GANDLF/utils/tensor.py
@@ -526,10 +526,8 @@ def get_image_from_tensor(input_tensor: torch.Tensor) -> sitk.Image:
     """
     arr = input_tensor.cpu().numpy()
     return_image = sitk.GetImageFromArray(arr)
-    # this is specifically the case for 2D rgb images
-    if arr.shape[1] == 3:
-        return_image = sitk.GetImageFromArray(arr)
-    elif arr.shape[0] == 1:
+    # this is specifically the case for 3D images
+    if (arr.shape[0] == 1) and (arr.shape[1] > 3):
         return_image = sitk.GetImageFromArray(arr[0])
 
     return return_image