[utils] Implicitly convert unsigned types which are not supported by …

…PyTorch
BioMedIA · Aug 3, 2023 · 03d763f · 03d763f
1 parent 28ccfa9
commit 03d763f
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 1 deletion.
diff --git a/src/deepali/utils/imageio/meta.py b/src/deepali/utils/imageio/meta.py
@@ -64,6 +64,10 @@ def read_meta_image(arg: Union[PathUri, bytes, io.BufferedReader]) -> Tuple[Tens
         spacing=spacing,
         direction=matrix,
     )
+    if data.dtype == np.uint16:
+        data = data.astype(np.int32)
+    elif data.dtype == np.uint32:
+        data = data.astype(np.int64)
     return torch.from_numpy(data), grid
 
 

diff --git a/src/deepali/utils/imageio/nifti.py b/src/deepali/utils/imageio/nifti.py
@@ -88,6 +88,10 @@ def read_nifti_image(path: PathUri) -> Tuple[Tensor, Grid]:
     grid = Grid(size=size, origin=origin, spacing=spacing, direction=direction)
     if data.ndim == grid.ndim:
         data = np.expand_dims(data, 0)
+    if data.dtype == np.uint16:
+        data = data.astype(np.int32)
+    elif data.dtype == np.uint32:
+        data = data.astype(np.int64)
     return torch.from_numpy(data), grid
 
 

diff --git a/src/deepali/utils/simpleitk/torch.py b/src/deepali/utils/simpleitk/torch.py
@@ -48,7 +48,7 @@ def tensor_from_image(
     r"""Create image data tensor from ``SimpleITK.Image``."""
     if image.GetPixelID() == sitk.sitkUInt16:
         image = sitk.Cast(image, sitk.sitkInt32)
-    elif image.GetPixelID() == sitk.sitkUInt16:
+    elif image.GetPixelID() == sitk.sitkUInt32:
         image = sitk.Cast(image, sitk.sitkInt64)
     data = torch.from_numpy(sitk.GetArrayFromImage(image))
     data = data.unsqueeze(0)