"""
GIF export utilities for AFM image stacks.
This module provides functions for generating animated GIFs from AFM image stacks,
with optional timestamps and scale bars. Frames can be normalized automatically or
scaled using a fixed z-range.
Dependencies
------------
- matplotlib
- numpy
- Pillow (PIL)
"""
import logging
from pathlib import Path
from matplotlib import colormaps as cm
from PIL import Image
from playnano.io.render_utils import (
DEFAULT_FONT_SCALE,
render_frame,
resolve_stack_data,
)
from playnano.utils.colormaps import DEFAULT_CMAP
from playnano.utils.io_utils import (
compute_zscale_range,
prepare_output_directory,
sanitize_output_name,
)
logger = logging.getLogger(__name__)
[docs]
def create_gif_with_scale_and_timestamp(
image_stack,
pixel_sizes_nm,
timestamps=None,
scale_bar_length_nm=100,
output_path="output",
fps: float = 5.0,
cmap_name=DEFAULT_CMAP,
zmin: float | str | None = None,
zmax: float | str | None = None,
draw_ts: bool = True,
draw_scale: bool = True,
font_scale: float = DEFAULT_FONT_SCALE,
):
"""
Create an animated GIF from an AFM image stack with optional overlays.
Frames are normalized, colorized using a matplotlib colormap, and annotated
with a scale bar and timestamps before being compiled into a GIF.
Parameters
----------
image_stack : np.ndarray
3D array of shape (N, H, W) representing the AFM image stack.
pixel_sizes_nm : list
List containing the size of a pixel in nanometers for each frame.
timestamps : list[float] or tuple[float], optional
Timestamps for each frame in seconds. If ``None`` or invalid,
frame indices are used.
scale_bar_length_nm : int
Length of the scale bar in nanometers. Default is 100.
output_path : str
Path where the GIF will be saved. Default is 'output'.
fps : float
Playback frame rate in frames per second. Default is 5.0.
GIF format has limited temporal resolution; the actual frame
delay is rounded to the nearest millisecond.
cmap_name : str
Name of the matplotlib colormap to apply. Default is 'afmhot'.
zmin : float or str or None, optional
Minimum z-value mapped to colormap low end.
The string literal ``"auto"`` uses the 1st percentile.
zmax : float or str or None, optional
Maximum z-value mapped to colormap high end.
The string literal ``"auto"`` uses the 99th percentile.
draw_ts : bool
Whether to draw timestamps. Default is True.
draw_scale : bool
Whether to draw a scale bar. Default is True.
font_scale : float
Base font scale for annotations.
Raises
------
ValueError
If ``zmin`` equals ``zmax`` or ``timestamps`` have incorrect shape.
Returns
-------
None
Notes
-----
- Timestamps and scale bars are drawn in white.
- Frames are normalized globally if ``zmin`` and ``zmax`` are provided;
otherwise, per-frame.
"""
frames = []
cmap = cm.get_cmap(cmap_name)
# Validate timestamps
if (
timestamps is not None
and isinstance(timestamps, (list, tuple))
and len(timestamps) == len(image_stack)
):
has_valid_timestamps = True
else:
has_valid_timestamps = False
logger.warning(
"Invalid timestamps provided; frame indices will be used instead."
)
# Validate pixel sizes
if not (
pixel_sizes_nm is not None
and isinstance(pixel_sizes_nm, list)
and len(pixel_sizes_nm) == len(image_stack)
):
draw_scale = False
logger.warning("Invalid pixel_sizes_nm list; scale bar will be omitted.")
if zmin is not None or zmax is not None:
zmin_val, zmax_val = compute_zscale_range(image_stack, zmin, zmax)
else:
zmin_val, zmax_val = None, None
for i, frame in enumerate(image_stack):
# Determine timestamps
if has_valid_timestamps:
try:
timestamp = float(timestamps[i])
except (TypeError, ValueError, IndexError):
timestamp = i
else:
logger.warning(
f"Invalid timestamps provided, using frame index {i} as timestamp."
)
timestamp = i
pixel_size_nm = pixel_sizes_nm[i] if draw_scale else 1.0
frame_annotated = render_frame(
frame=frame,
cmap=cmap,
timestamp=timestamp,
pixel_size_nm=pixel_size_nm,
scale_bar_length_nm=scale_bar_length_nm,
font_scale=font_scale,
draw_ts=draw_ts,
draw_scale=draw_scale,
zmin_val=zmin_val,
zmax_val=zmax_val,
)
frames.append(Image.fromarray(frame_annotated))
duration_ms = int((1.0 / fps) * 1000) if fps > 0 else 500
frames[0].save(
output_path,
save_all=True,
append_images=frames[1:],
duration=duration_ms,
loop=0,
)
logger.info(f"GIF saved to {output_path}")
[docs]
def export_gif(
afm_stack,
make_gif: bool,
output_folder: str | None,
output_name: str | None,
scale_bar_nm: int | None,
cmap_name: str = DEFAULT_CMAP,
raw: bool = False,
fps: float = 5.0,
zmin: float | None = None,
zmax: float | None = None,
draw_ts: bool = True,
draw_scale: bool = True,
font_scale: float = DEFAULT_FONT_SCALE,
) -> None:
"""
Export an AFM image stack as an annotated GIF.
Parameters
----------
afm_stack : AFMImageStack
AFM stack object containing raw and/or processed data.
make_gif : bool
Whether to generate the GIF. If ``False``, the function exits immediately.
output_folder : str or None
Directory to save the GIF. Defaults to ``"output"`` if ``None``.
output_name : str or None
Base name for the GIF file. If ``None``, derived from the stack file name.
scale_bar_nm : int or None
Length of the scale bar in nanometers. Defaults to 100 nm.
raw : bool
If ``True``, export raw (unprocessed) data; otherwise export processed data
if available. Default is False.
fps : float
Frame rate for the GIF in frames per second. Default is 5.0.
zmin : float or None, optional
Minimum z-value mapped to colormap low end. The string literal ``"auto"``
can also be used to automatically set the 1st percentile. ``None`` uses
the minimum value of the data.
zmax : float or None, optional
Maximum z-value mapped to colormap high end. The string literal ``"auto"``
can also be used to automatically set the 99th percentile. ``None`` uses
the maximum value of the data.
draw_ts : bool
Whether to draw timestamps on each frame. Default is True.
draw_scale : bool
Whether to draw a scale bar on each frame. Default is True.
font_scale : float
Base font scale for annotations at the reference height.
Returns
-------
None
Notes
-----
- Uses processed data if available; otherwise falls back to raw data.
- Timestamps and pixel size are read from ``afm_stack`` metadata although if raw
data is exported after an edit_stack processing step then the timestamps in
``afm_stack.state_backups['frame_metadata_before_edit']`` are retrieved and used.
- Output file name includes ``"_filtered"`` if processed data is exported.
"""
if not make_gif:
return
out_dir = prepare_output_directory(output_folder, default="output")
base = sanitize_output_name(output_name, Path(afm_stack.file_path).stem)
stack_data, meta_src, is_filtered = resolve_stack_data(afm_stack, raw)
if is_filtered:
base = f"{base}_filtered"
timestamps = [md["timestamp"] for md in meta_src]
pixels_to_nm = [
md.get("frame_pixel_size_nm", afm_stack.pixel_size_nm) for md in meta_src
]
gif_path = out_dir / f"{base}.gif"
bar_nm = scale_bar_nm if scale_bar_nm is not None else 100
logger.debug(f"[export] Writing GIF → {gif_path}")
create_gif_with_scale_and_timestamp(
stack_data,
pixels_to_nm,
timestamps,
output_path=gif_path,
scale_bar_length_nm=bar_nm,
cmap_name=cmap_name,
zmin=zmin,
zmax=zmax,
draw_ts=draw_ts,
draw_scale=draw_scale,
font_scale=font_scale,
fps=fps,
)
logger.debug(f"[export] GIF written to {gif_path}")