"""Movement interpretation of PhysicsNeMo mesh-stage predictions.
:class:`WorkflowInferMovement` wraps a
:class:`physiotwin4d.WorkflowInferPhysicsNeMo` whose targets are per-point
displacements from the subject's reference mesh, and turns those raw predictions
into geometry: deformed meshes (``reference + displacement``), error statistics
in millimetres, and rasterized deformation / surface-normal fields.
The generic workflow stays target-agnostic; everything that assumes "the target
is a 3-vector displacement in mm" lives here. Composition, not inheritance, so
one decoder serves both the MeshGraphNet and the MLP inference methods.
"""
from __future__ import annotations
import csv
import logging
from collections.abc import Sequence
from pathlib import Path
from typing import Any, Literal, Optional, cast
import itk
import numpy as np
import pyvista as pv
from . import physicsnemo_tools as pnt
from .physiotwin4d_base import PhysioTwin4DBase
from .transform_tools import TransformTools
from .workflow_convert_vtk_to_usd import WorkflowConvertVTKToUSD
from .workflow_infer_physicsnemo import WorkflowInferPhysicsNeMo
[docs]
class WorkflowInferMovement(PhysioTwin4DBase):
"""Reconstruct geometry from displacement predictions.
The displacements are added to the caller's reference mesh when one is
available — the manifest's ``reference_mesh``, or the ``reference_mesh``
argument of the single-subject methods — which keeps the result in that
mesh's world frame. With no reference mesh they are added to the mesh
reconstructed from the PCA coefficients alone, which stays in the model's
PCA frame.
Args:
inference_workflow: A loaded :class:`WorkflowInferPhysicsNeMo` whose
model predicts three-component displacements.
log_level: Logging level. Default: ``logging.INFO``.
Raises:
ValueError: If the wrapped model does not predict exactly three
components, which cannot be a displacement.
"""
[docs]
def __init__(
self,
inference_workflow: WorkflowInferPhysicsNeMo,
log_level: int | str = logging.INFO,
) -> None:
super().__init__(class_name=self.__class__.__name__, log_level=log_level)
if inference_workflow.n_target != 3:
raise ValueError(
"WorkflowInferMovement needs a 3-component target, but the "
f"model predicts {inference_workflow.n_target} components."
)
self.inference_workflow = inference_workflow
self.model_directory = inference_workflow.model_directory
def _reference_points(
self, pca_coeffs: np.ndarray, reference_mesh: Optional[pv.DataSet]
) -> np.ndarray:
"""Return the points the displacements are added to.
Raises:
ValueError: If the supplied mesh has the wrong point count.
"""
if reference_mesh is None:
return self.inference_workflow.reference_points_from_coefficients(
pca_coeffs
)
points = np.asarray(reference_mesh.points, dtype=np.float32)
n_expected = self.inference_workflow.template_mesh.n_points
if points.shape[0] != n_expected:
raise ValueError(
f"Reference mesh has {points.shape[0]} points, expected "
f"{n_expected} (template topology)."
)
return points
# ─────────────────────────── Public API ────────────────────────────────
[docs]
def process(
self,
subject_manifest: Path,
stages: Optional[list[float]] = None,
output_directory: Optional[Path] = None,
) -> dict[str, Any]:
"""Predict a subject's deformed meshes from a manifest.
When ``stages`` is ``None`` every phase in the manifest is predicted and
compared against its ground truth (``reference + stored displacement``);
when ``stages`` is given those arbitrary stages are predicted without
comparison. The displacements are added to the manifest's
``reference_mesh`` points.
Args:
subject_manifest: Path to the subject manifest JSON.
stages: Optional list of stages to predict.
output_directory: Output directory; defaults to
``<model_directory>/<subject_id>``.
Returns:
Dict with ``subject_id``, ``predicted_surfaces`` (paths), and, in the
phase mode, ``statistics``, ``statistics_file`` and ``rmse_surface``.
"""
workflow = self.inference_workflow
manifest = pnt.parse_manifest(subject_manifest)
pca_coeffs = pnt.load_pca_coefficients(manifest.pca_coefficients)
ref_mesh = cast(pv.DataSet, pv.read(str(manifest.reference_mesh)))
ref_points = self._reference_points(pca_coeffs, ref_mesh)
out_dir = (
Path(output_directory)
if output_directory is not None
else self.model_directory / manifest.subject_id
)
out_dir.mkdir(parents=True, exist_ok=True)
sid = manifest.subject_id
self.log_section("INFER MOVEMENT [%s]", sid)
suffix = ".vtp" if isinstance(ref_mesh, pv.PolyData) else ".vtu"
requested = stages if stages is not None else [p.stage for p in manifest.phases]
surfaces: list[Path] = []
stats: list[dict] = []
sq_err_sum = np.zeros(ref_points.shape[0], dtype=np.float64)
for index, stage in enumerate(requested):
pred_points = ref_points + workflow.predict(pca_coeffs, stage)
pred_mesh = ref_mesh.copy(deep=True)
pred_mesh.points = pred_points
path = out_dir / f"{sid}_s{int(stage * 100):03d}_pred{suffix}"
pred_mesh.save(str(path))
surfaces.append(path)
if stages is not None:
self.log_info("stage %.3f -> %s", stage, path.name)
continue
stored = pnt.load_target_array(
manifest.phases[index].mesh, manifest.target_array
)
if stored.shape != ref_points.shape:
raise ValueError(
f"Stored '{manifest.target_array}' targets in "
f"{manifest.phases[index].mesh} have shape {stored.shape}, "
f"expected {ref_points.shape} to be displacements of the "
"reference mesh."
)
actual = ref_points + stored
euclidean = np.linalg.norm(pred_points - actual, axis=1)
sq_err_sum += euclidean.astype(np.float64) ** 2
stats.append(self._error_row(sid, stage, pred_points, actual))
self.log_info(
"stage %.3f: mean=%.3f mm max=%.3f mm",
stage,
stats[-1]["mean_error_mm"],
stats[-1]["max_error_mm"],
)
result: dict[str, Any] = {"subject_id": sid, "predicted_surfaces": surfaces}
if stages is None:
point_rmse = np.sqrt(sq_err_sum / len(requested)).astype(np.float32)
rmse_mesh = ref_mesh.copy(deep=True)
rmse_mesh.point_data["RMSE_mm"] = point_rmse
rmse_file = out_dir / f"{sid}_rmse{suffix}"
rmse_mesh.save(str(rmse_file))
stats_file = out_dir / "statistics_per_phase.csv"
with stats_file.open("w", newline="", encoding="utf-8") as fh:
writer = csv.DictWriter(fh, fieldnames=list(stats[0].keys()))
writer.writeheader()
writer.writerows(stats)
result["rmse_surface"] = rmse_file
result["statistics"] = stats
result["statistics_file"] = stats_file
return result
[docs]
def predict_single(
self,
shape_parameters: Path,
stage: float,
reference_mesh: Optional[Path] = None,
ground_truth: Optional[Path] = None,
output_directory: Optional[Path] = None,
) -> dict[str, Any]:
"""Predict one subject at one stage without a manifest.
Without ``reference_mesh`` the subject reference is reconstructed from
the PCA shape parameters (``P = mean + Σ b_i·std_i·eigenvector_i``) in
the SSM/PCA frame, so the prediction is self-consistent with no
reference mesh file at all; with one, the prediction stays in that
mesh's world frame.
Args:
shape_parameters: JSON file with the subject PCA coefficient vector.
stage: Target stage to predict.
reference_mesh: The subject's reference mesh; omit to displace the
PCA reconstruction instead.
ground_truth: Optional mesh whose points are the true stage
positions, for error reporting.
output_directory: Output directory; defaults to
``<model_directory>/single_prediction``.
Returns:
Dict with ``predicted_surface`` (path), ``predicted_points``, and,
when ``ground_truth`` is supplied, ``statistics``.
"""
workflow = self.inference_workflow
coeffs = pnt.load_pca_coefficients(shape_parameters)
ref_mesh = (
cast(pv.DataSet, pv.read(str(reference_mesh)))
if reference_mesh is not None
else None
)
ref_points = self._reference_points(coeffs, ref_mesh)
pred_points = ref_points + workflow.predict(coeffs, stage)
out_dir = (
Path(output_directory)
if output_directory is not None
else self.model_directory / "single_prediction"
)
out_dir.mkdir(parents=True, exist_ok=True)
template = workflow.template_mesh
pred_mesh = template.copy(deep=True)
pred_mesh.points = pred_points
suffix = ".vtp" if isinstance(template, pv.PolyData) else ".vtu"
stem = Path(shape_parameters).stem
path = out_dir / f"{stem}_pred_s{int(stage * 100):03d}{suffix}"
pred_mesh.save(str(path))
self.log_info("single prediction stage %.3f -> %s", stage, path.name)
result: dict[str, Any] = {
"predicted_surface": path,
"predicted_points": pred_points,
}
if ground_truth is not None:
actual = np.asarray(pv.read(str(ground_truth)).points, dtype=np.float32)
result["statistics"] = self._error_row(stem, stage, pred_points, actual)
self.log_info(
"ground-truth error: mean=%.3f mm max=%.3f mm",
result["statistics"]["mean_error_mm"],
result["statistics"]["max_error_mm"],
)
return result
[docs]
def process_time_series(
self,
shape_parameters: Path,
stages: Sequence[float],
output_directory: Path,
reference_mesh: Optional[Path] = None,
ground_truth: Optional[Sequence[Path]] = None,
reference_image: Optional[itk.Image] = None,
warp_interpolation: str = "linear",
warp_background_value: float = 0.0,
smoothing_sigma_mm: float = 10.0,
usd_project_name: Optional[str] = None,
anatomy_type: Optional[str] = None,
separate_by_connectivity: bool = False,
) -> dict[str, Any]:
"""Predict one subject across a whole time series and write its geometry.
One prediction per entry of ``stages``, each written as a mesh. When
``reference_image`` is supplied, each stage also gets a deformation
field, which is Gaussian-smoothed into a continuous
:class:`itk.DisplacementFieldTransform` and used to carry
``reference_image`` into that stage's frame. The smoothing spreads a
surface-shell field into the volume, so the warped image is an
interpolation of the surface motion, not an independent registration.
Args:
shape_parameters: JSON file with the subject PCA coefficient vector.
stages: Stages to predict, in the order they are to be animated.
output_directory: Directory every artifact is written to.
reference_mesh: The subject's reference mesh; omit to displace the
PCA reconstruction instead.
ground_truth: One mesh per stage whose points are the true stage
positions, for error reporting. Must align with ``stages``.
reference_image: Image carried through each stage's deformation, and
the grid the deformation field is rasterized on. Omit to write
meshes only.
warp_interpolation: Interpolation used to resample
``reference_image``: ``"linear"`` for intensity images,
``"nearest"`` for labelmaps and masks.
warp_background_value: Value written where a stage's grid samples
outside ``reference_image``. ``0.0`` suits labelmaps; CT needs
``-1000.0``, which is air in Hounsfield units.
smoothing_sigma_mm: Gaussian sigma, in millimeters, that turns the
sparse surface-shell field into a continuous deformation.
usd_project_name: When given, the stage meshes are also written as
one animated USD under this name, one time sample per stage.
anatomy_type: Anatomy whose materials color that USD.
separate_by_connectivity: Whether that USD splits each frame into
separate objects by connectivity.
Returns:
Dict with ``stages``, ``predicted_surfaces``, ``warped_images``,
``transforms``, ``usd_file``, ``statistics`` and
``statistics_file``. Entries that were not requested are empty
lists or ``None``.
Raises:
ValueError: If ``stages`` is empty, or ``ground_truth`` is given
with a different length.
"""
if not stages:
raise ValueError("process_time_series needs at least one stage.")
if ground_truth is not None and len(ground_truth) != len(stages):
raise ValueError(
f"ground_truth has {len(ground_truth)} entries but there are "
f"{len(stages)} stages."
)
workflow = self.inference_workflow
coeffs = pnt.load_pca_coefficients(shape_parameters)
ref_mesh = (
cast(pv.DataSet, pv.read(str(reference_mesh)))
if reference_mesh is not None
else None
)
ref_points = self._reference_points(coeffs, ref_mesh)
template = ref_mesh if ref_mesh is not None else workflow.template_mesh
out_dir = Path(output_directory)
out_dir.mkdir(parents=True, exist_ok=True)
stem = Path(shape_parameters).stem
suffix = ".vtp" if isinstance(template, pv.PolyData) else ".vtu"
self.log_section("INFER MOVEMENT TIME SERIES [%s]", stem)
transform_tools = TransformTools(log_level=self.log_level)
stage_meshes: list[pv.DataSet] = []
surfaces: list[Path] = []
warped_images: list[Path] = []
transforms: list[itk.Transform] = []
stats: list[dict] = []
for index, stage in enumerate(stages):
tag = f"s{int(stage * 100):03d}"
pred_points = ref_points + workflow.predict(coeffs, stage)
pred_mesh = template.copy(deep=True)
pred_mesh.points = pred_points
surface_file = out_dir / f"{stem}_{tag}_pred{suffix}"
pred_mesh.save(str(surface_file))
stage_meshes.append(pred_mesh)
surfaces.append(surface_file)
if reference_image is not None:
field = self.create_deformation_field(
shape_parameters=shape_parameters,
stage=stage,
reference_image=reference_image,
reference_mesh=reference_mesh,
direction="inverse",
)
transform = transform_tools.smooth_deformation_field_transform(
field["deformation_field"],
sigma=smoothing_sigma_mm,
weight_image=field["weight_image"],
)
transforms.append(transform)
warped = transform_tools.transform_image(
reference_image,
transform,
reference_image=reference_image,
interpolation_method=warp_interpolation,
background_value=warp_background_value,
)
warped_file = out_dir / f"{stem}_{tag}_warped.mha"
itk.imwrite(warped, str(warped_file), compression=True)
warped_images.append(warped_file)
if ground_truth is not None:
actual = np.asarray(
pv.read(str(ground_truth[index])).points, dtype=np.float32
)
stats.append(self._error_row(stem, stage, pred_points, actual))
self.log_info(
"stage %.3f: mean=%.3f mm max=%.3f mm",
stage,
stats[-1]["mean_error_mm"],
stats[-1]["max_error_mm"],
)
else:
self.log_info("stage %.3f -> %s", stage, surface_file.name)
statistics_file: Optional[Path] = None
if stats:
statistics_file = out_dir / "statistics_per_stage.csv"
with statistics_file.open("w", newline="", encoding="utf-8") as fh:
writer = csv.DictWriter(fh, fieldnames=list(stats[0].keys()))
writer.writeheader()
writer.writerows(stats)
usd_file: Optional[Path] = None
if usd_project_name is not None:
usd_workflow = WorkflowConvertVTKToUSD(
input_meshes=stage_meshes,
usd_project_name=usd_project_name,
output_directory=out_dir,
appearance="anatomy" if anatomy_type is not None else "solid",
anatomy_type=anatomy_type,
separate_by_connectivity=separate_by_connectivity,
frames_per_second=float(len(stage_meshes)),
log_level=self.log_level,
)
usd_file = Path(usd_workflow.process()["usd_file"])
return {
"stages": list(stages),
"predicted_surfaces": surfaces,
"warped_images": warped_images,
"transforms": transforms,
"usd_file": usd_file,
"statistics": stats,
"statistics_file": statistics_file,
}
@staticmethod
def _vector_image_like(array: np.ndarray, reference_image: itk.Image) -> itk.Image:
"""Wrap a ``(z, y, x, 3)`` array as an ITK vector image on ``reference``'s grid."""
image = itk.image_from_array(np.ascontiguousarray(array), is_vector=True)
image.SetSpacing(reference_image.GetSpacing())
image.SetOrigin(reference_image.GetOrigin())
image.SetDirection(reference_image.GetDirection())
return image
@staticmethod
def _scalar_image_like(array: np.ndarray, reference_image: itk.Image) -> itk.Image:
"""Wrap a ``(z, y, x)`` array as an ITK scalar image on ``reference``'s grid."""
image = itk.image_from_array(np.ascontiguousarray(array))
image.SetSpacing(reference_image.GetSpacing())
image.SetOrigin(reference_image.GetOrigin())
image.SetDirection(reference_image.GetDirection())
return image
@staticmethod
def _error_row(
subject_id: str, stage: float, pred: np.ndarray, actual: np.ndarray
) -> dict:
"""Per-phase error statistics between predicted and actual points."""
errors = pred - actual
euclidean = np.linalg.norm(errors, axis=1)
return {
"subject_id": subject_id,
"stage": stage,
"n_points": int(len(euclidean)),
"mean_error_mm": float(euclidean.mean()),
"median_error_mm": float(np.median(euclidean)),
"max_error_mm": float(euclidean.max()),
"rms_error_mm": float(np.sqrt(np.mean(euclidean**2))),
"std_error_mm": float(euclidean.std()),
"mean_abs_error_x_mm": float(np.abs(errors[:, 0]).mean()),
"mean_abs_error_y_mm": float(np.abs(errors[:, 1]).mean()),
"mean_abs_error_z_mm": float(np.abs(errors[:, 2]).mean()),
}