massive-cloud/alg/delivery_v1.2.1/underframe_deflection/visualization.py

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)