"""Optional output writers for a :class:`DeflectionResult`. The functions in this module are entirely optional — calling :func:`underframe_deflection.compute_deflection` does not write anything to disk. Pull in only the writers you need. Plot rendering requires matplotlib (install with the ``viz`` extra: ``pip install underframe-deflection[viz]``); JSON/CSV writers have no extra deps. All writers expect the ``DeflectionResult`` to come from the same package version that produced it; do not pass in a result whose ``_artifacts`` field has been tampered with. """ from __future__ import annotations from pathlib import Path from typing import TYPE_CHECKING from .api import DeflectionResult from .beams import candidate_rows from .io import ensure_output_dir, write_csv, write_json from .ply_overlay import export_measurement_overlay_plys from .structures import TrackSolution if TYPE_CHECKING: pass # --------------------------------------------------------------------------- # JSON / CSV writers (stdlib only) # --------------------------------------------------------------------------- def write_report_json(result: DeflectionResult, path: str | Path) -> Path: """Write the full report to ``path`` as JSON. The structure matches ``DeflectionResult.to_dict()``. Schema is versioned by ``result.schema_version``. """ output_path = Path(path) ensure_output_dir(output_path.parent) write_json(output_path, result.to_dict()) return output_path def write_curves_csv(result: DeflectionResult, path: str | Path) -> Path: """Write the sampled deflection curves to a CSV file. Columns: ``x``, plus per-side beam/reference Y/Z, deflection, and confidence. NaN values indicate gaps where the curve is undefined. """ output_path = Path(path) ensure_output_dir(output_path.parent) artifacts = _require_artifacts(result) rows = artifacts.deflection_rows fieldnames = ( list(rows[0].keys()) if rows else [ "x", "left_beam_y", "left_beam_z", "left_ref_y", "left_ref_z", "left_deflection", "left_confidence", "right_beam_y", "right_beam_z", "right_ref_y", "right_ref_z", "right_deflection", "right_confidence", ] ) write_csv(output_path, fieldnames=fieldnames, rows=rows) return output_path def write_beam_tracks_csv(result: DeflectionResult, path: str | Path) -> Path: """Write the per-slice beam candidate table to a CSV file. Useful when debugging tracking behaviour or visualising candidate bands in external tools. """ output_path = Path(path) ensure_output_dir(output_path.parent) artifacts = _require_artifacts(result) rows = _annotate_beam_track_rows( candidate_rows(artifacts.slices), artifacts.left_tracks, artifacts.right_tracks, ) fieldnames = ( list(rows[0].keys()) if rows else [ "slice_index", "candidate_id", "side", "x_start", "x_end", "x_center", "y_min", "y_max", "y_center", "z_low", "width", "point_count", "low_point_count", "flatness", "normal_alignment", "score", "y_mid", "selected_rank", "selected_side", ] ) write_csv(output_path, fieldnames=fieldnames, rows=rows) return output_path # --------------------------------------------------------------------------- # Plots (matplotlib — optional dependency) # --------------------------------------------------------------------------- def render_plots( result: DeflectionResult, out_dir: str | Path, *, dpi: int | None = None, max_scatter_points: int | None = None, ) -> dict[str, Path]: """Render diagnostic PNG plots into ``out_dir``. Returns a dict mapping plot name to the written file path. Requires matplotlib; the import is deferred so this module can be imported without it as long as ``render_plots`` is not called. """ out_dir_path = ensure_output_dir(out_dir) artifacts = _require_artifacts(result) cfg = artifacts.resolved_config["visualization"] resolved_dpi = dpi if dpi is not None else cfg["dpi"] resolved_limit = ( max_scatter_points if max_scatter_points is not None else cfg["max_scatter_points"] ) from .visualization import ( plot_beam_candidates, plot_deflection_curves, plot_frame_hypotheses, plot_representatives_and_lines, plot_roi_overview, ) paths: dict[str, Path] = {} paths["frame_hypotheses"] = out_dir_path / "frame_hypotheses.png" plot_frame_hypotheses(artifacts.frame_hypotheses, paths["frame_hypotheses"], dpi=resolved_dpi) paths["bottom_roi"] = out_dir_path / "bottom_roi.png" plot_roi_overview( artifacts.work_vehicle, artifacts.work_roi, paths["bottom_roi"], dpi=resolved_dpi, limit=resolved_limit, ) paths["beam_tracks"] = out_dir_path / "beam_tracks.png" plot_beam_candidates( artifacts.slices, artifacts.best_left, artifacts.best_right, paths["beam_tracks"], dpi=resolved_dpi, ) paths["representative_lines"] = out_dir_path / "representative_lines.png" plot_representatives_and_lines( artifacts.left_reps, artifacts.right_reps, artifacts.left_line, artifacts.right_line, artifacts.left_reference, artifacts.right_reference, paths["representative_lines"], dpi=resolved_dpi, ) paths["deflection_curves"] = out_dir_path / "deflection_curves.png" plot_deflection_curves(artifacts.deflection_rows, paths["deflection_curves"], dpi=resolved_dpi) return paths # --------------------------------------------------------------------------- # PLY overlays # --------------------------------------------------------------------------- def export_overlay_plys(result: DeflectionResult, out_dir: str | Path) -> dict[str, str]: """Export coloured PLY overlays of the ROI, fitted lines, and measurements. Produces ``roi_cropped_cloud.ply`` and three measurement overlays (full / left-only / right-only). Returns a mapping of name to path. """ out_dir_path = ensure_output_dir(out_dir) artifacts = _require_artifacts(result) return export_measurement_overlay_plys( output_dir=out_dir_path, roi_points=artifacts.filtered_roi, left_track=artifacts.best_left, right_track=artifacts.best_right, left_line=artifacts.left_line, right_line=artifacts.right_line, left_reference=artifacts.left_reference, right_reference=artifacts.right_reference, left_measurement=artifacts.left_summary.get("selected_measurement"), right_measurement=artifacts.right_summary.get("selected_measurement"), config=artifacts.resolved_config, ) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _require_artifacts(result: DeflectionResult): if result._artifacts is None: raise ValueError( "DeflectionResult has no internal artifacts; it was likely " "reconstructed from a serialized report. Re-run compute_deflection " "to obtain a writable result." ) return result._artifacts def _annotate_beam_track_rows( rows: list[dict], left_tracks: list[TrackSolution], right_tracks: list[TrackSolution], ) -> list[dict]: rank_map: dict[str, tuple[str, int]] = {} for track in left_tracks + right_tracks: for candidate in track.candidates: key = f"{track.side}:{candidate.candidate_id}" if key not in rank_map: rank_map[key] = (track.side, track.rank) annotated = [] for row in rows: key = f"{row['side']}:{row['candidate_id']}" annotation = rank_map.get(key) annotated.append( row | { "selected_rank": annotation[1] if annotation else "", "selected_side": annotation[0] if annotation else "", } ) return annotated __all__ = [ "write_report_json", "write_curves_csv", "write_beam_tracks_csv", "render_plots", "export_overlay_plys", ]