Fix typing errors

TorchIO-project · Sep 23, 2024 · a7ea6a8 · a7ea6a8
1 parent 0475b39
commit a7ea6a8
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/src/torchio/transforms/augmentation/intensity/random_ghosting.py b/src/torchio/transforms/augmentation/intensity/random_ghosting.py
@@ -182,6 +182,10 @@ def __init__(
         self.args_names = ['num_ghosts', 'axis', 'intensity', 'restore']
 
     def apply_transform(self, subject: Subject) -> Subject:
+        axis: Union[int, Dict[str, int]]
+        num_ghosts: Union[int, Dict[str, int]]
+        intensity: Union[float, Dict[str, float]]
+        restore: Union[Optional[float], Dict[str, Optional[float]]]
         for name, image in self.get_images_dict(subject).items():
             if self.arguments_are_dict():
                 assert isinstance(self.axis, dict)
@@ -253,8 +257,8 @@ def _get_planes_to_modify(
     ) -> torch.Tensor:
         slices = [slice(None)] * spectrum.ndim
         slices[axis] = slice(None, None, num_ghosts)
-        slices = tuple(slices)
-        return spectrum[slices]
+        slices_tuple = tuple(slices)
+        return spectrum[slices_tuple]
 
     @staticmethod
     def _get_slices_to_restore(
@@ -271,6 +275,6 @@ def _get_slices_to_restore(
             size_restore = int(np.round(restore_center * dim_shape))
             slice_ = slice(mid_idx - size_restore // 2, mid_idx + size_restore // 2)
         slices[axis] = slice_
-        slices = tuple(slices)
-        restore_tensor = spectrum[slices]
-        return restore_tensor, slices
+        slices_tuple = tuple(slices)
+        restore_tensor = spectrum[slices_tuple]
+        return restore_tensor, slices_tuple