from __future__ import annotations from pathlib import Path import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np from .beams import candidate_rows from .fitting import evaluate_line from .reference import evaluate_reference_line from .structures import FittedLine, FrameHypothesis, RepresentativePoint, SliceResult, TrackSolution def _sample_points(points: np.ndarray, limit: int) -> np.ndarray: if len(points) <= limit: return points indices = np.linspace(0, len(points) - 1, limit, dtype=int) return points[indices] def plot_frame_hypotheses( hypotheses: list[FrameHypothesis], output_path: str | Path, dpi: int, ) -> None: fig, ax = plt.subplots(figsize=(10, 4.5)) names = [item.name for item in hypotheses] scores = [item.score for item in hypotheses] ax.bar(names, scores, color="#2f5c7e") ax.set_title("Frame Hypothesis Scores") ax.set_ylabel("Score") ax.tick_params(axis="x", rotation=30) fig.tight_layout() fig.savefig(output_path, dpi=dpi) plt.close(fig) def plot_roi_overview( work_points_vehicle: np.ndarray, roi_points_vehicle: np.ndarray, output_path: str | Path, dpi: int, limit: int, ) -> None: full_sample = _sample_points(work_points_vehicle, limit) roi_sample = _sample_points(roi_points_vehicle, limit) fig, axes = plt.subplots(1, 2, figsize=(12, 5)) axes[0].scatter(full_sample[:, 0], full_sample[:, 2], s=1, c="#d3d3d3", alpha=0.6) axes[0].scatter(roi_sample[:, 0], roi_sample[:, 2], s=1.5, c="#1f77b4", alpha=0.7) axes[0].set_title("ROI in X-Z") axes[0].set_xlabel("x (mm)") axes[0].set_ylabel("z (mm)") axes[1].scatter(full_sample[:, 1], full_sample[:, 2], s=1, c="#d3d3d3", alpha=0.6) axes[1].scatter(roi_sample[:, 1], roi_sample[:, 2], s=1.5, c="#1f77b4", alpha=0.7) axes[1].set_title("ROI in Y-Z") axes[1].set_xlabel("y (mm)") axes[1].set_ylabel("z (mm)") fig.tight_layout() fig.savefig(output_path, dpi=dpi) plt.close(fig) def plot_beam_candidates( slices: list[SliceResult], left_track: TrackSolution | None, right_track: TrackSolution | None, output_path: str | Path, dpi: int, ) -> None: rows = candidate_rows(slices) fig, axes = plt.subplots(1, 2, figsize=(13, 5)) if rows: x = np.asarray([row["x_center"] for row in rows]) y = np.asarray([row["y_center"] for row in rows]) z = np.asarray([row["z_low"] for row in rows]) colors = ["#4a90e2" if row["side"] == "left" else "#ff7f50" for row in rows] axes[0].scatter(x, y, c=colors, s=10, alpha=0.45) axes[1].scatter(x, z, c=colors, s=10, alpha=0.45) for track, color in ((left_track, "#0b5394"), (right_track, "#b45f06")): if not track: continue track_x = [candidate.x_center for candidate in track.candidates] track_y = [candidate.y_center for candidate in track.candidates] track_z = [candidate.z_low for candidate in track.candidates] axes[0].plot(track_x, track_y, color=color, linewidth=2.5) axes[1].plot(track_x, track_z, color=color, linewidth=2.5) axes[0].set_title("Beam Candidate Bands and Tracks (X-Y)") axes[0].set_xlabel("x (mm)") axes[0].set_ylabel("y (mm)") axes[1].set_title("Beam Candidate Bands and Tracks (X-Z)") axes[1].set_xlabel("x (mm)") axes[1].set_ylabel("z (mm)") fig.tight_layout() fig.savefig(output_path, dpi=dpi) plt.close(fig) def plot_representatives_and_lines( left_points: list[RepresentativePoint], right_points: list[RepresentativePoint], left_line: FittedLine | None, right_line: FittedLine | None, left_reference, right_reference, output_path: str | Path, dpi: int, ) -> None: fig, axes = plt.subplots(1, 2, figsize=(13, 5)) for reps, color, label in ( (left_points, "#0b5394", "Left beam"), (right_points, "#b45f06", "Right beam"), ): if not reps: continue x = np.asarray([item.x for item in reps]) y = np.asarray([item.y for item in reps]) z = np.asarray([item.z for item in reps]) axes[0].scatter(x, y, s=12, color=color, alpha=0.7, label=f"{label} reps") axes[1].scatter(x, z, s=12, color=color, alpha=0.7, label=f"{label} reps") for line, reference, color, label in ( (left_line, left_reference, "#0b5394", "Left beam"), (right_line, right_reference, "#b45f06", "Right beam"), ): if line is None: continue x_grid = np.linspace(float(line.x_support.min()), float(line.x_support.max()), 300) y, z, _ = evaluate_line(line, x_grid) axes[0].plot(x_grid, y, color=color, linewidth=2.2, label=f"{label} fit") axes[1].plot(x_grid, z, color=color, linewidth=2.2, label=f"{label} fit") ref_eval, _ = evaluate_reference_line(reference, x_grid) if np.ndim(ref_eval) == 2 and not np.isnan(ref_eval[:, 0]).all(): axes[0].plot(x_grid, ref_eval[:, 0], color=color, linestyle="--", linewidth=1.4) axes[1].plot(x_grid, ref_eval[:, 1], color=color, linestyle="--", linewidth=1.4) axes[0].set_title("Representative Points and Lines (X-Y)") axes[0].set_xlabel("x (mm)") axes[0].set_ylabel("y (mm)") axes[1].set_title("Representative Points and Lines (X-Z)") axes[1].set_xlabel("x (mm)") axes[1].set_ylabel("z (mm)") handles, labels = axes[0].get_legend_handles_labels() if handles: fig.legend(handles, labels, loc="upper center", ncol=4) fig.tight_layout() fig.savefig(output_path, dpi=dpi) plt.close(fig) def plot_deflection_curves( deflection_rows: list[dict], output_path: str | Path, dpi: int, ) -> None: fig, ax = plt.subplots(figsize=(11, 4.5)) if deflection_rows: rows = deflection_rows x = np.asarray([row["x"] for row in rows], dtype=np.float64) left = np.asarray([row["left_deflection"] for row in rows], dtype=np.float64) right = np.asarray([row["right_deflection"] for row in rows], dtype=np.float64) if not np.isnan(left).all(): ax.plot(x, left, color="#0b5394", linewidth=2.0, label="Left deflection") if not np.isnan(right).all(): ax.plot(x, right, color="#b45f06", linewidth=2.0, label="Right deflection") ax.axhline(0.0, color="#666666", linewidth=1.0, linestyle=":") ax.set_title("Deflection Curves") ax.set_xlabel("x (mm)") ax.set_ylabel("Deflection Δz (mm)") handles, labels = ax.get_legend_handles_labels() if handles: ax.legend(loc="best") fig.tight_layout() fig.savefig(output_path, dpi=dpi) plt.close(fig)