diff --git a/pyproject.toml b/pyproject.toml index 7016afc01..9e64d326e 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/torchio/data/image.py b/src/torchio/data/image.py index 637c7fddc..f3959bf06 100644 --- a/src/torchio/data/image.py +++ b/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. diff --git a/src/torchio/external/imports.py b/src/torchio/external/imports.py index 2e083db66..343cb24f5 100644 --- a/src/torchio/external/imports.py +++ b/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) diff --git a/src/torchio/visualization.py b/src/torchio/visualization.py index 2af006a58..634f14fbc 100644 --- a/src/torchio/visualization.py +++ b/src/torchio/visualization.py @@ -1,6 +1,7 @@ from __future__ import annotations import warnings +from pathlib import Path from typing import TYPE_CHECKING import numpy as np @@ -8,9 +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: @@ -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()