Standardize Axes in Random Transforms. Add Random Axis to RandomMotion

src/torchio/data/image.py

@@ -52,6 +52,16 @@
 PROTECTED_KEYS = DATA, AFFINE, TYPE, PATH, STEM
 TypeBound = Tuple[float, float]
 TypeBounds = Tuple[TypeBound, TypeBound, TypeBound]
+FLIP_AXIS = {
+ 'L': 'R',
+ 'R': 'L',
+ 'A': 'P',
+ 'P': 'A',
+ 'I': 'S',
+ 'S': 'I',
+ 'T': 'B',
+ 'B': 'T',
+}
 
 deprecation_message = (
  'Setting the image data with the property setter is deprecated. Use the'
@@ -378,7 +388,7 @@ def axis_name_to_index(self, axis: str) -> int:
  versions and first letters are also valid, as only the first
  letter will be used.
 
- .. note:: If you are working with animals, you should probably use
+ .. note:: If you are working with animals, you should use
  ``'Superior'``, ``'Inferior'``, ``'Anterior'`` and ``'Posterior'``
  for ``'Dorsal'``, ``'Ventral'``, ``'Rostral'`` and ``'Caudal'``,
  respectively.
@@ -392,6 +402,15 @@ def axis_name_to_index(self, axis: str) -> int:
  if not isinstance(axis, str):
  raise ValueError('Axis must be a string')
  axis = axis[0].upper()
+ if axis not in 'LRPAISTB':
+ message = (
+ 'Incorrect axis naming. Please use one of: "Left", "Right", '
+ '"Anterior", "Posterior", "Inferior", "Superior". '
+ 'Lower-case versions and first letters are also valid '
+ '(i.e., "L", "r", etc). For 2D images, use "Top" and "Bottom" '
+ 'to refer to the vertical (2nd) axis.'
+ )
+ raise ValueError(message)
 
  # Generally, TorchIO tensors are (C, W, H, D)
  if axis in 'TB': # Top, Bottom
@@ -400,31 +419,12 @@ def axis_name_to_index(self, axis: str) -> int:
  try:
  index = self.orientation.index(axis)
  except ValueError:
- index = self.orientation.index(self.flip_axis(axis))
+ index = self.orientation.index(FLIP_AXIS[axis])
  # Return negative indices so that it does not matter whether we
  # refer to spatial dimensions or not
  index = -3 + index
  return index
 
- @staticmethod
- def flip_axis(axis: str) -> str:
- """Return the opposite axis label. For example, ``'L'`` -> ``'R'``.
-
- Args:
- axis: Axis label, such as ``'L'`` or ``'left'``.
- """
- labels = 'LRPAISTBDV'
- first = labels[::2]
- last = labels[1::2]
- flip_dict = {a: b for a, b in zip(first + last, last + first)}
- axis = axis[0].upper()
- flipped_axis = flip_dict.get(axis)
- if flipped_axis is None:
- values = ', '.join(labels)
- message = f'Axis not understood. Please use one of: {values}'
- raise ValueError(message)
- return flipped_axis
-
  def get_spacing_string(self) -> str:
  strings = [f'{n:.2f}' for n in self.spacing]
  string = f'({", ".join(strings)})'

src/torchio/data/io.py

@@ -14,11 +14,9 @@
 from ..typing import TypeData
 from ..typing import TypeDataAffine
 from ..typing import TypeDirection
-from ..typing import TypeDoubletInt
 from ..typing import TypePath
 from ..typing import TypeQuartetInt
 from ..typing import TypeTripletFloat
-from ..typing import TypeTripletInt
 
 
 # Matrices used to switch between LPS and RAS
@@ -87,26 +85,43 @@ def _read_dicom(directory: TypePath):
 
 
 def read_shape(path: TypePath) -> TypeQuartetInt:
- reader = sitk.ImageFileReader()
- reader.SetFileName(str(path))
- reader.ReadImageInformation()
- num_channels = reader.GetNumberOfComponents()
- num_dimensions = reader.GetDimension()
+ try:
+ reader = sitk.ImageFileReader()
+ reader.SetFileName(str(path))
+ reader.ReadImageInformation()
+ num_channels = reader.GetNumberOfComponents()
+ num_dimensions = reader.GetDimension()
+ shape = reader.GetSize()
+ except RuntimeError as e: # try with NiBabel
+ message = f'Error loading image with SimpleITK:\n{e}\n\nTrying NiBabel...'
+ warnings.warn(message, stacklevel=2)
+ try:
+ obj: SpatialImage = nib.load(str(path)) # type: ignore[assignment]
+ except nib.loadsave.ImageFileError as e:
+ message = (
+ f'File "{path}" not understood.'
+ ' Check supported formats by at'
+ ' https://simpleitk.readthedocs.io/en/master/IO.html#images'
+ ' and https://nipy.org/nibabel/api.html#file-formats'
+ )
+ raise RuntimeError(message) from e
+ num_dimensions = obj.ndim
+ shape = obj.shape
+ num_channels = 1 if num_dimensions < 4 else shape[-1]
  assert 2 <= num_dimensions <= 4
  if num_dimensions == 2:
- spatial_shape_2d: TypeDoubletInt = reader.GetSize()
- assert len(spatial_shape_2d) == 2
- si, sj = spatial_shape_2d
+ assert len(shape) == 2
+ si, sj = shape
  sk = 1
  elif num_dimensions == 4:
  # We assume bad NIfTI file (channels encoded as spatial dimension)
- spatial_shape_4d: TypeQuartetInt = reader.GetSize()
- assert len(spatial_shape_4d) == 4
- si, sj, sk, num_channels = spatial_shape_4d
+ assert len(shape) == 4
+ si, sj, sk, num_channels = shape
  elif num_dimensions == 3:
- spatial_shape_3d: TypeTripletInt = reader.GetSize()
- assert len(spatial_shape_3d) == 3
- si, sj, sk = spatial_shape_3d
+ assert len(shape) == 3
+ si, sj, sk = shape
+ else:
+ raise ValueError(f'Unsupported number of dimensions: {num_dimensions}')
  shape = num_channels, si, sj, sk
  return shape
 

src/torchio/transforms/augmentation/intensity/random_ghosting.py

@@ -1,6 +1,5 @@
 from collections import defaultdict
 from typing import Dict
-from typing import Iterable
 from typing import Tuple
 from typing import Union
 
@@ -60,16 +59,7 @@ def __init__(
  **kwargs,
  ):
  super().__init__(**kwargs)
- if not isinstance(axes, tuple):
- try:
- axes = tuple(axes) # type: ignore[arg-type]
- except TypeError:
- axes = (axes,) # type: ignore[assignment]
- assert isinstance(axes, Iterable)
- for axis in axes:
- if not isinstance(axis, str) and axis not in (0, 1, 2):
- raise ValueError(f'Axes must be in (0, 1, 2), not "{axes}"')
- self.axes = axes
+ self.axes = self.parse_axes(axes)
  self.num_ghosts_range = self._parse_range(
  num_ghosts,
  'num_ghosts',
@@ -84,16 +74,13 @@ def __init__(
  self.restore = _parse_restore(restore)
 
  def apply_transform(self, subject: Subject) -> Subject:
+ axes = self.ensure_axes_indices(subject, self.axes)
  arguments: Dict[str, dict] = defaultdict(dict)
- if any(isinstance(n, str) for n in self.axes):
- subject.check_consistent_orientation()
- for name, image in self.get_images_dict(subject).items():
- is_2d = image.is_2d()
- axes = [a for a in self.axes if a != 2] if is_2d else self.axes
+ for name, _ in self.get_images_dict(subject).items():
  min_ghosts, max_ghosts = self.num_ghosts_range
  params = self.get_params(
+ axes,
  (int(min_ghosts), int(max_ghosts)),
- axes, # type: ignore[arg-type]
  self.intensity_range,
  )
  num_ghosts_param, axis_param, intensity_param = params
@@ -108,8 +95,8 @@ def apply_transform(self, subject: Subject) -> Subject:
 
  def get_params(
  self,
- num_ghosts_range: Tuple[int, int],
  axes: Tuple[int, ...],
+ num_ghosts_range: Tuple[int, int],
  intensity_range: Tuple[float, float],
  ) -> Tuple:
  ng_min, ng_max = num_ghosts_range
@@ -118,6 +105,17 @@ def get_params(
  intensity = self.sample_uniform(*intensity_range)
  return num_ghosts, axis, intensity
 
+ @staticmethod
+ def parse_restore(restore):
+ try:
+ restore = float(restore)
+ except ValueError as e:
+ raise TypeError(f'Restore must be a float, not "{restore}"') from e
+ if not 0 <= restore <= 1:
+ message = f'Restore must be a number between 0 and 1, not {restore}'
+ raise ValueError(message)
+ return restore
+
 
 class Ghosting(IntensityTransform, FourierTransform):
  r"""Add MRI ghosting artifact.