181 lines
6.6 KiB
Python
181 lines
6.6 KiB
Python
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)
|