"""CLI for beamforming a zea dataset with a pipeline defined in a YAML config file.
Usage:
python -m zea.data.process --dataset <path> --config <config.yaml>
"""
import types
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from typing import Annotated, Sequence
import numpy as np
import tyro
from keras import ops
from zea import io_lib, log
from zea.backend import jit
from zea.cli_args import DEFAULT_DEVICE, DEVICE_HELP, SUPPORTED_FORMATS, ProcessArgs
from zea.config import Config
from zea.data.dataloader import Dataloader
from zea.data.datasets import Dataset
from zea.data.file import File, _GroupProxy
from zea.data.file_operations import output_blocked
from zea.data.spec import strip_track_prefix
from zea.func import translate
from zea.internal.checks import _NON_IMAGE_DATA_TYPES
from zea.internal.device import init_device
from zea.ops.imports import import_ops_modules
from zea.ops.pipeline import Pipeline
from zea.utils import FunctionTimer, ProgressBar
def _axis_selections_from_params(parameters) -> dict | None:
"""Return HDF5 axis_selections for transmit-axis pre-filtering, or None."""
_tx = getattr(parameters, "selected_transmits", None)
if _tx is not None:
return {1: sorted(int(t) for t in _tx)}
return None
sitk: types.ModuleType | None = None
try:
import SimpleITK as _sitk
sitk = _sitk
except ImportError:
pass
def _get_config_parameters(config: Config) -> dict:
"""Return the config parameters dict, handling missing or empty sections."""
params = getattr(config, "parameters", None)
if params is None:
return {}
return params.as_dict() if hasattr(params, "as_dict") else dict(params)
def _key_requires_pipeline(key: str) -> bool:
"""Return True if ``key`` holds raw RF/pre-beamformed data that needs a pipeline.
Normalizes the key so aliases like ``raw_data`` match ``_NON_IMAGE_DATA_TYPES``.
"""
normalized = (key or "").strip()
normalized = strip_track_prefix(normalized)
normalized = normalized.removeprefix("data/").removesuffix("/values")
return normalized in _NON_IMAGE_DATA_TYPES
def _resolve_track(file: File, key: str, track: str | None) -> tuple[int | None, str]:
"""Pick the track to process and return ``(track_index, data_key)``.
Resolution is delegated to :meth:`File.get_track`, which matches ``track`` against
the file's labels first and reads it as an index only when no label matches, so a
numeric label still wins. Files with a single track have nothing to address and
return ``(None, key)`` unchanged.
Raises:
ValueError: If the file has several tracks and none was selected, or the
requested track does not exist in this file.
"""
if file._n_tracks <= 1:
return None, key
labels = file.track_labels
if track is None:
raise ValueError(
f"{file.path} has {len(labels)} tracks {labels} but no track was selected. "
"Pass --track <label|index> to pick one."
)
try:
index = file.get_track(track).index
except KeyError as exc:
raise ValueError(f"No track {track!r} in {file.path}. Available tracks: {labels}.") from exc
return index, f"tracks/track_{index}/data/{strip_track_prefix(key).removeprefix('data/')}"
def _load_track_parameters(file: File, track_index: int | None):
"""Load the acquisition parameters of ``track_index`` (or of the whole file)."""
if track_index is None:
return file.load_parameters()
return file.tracks[track_index].load_parameters()
def _check_track_consistency(file: File, track_index: int | None, track_label: str | None) -> None:
"""Raise unless ``file`` carries the same track at ``track_index`` as the first file.
The data key resolved from the first file addresses its track by index, so a file
whose tracks are ordered differently is silently read at the wrong one. The label is
what reveals it: ``--track 1`` resolves to index 1 in every file however they are
ordered, so the index alone can never catch a reordering.
"""
if track_index is None:
return
labels = file.track_labels
label = labels[track_index] if track_index < len(labels) else None
if label != track_label:
raise ValueError(
f"Track at index {track_index} is {label!r} in {file.path} but {track_label!r} "
"in the first file; tracks must be ordered consistently across the dataset."
)
def _resolve_save_as(save_as: str | None, n_frames: int) -> str:
"""Pick the output format for a file of ``n_frames`` frames.
An explicit ``--save-as`` always wins. Left out, the format follows the data:
a single frame is a still, so it is written as a PNG rather than as a
one-frame GIF, and anything longer stays a GIF. The count is the number of
frames actually written, so a file that simply holds one frame gets a PNG
too, not only a run limited with ``--n-frames 1``.
"""
if save_as is not None:
return save_as
return "png" if n_frames == 1 else "gif"
def _save_frames_as_png(frames: np.ndarray, save_path: Path, overwrite: bool) -> None:
"""Write every frame of ``frames`` as its own PNG.
A single frame is written to ``save_path`` unchanged, which is what makes
``--save-as png --n-frames 1`` produce ``<filestem>.png``. Several frames are
numbered ``<filestem>_0000.png``, ``…_0001.png``, … so that the frames of one
file do not overwrite each other.
"""
paths = (
[save_path]
if len(frames) == 1
else [save_path.with_name(f"{save_path.stem}_{i:04d}.png") for i in range(len(frames))]
)
for frame, path in zip(frames, paths):
if output_blocked(path, overwrite):
log.warning(f"File {path} already exists. Use --overwrite to replace it.")
continue
io_lib.save_image(frame, path)
log.info(f"Saved {log.yellow(path)}")
def _run_passthrough(
dataset_path: str,
key: str,
n_frames: int | None,
save_dir: Path,
save_as: str | None,
overwrite: bool,
track_index: int | None = None,
track_label: str | None = None,
revision: str | None = None,
) -> None:
"""Save data frames directly without a beamforming pipeline."""
if save_as is not None and save_as not in ("gif", "mp4", "png", "hdf5"):
raise ValueError(f"Passthrough mode only supports gif/mp4/png/hdf5, got {save_as!r}")
save_dir.mkdir(parents=True, exist_ok=True)
with Dataset(dataset_path, lazy=True, _suggest_lazy=False, revision=revision) as ds:
pbar = ProgressBar(len(ds))
for i in range(len(ds)):
f = ds[i] # lazy download for hf:// paths; returns cached File handle
# ``key`` addresses the track by index, so every file must agree on which
# track sits there — the pipeline path checks the same thing.
_check_track_consistency(f, track_index, track_label)
data_key = f.format_key(key)
_dset = f.dataset(data_key)
# A map-style key such as "image" resolves to a group, not a dataset: read its
# values child, the same way load_file does.
if isinstance(_dset, _GroupProxy):
_dset = _dset.values
arr = np.asarray(_dset[:n_frames] if n_frames is not None else _dset[:])
filestem = f.stem
# Ensure (N, H, W) — squeeze any leading single-element dims
while arr.ndim > 3 and arr.shape[0] == 1:
arr = arr[0]
if arr.ndim == 2:
arr = arr[np.newaxis] # add frame axis
if arr.dtype != np.uint8:
lo, hi = float(arr.min()), float(arr.max())
arr = (
((arr - lo) / (hi - lo) * 255).astype(np.uint8)
if hi > lo
else np.zeros_like(arr, dtype=np.uint8)
)
fmt = _resolve_save_as(save_as, len(arr))
save_path = save_dir / f"{filestem}.{fmt}"
if fmt == "png":
# One path per frame, so the overwrite guard lives in the helper.
_save_frames_as_png(arr, save_path, overwrite)
elif output_blocked(save_path, overwrite):
log.warning(f"File {save_path} already exists. Use --overwrite to replace it.")
else:
if fmt in ("gif", "mp4"):
io_lib.save_video(arr, save_path, fps=20)
elif fmt == "hdf5":
File.create(save_path, data={"image": {"values": arr}}, overwrite=overwrite)
log.info(f"Saved {log.yellow(save_path)}")
pbar.add(1)
[docs]
def run_processing(
dataset_path: str,
config_path: str,
key: str,
n_frames: int | None,
save_dir: Path,
save_as: str | None = None,
keep_keys=("maxval",),
timings=False,
num_threads=16,
overwrite=False,
keep_dynamic_range=False,
revision: str | None = None,
config_revision: str | None = None,
track: str | None = None,
imports: Sequence[str] = (),
trust_remote_code: bool = False,
) -> None:
if keep_dynamic_range and save_as != "hdf5":
raise ValueError("--keep_dynamic_range is only supported with --save_as hdf5.")
if save_as == "nii.gz" and sitk is None:
raise ValueError("SimpleITK is not installed; cannot save as nii.gz.")
if save_as is not None and save_as not in SUPPORTED_FORMATS:
raise ValueError(f"save_as must be one of {SUPPORTED_FORMATS}, got {save_as!r}")
if config_revision is None:
config_revision = revision
config_revision = config_revision or None
# Imported after init_device has run (custom operations may touch keras at import
# time) and before anything resolves an operation name.
import_ops_modules(imports, trust_remote_code=trust_remote_code)
config = Config.from_path(config_path, revision=config_revision)
config_params = _get_config_parameters(config)
# Peek at the first file for the track to read (the Dataloader takes one key for
# the whole dataset, so ``data_key`` has to name it) and for selected_transmits,
# letting the dataloader pre-filter the transmit axis at HDF5 read time. lazy=True
# streams hf:// paths, fetching only the chunks the peek touches; validate=False
# because the Dataloader validates these same files anyway.
track_index: int | None = None
track_label: str | None = None
data_key = key
axis_selections: dict | None = None
with Dataset(
dataset_path, validate=False, lazy=True, _suggest_lazy=False, revision=revision
) as _peek_ds:
if len(_peek_ds) > 0:
_peek_f = _peek_ds[0]
track_index, data_key = _resolve_track(_peek_f, key, track)
if track_index is not None:
# Baseline for the per-file check below: every other file must carry the
# same track at this index, since ``data_key`` addresses it by index.
track_label = _peek_f.track_labels[track_index]
if _key_requires_pipeline(data_key):
try:
_peek_params = _load_track_parameters(_peek_f, track_index)
_peek_params.update(config_params)
axis_selections = _axis_selections_from_params(_peek_params)
except Exception:
pass # fall back to runtime slicing if the peek fails
try:
pipeline = Pipeline.from_path(
config_path,
with_batch_dim=False,
trust_remote_code=trust_remote_code,
revision=config_revision,
)
# A missing/unbuildable pipeline is recoverable for keys that need no beamforming.
# Failures from the config's own ``imports`` are not: those raise RemoteCodeError /
# OpsModuleImportError, which deliberately fall outside this except clause.
except (ValueError, KeyError) as exc:
if _key_requires_pipeline(data_key):
raise
log.warning(
f"No pipeline found in config ({exc}). "
f"Key '{data_key}' does not require beamforming — saving data as-is."
)
save_dir.mkdir(parents=True, exist_ok=True)
_run_passthrough(
dataset_path,
data_key,
n_frames,
save_dir,
save_as,
overwrite,
track_index=track_index,
track_label=track_label,
revision=revision,
)
return
save_dir.mkdir(parents=True, exist_ok=True)
dataloader = Dataloader(
dataset_path,
key=data_key,
batch_size=None,
shuffle=False,
return_metadata=True,
limit_n_frames=n_frames,
n_frames=None,
num_threads=num_threads,
sort_files=True,
dtype="float32",
axis_selections=axis_selections,
revision=revision,
)
iterator = iter(dataloader)
total_batches = len(dataloader)
get_data = lambda: next(iterator)
prepare_parameters = pipeline.prepare_parameters
pipeline_call = pipeline.__call__
if timings:
timer = FunctionTimer()
get_data = timer(get_data, name="dataloader")
prepare_parameters = timer(prepare_parameters, name="prepare_parameters")
pipeline_call = timer(pipeline_call, name="pipeline")
_DEFAULT_FPS = 20
prev_file_path = None
data_output = []
filestem = None
parameters = None
params = None
fps = _DEFAULT_FPS
def save_video_worker(
video: np.ndarray,
save_path: Path,
parameters,
fps: int,
fmt: str,
):
if fmt == "png":
# One path per frame, so the overwrite guard lives in the helper.
_save_frames_as_png(video, save_path, overwrite)
return
if output_blocked(save_path, overwrite):
log.warning(f"File {save_path} already exists. Use --overwrite to replace it.")
return
if fmt in ["mp4", "gif"]:
io_lib.save_video(video, save_path, fps=fps)
elif fmt == "hdf5":
scan_dict = parameters.to_scan_dict()
probe_dict = parameters.to_probe_dict()
File.create(
save_path,
data={"image": {"values": video}},
scan=scan_dict or None,
probe=probe_dict or None,
overwrite=overwrite,
)
elif fmt == "nii.gz":
assert sitk is not None, "SimpleITK must be installed to save as nii.gz"
sitk.WriteImage(sitk.GetImageFromArray(video), str(save_path))
log.info(f"Saved NIfTI to {log.yellow(save_path)}")
pbar = ProgressBar(total_batches)
@jit
def to_8bit(image, dynamic_range):
image = ops.nan_to_num(image, nan=dynamic_range[0])
image = translate(image, dynamic_range, (0, 255))
image = ops.clip(image, 0, 255)
image = ops.cast(image, "uint8")
return image
with ThreadPoolExecutor(max_workers=1) as executor:
save_future = None
for i in range(total_batches + 1):
if i < total_batches:
frame, metadata = get_data()
file_path = metadata["file"]["fullpath"]
else:
file_path = None # sentinel to flush the last file
if file_path != prev_file_path:
if prev_file_path is not None:
video = ops.convert_to_numpy(data_output)
fmt = _resolve_save_as(save_as, len(video))
save_path = save_dir / f"{filestem}.{fmt}"
if save_future is not None:
save_future.result()
save_future = executor.submit(
save_video_worker, video, save_path, parameters, fps, fmt
)
data_output = []
if file_path is None:
break
prev_file_path = file_path
# Already validated by the Dataloader that handed us this path.
with File(file_path, validate=False) as f:
filestem = f.stem
# Re-check rather than trust the peek: a file whose tracks are
# ordered differently would otherwise be read at the wrong index.
_check_track_consistency(f, track_index, track_label)
parameters = _load_track_parameters(f, track_index)
parameters.update(config_params)
try:
fps = int(round(parameters.frames_per_second))
except (ValueError, AttributeError):
fps = _DEFAULT_FPS
params = prepare_parameters(parameters)
# Sentinel iteration (no more data — also covers an empty dataset
# where total_batches == 0); nothing to process, so stop here.
if file_path is None:
break
output = pipeline_call(data=frame, **params) # ty: ignore[invalid-argument-type]
processed_frame = output["data"]
if not keep_dynamic_range:
dr = getattr(parameters, "dynamic_range", None)
dynamic_range = tuple(dr) if dr is not None else (-60, 0)
processed_frame = to_8bit(processed_frame, dynamic_range)
data_output.append(processed_frame)
pbar.add(1)
for keep_key in keep_keys:
if keep_key in output:
params[keep_key] = output[keep_key] # ty: ignore[invalid-assignment]
if timings:
for tname in timer.timings.keys():
timer.append_to_yaml(save_dir / f"timings_{tname}.yaml", tname)
# Re-raise anything the last save hit. Every other save is awaited before
# the next one is submitted, but the final one is not: leaving its result
# unread swallows the exception, and a run that wrote nothing at all still
# exits 0.
if save_future is not None:
save_future.result()
if timings:
timer.print()
@dataclass
class _StandaloneProcessArgs(ProcessArgs):
"""``ProcessArgs`` plus the ``--device`` flag that ``zea`` exposes globally."""
device: Annotated[
str,
tyro.conf.arg(help=DEVICE_HELP),
] = DEFAULT_DEVICE
[docs]
def main() -> None:
"""Entry point for ``python -m zea.data.process``, equivalent to ``zea process``."""
args = tyro.cli(_StandaloneProcessArgs)
init_device(args.device)
args.run()
if __name__ == "__main__":
main()