Add many features

Author: fepegar

src/torchio/data/image.py

@@ -742,29 +742,41 @@ def to_gif(
             reverse=reverse,
         )
 
+    def to_ras(self) -> Image:
+        if self.orientation != tuple('RAS'):
+            from ..transforms.preprocessing.spatial.to_canonical import ToCanonical
+
+            return ToCanonical()(self)
+        return self
+
     def to_video(
         self,
         output_path: TypePath,
-        duration: float | None = None,
-        frame_rate: float | None = None,
+        frame_rate: float | None = 15,
+        seconds: float | None = None,
+        direction: str = 'I',
+        verbosity: str = 'error',
     ) -> None:
-        """Save a video of the image.
-
-        TODO: add usage example, with recommended transforms (rescale, LPS, etc.)
+        """Create a video showing all image slices along a specified direction.
 
         Args:
-            axis: Spatial axis (0, 1 or 2).
-            duration: Duration of the full video in seconds.
-            frame_rate: Number of frames per second.
             output_path: Path to the output video file.
+            frame_rate: Number of frames per second (FPS).
+            seconds: Target duration of the full video.
+            direction:
+            verbosity:
+
+        .. note:: Only ``frame_rate`` or ``seconds`` may (and must) be specified.
         """
         from ..visualization import make_video  # avoid circular import
 
         make_video(
-            self.data,
+            self.to_ras(),
             output_path,
-            duration=duration,
             frame_rate=frame_rate,
+            seconds=seconds,
+            direction=direction,
+            verbosity=verbosity,
         )
 
     def get_center(self, lps: bool = False) -> TypeTripletFloat:

src/torchio/visualization.py

@@ -9,10 +9,15 @@
 
 from .data.image import Image
 from .data.image import LabelMap
+from .data.image import ScalarImage
 from .data.subject import Subject
 from .external.imports import get_ffmpeg
 from .transforms.preprocessing.intensity.rescale import RescaleIntensity
+from .transforms.preprocessing.intensity.to import To
+from .transforms.preprocessing.spatial.ensure_shape_multiple import EnsureShapeMultiple
+from .transforms.preprocessing.spatial.resample import Resample
 from .transforms.preprocessing.spatial.to_canonical import ToCanonical
+from .transforms.preprocessing.spatial.to_orientation import ToOrientation
 from .types import TypePath
 
 if TYPE_CHECKING:
@@ -287,33 +292,94 @@ def make_gif(
 
 
 def make_video(
-    tensor: torch.Tensor,
+    image: ScalarImage,
     output_path: TypePath,
-    duration: float | None = None,
+    seconds: float | None = None,
     frame_rate: float | None = None,
+    direction: str = 'I',
+    verbosity: str = 'error',
 ) -> None:
-    """Encode a 3D array into an MP4 video."""
     ffmpeg = get_ffmpeg()
 
-    if duration is None and frame_rate is None:
-        message = 'Either duration or frame_rate must be provided.'
+    if seconds is None and frame_rate is None:
+        message = 'Either seconds or frame_rate must be provided.'
         raise ValueError(message)
-    if duration is not None and frame_rate is not None:
-        message = 'Provide either duration or frame_rate, not both.'
+    if seconds is not None and frame_rate is not None:
+        message = 'Provide either seconds or frame_rate, not both.'
         raise ValueError(message)
-    if len(tensor) > 1:
+    if image.num_channels > 1:
         message = 'Only single-channel tensors are supported for video output for now.'
         raise ValueError(message)
-    frames = tensor.numpy()[0].T
-    num_frames = len(frames)
-    if duration is not None:
-        frame_rate = num_frames / duration
+    tmin, tmax = image.data.min(), image.data.max()
+    if tmin < 0 or tmax > 255:
+        message = (
+            'The tensor must be in the range [0, 256) for video output.'
+            ' The image data will be rescaled to this range.'
+        )
+        warnings.warn(message, RuntimeWarning, stacklevel=2)
+        image = RescaleIntensity((0, 255))(image)
+    if image.data.dtype != torch.uint8:
+        message = (
+            'Only uint8 tensors are supported for video output. The image data'
+            ' will be cast to uint8.'
+        )
+        warnings.warn(message, RuntimeWarning, stacklevel=2)
+        image = To(torch.uint8)(image)
+
+    # Reorient so the output looks like in typical visualization software
+    direction = direction.upper()
+    if direction == 'I':  # axial top to bottom
+        target = 'IPL'
+    elif direction == 'S':  # axial bottom to top
+        target = 'SPL'
+    elif direction == 'A':  # coronal back to front
+        target = 'AIL'
+    elif direction == 'P':  # coronal front to back
+        target = 'PIL'
+    elif direction == 'R':  # sagittal left to right
+        target = 'RIP'
+    elif direction == 'L':  # sagittal right to left
+        target = 'LIP'
+    else:
+        message = (
+            'Direction must be one of "I", "S", "P", "A", "R" or "L".'
+            f' Got {direction!r}.'
+        )
+        raise ValueError(message)
+    image = ToOrientation(target)(image)
+
+    # Check isotropy
+    spacing_f, spacing_h, spacing_w = image.spacing
+    if spacing_h != spacing_w:
+        message = (
+            'The height and width spacings should be the same video output.'
+            f' Got {spacing_h:.2f} and {spacing_w:.2f}.'
+            f' Resampling both to {spacing_f:.2f}.'
+        )
+        warnings.warn(message, RuntimeWarning, stacklevel=2)
+        spacing_iso = min(spacing_h, spacing_w)
+        target = spacing_f, spacing_iso, spacing_iso
+        image = Resample(target)(image)  # type: ignore[assignment]
+
+    # Check that height and width are multiples of 2 for H.265 encoding
+    num_frames, height, width = image.spatial_shape
+    if height % 2 != 0 or width % 2 != 0:
+        message = (
+            f'The height ({height}) and width ({width}) must be even.'
+            ' The image will be cropped to the nearest even number.'
+        )
+        warnings.warn(message, RuntimeWarning, stacklevel=2)
+        image = EnsureShapeMultiple((1, 2, 2), method='crop')(image)
+
+    if seconds is not None:
+        frame_rate = num_frames / seconds
 
     output_path = Path(output_path)
     if output_path.suffix.lower() != '.mp4':
         message = 'Only .mp4 files are supported for video output.'
-        raise ValueError(message)
+        raise NotImplementedError(message)
 
+    frames = image.numpy()[0]
     first = frames[0]
     height, width = first.shape
 
@@ -327,8 +393,10 @@ def make_video(
         )
         .output(
             str(output_path),
-            vcodec='libx264',
+            vcodec='libx265',
             pix_fmt='yuv420p',
+            loglevel=verbosity,
+            **{'x265-params': f'log-level={verbosity}'},
         )
         .overwrite_output()
         .run_async(pipe_stdin=True)