ensure 2d rgb cases are picked up correctly and all aux cases are han…

…dled
wangkun543604 · Jun 13, 2023 · 832a998 · 832a998
1 parent 590b137
commit 832a998
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/GANDLF/utils/tensor.py b/GANDLF/utils/tensor.py
@@ -524,4 +524,12 @@ def get_image_from_tensor(input_tensor: torch.Tensor) -> sitk.Image:
     Returns:
         sitk.Image: The converted sitk image.
     """
-    return sitk.GetImageFromArray(input_tensor.cpu().numpy())
+    arr = input_tensor.cpu().numpy()
+    # this is specifically the case for 2D rgb images
+    return_image = sitk.GetImageFromArray(arr)
+    if arr.shape[1] == 3:
+        return_image = sitk.GetImageFromArray(arr)
+    elif arr.shape[0] == 1:
+        return_image = sitk.GetImageFromArray(arr[0])
+
+    return return_image