view_camera_log.py

"""Utility to inspect camera logs saved by run_bullet_simulation."""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

try:  # pragma: no cover
    import cv2
except ImportError:  # pragma: no cover
    cv2 = None  # type: ignore[assignment]


def _cv_to_rgb(image: np.ndarray) -> np.ndarray:
    if cv2 is None:
        return image
    if image.ndim == 2 or image.shape[2] != 3:
        return image
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Visualize frames stored in a camera log (.npz) generated by run_bullet_simulation."
    )
    parser.add_argument("log_path", type=Path, help="Path to the camera log .npz file.")
    parser.add_argument(
        "--stride",
        type=int,
        default=1,
        help="Display one every N frames (default: 1).",
    )
    parser.add_argument(
        "--pause",
        type=float,
        default=0.05,
        help="Delay between frames in seconds when playing as a loop.",
    )
    parser.add_argument(
        "--headless",
        action="store_true",
        help="Print basic stats without opening a Matplotlib window.",
    )
    args = parser.parse_args()

    if not args.log_path.exists():
        parser.error(f"{args.log_path} does not exist")

    log = np.load(args.log_path)
    images = log["image"]
    times = log["time"]
    positions = log["position"]
    rotations = log["rotation"]
    stride = max(1, args.stride)
    print(f"Loaded {len(times)} frames from {args.log_path}")
    print(f"Positions range: x {positions[:,0].min():.2f}..{positions[:,0].max():.2f}, "
          f"y {positions[:,1].min():.2f}..{positions[:,1].max():.2f}, "
          f"z {positions[:,2].min():.2f}..{positions[:,2].max():.2f}")
    roll = np.arctan2(rotations[:, 2, 1], rotations[:, 2, 2])
    pitch = np.arctan2(-rotations[:, 2, 0], np.sqrt(rotations[:, 2, 1] ** 2 + rotations[:, 2, 2] ** 2))
    yaw = np.arctan2(rotations[:, 1, 0], rotations[:, 0, 0])
    print(
        "Euler (deg) ranges:"
        f" roll {np.rad2deg(roll).min():.1f}..{np.rad2deg(roll).max():.1f},"
        f" pitch {np.rad2deg(pitch).min():.1f}..{np.rad2deg(pitch).max():.1f},"
        f" yaw {np.rad2deg(yaw).min():.1f}..{np.rad2deg(yaw).max():.1f}"
    )
    if args.headless:
        return

    fig, ax = plt.subplots(figsize=(4, 3))
    im = ax.imshow(images[0])
    ax.set_axis_off()
    ax.set_title(f"t={times[0]:.2f} s")
    plt.show(block=False)

    try:
        for idx in range(0, len(images), stride):
            im.set_data(images[idx])
            xyz = positions[idx]
            rpy = np.rad2deg([roll[idx], pitch[idx], yaw[idx]])
            ax.set_title(
                f"t={times[idx]:.2f}s | pos=({xyz[0]:.2f},{xyz[1]:.2f},{xyz[2]:.2f}) | "
                f"rpy=({rpy[0]:.1f},{rpy[1]:.1f},{rpy[2]:.1f})°"
            )
            fig.canvas.draw_idle()
            fig.canvas.flush_events()
            plt.pause(args.pause)
    except KeyboardInterrupt:  # pragma: no cover
        pass
    finally:
        plt.show()


if __name__ == "__main__":
    main()