Add support to export images as videos (#1346)

Author: fepegar

pyproject.toml

@@ -50,6 +50,7 @@ dependencies = [
 [project.optional-dependencies]
 csv = ["pandas>=1"]
 plot = ["distinctipy>=1.3.4", "matplotlib>=3.4"]
+video = ["ffmpeg-python>=0.2.0"]
 
 [project.scripts]
 tiohd = "torchio.cli.print_info:app"

src/torchio/data/image.py

@@ -742,6 +742,13 @@ 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 get_center(self, lps: bool = False) -> TypeTripletFloat:
         """Get image center in RAS+ or LPS+ coordinates.
 
@@ -880,6 +887,36 @@ def hist(self, **kwargs) -> None:
         x = self.data.flatten().numpy()
         plot_histogram(x, **kwargs)
 
+    def to_video(
+        self,
+        output_path: TypePath,
+        frame_rate: float | None = 15,
+        seconds: float | None = None,
+        direction: str = 'I',
+        verbosity: str = 'error',
+    ) -> None:
+        """Create a video showing all image slices along a specified direction.
+
+        Args:
+            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.to_ras(),  # type: ignore[arg-type]
+            output_path,
+            frame_rate=frame_rate,
+            seconds=seconds,
+            direction=direction,
+            verbosity=verbosity,
+        )
+
 
 class LabelMap(Image):
     """Image whose pixel values represent segmentation labels.

src/torchio/external/imports.py

@@ -1,26 +1,45 @@
+from __future__ import annotations
+
 from importlib import import_module
 from importlib.util import find_spec
+from shutil import which
 from types import ModuleType
 
 
-def _check_package(*, package: str, extra: str) -> None:
-    if find_spec(package) is None:
+def _check_module(*, module: str, extra: str, package: str | None = None) -> None:
+    if find_spec(module) is None:
+        name = module if package is None else package
         message = (
-            f'The `{package}` package is required for this.'
+            f'The `{name}` package is required for this.'
             f' Install TorchIO with the `{extra}` extra:'
             f' `pip install torchio[{extra}]`.'
         )
         raise ImportError(message)
 
 
-def _check_and_import(package: str, extra: str) -> ModuleType:
-    _check_package(package=package, extra=extra)
-    return import_module(package)
+def _check_and_import(module: str, extra: str, **kwargs) -> ModuleType:
+    _check_module(module=module, extra=extra, **kwargs)
+    return import_module(module)
 
 
 def get_pandas() -> ModuleType:
-    return _check_and_import(package='pandas', extra='csv')
+    return _check_and_import(module='pandas', extra='csv')
 
 
 def get_distinctipy() -> ModuleType:
-    return _check_and_import(package='distinctipy', extra='plot')
+    return _check_and_import(module='distinctipy', extra='plot')
+
+
+def get_ffmpeg() -> ModuleType:
+    ffmpeg = _check_and_import(module='ffmpeg', extra='video', package='ffmpeg-python')
+    _check_executable('ffmpeg')
+    return ffmpeg
+
+
+def _check_executable(executable: str) -> None:
+    if which(executable) is None:
+        message = (
+            f'The `{executable}` executable is required for this. Install it from your'
+            ' package manager or download it from the official website.'
+        )
+        raise FileNotFoundError(message)

src/torchio/visualization.py

@@ -1,16 +1,23 @@
 from __future__ import annotations
 
 import warnings
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 import numpy as np
 import torch
 
 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:
@@ -282,3 +289,122 @@ def make_gif(
         duration=frame_duration_ms,
         loop=loop,
     )
+
+
+def make_video(
+    image: ScalarImage,
+    output_path: TypePath,
+    seconds: float | None = None,
+    frame_rate: float | None = None,
+    direction: str = 'I',
+    verbosity: str = 'error',
+) -> None:
+    ffmpeg = get_ffmpeg()
+
+    if seconds is None and frame_rate is None:
+        message = 'Either seconds or frame_rate must be provided.'
+        raise ValueError(message)
+    if seconds is not None and frame_rate is not None:
+        message = 'Provide either seconds or frame_rate, not both.'
+        raise ValueError(message)
+    if image.num_channels > 1:
+        message = 'Only single-channel tensors are supported for video output for now.'
+        raise ValueError(message)
+    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 = spacing_f, spacing_iso, spacing_iso
+        image = Resample(target_spacing)(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 NotImplementedError(message)
+
+    frames = image.numpy()[0]
+    first = frames[0]
+    height, width = first.shape
+
+    process = (
+        ffmpeg.input(
+            'pipe:',
+            format='rawvideo',
+            pix_fmt='gray',
+            s=f'{width}x{height}',
+            framerate=frame_rate,
+        )
+        .output(
+            str(output_path),
+            vcodec='libx265',
+            pix_fmt='yuv420p',
+            loglevel=verbosity,
+            **{'x265-params': f'log-level={verbosity}'},
+        )
+        .overwrite_output()
+        .run_async(pipe_stdin=True)
+    )
+
+    for array in frames:
+        buffer = array.tobytes()
+        process.stdin.write(buffer)
+
+    process.stdin.close()
+    process.wait()