Files
malkuthro-ComfyUI-Koolook/nuke_CAM_exporter/testing_tools/plot_top_view.py
T
2025-12-03 23:18:50 +01:00

382 lines
13 KiB
Python

#!/usr/bin/env python3
"""
Quick top-view visualizer for camera paths.
Example:
py -3 nuke_CAM_exporter/testing_tools/plot_top_view.py ^
--config nuke_CAM_exporter/configs/camera_config_v01.json ^
--output nuke_CAM_exporter/testing_tools/top_view.png
"""
from __future__ import annotations
import argparse
import json
import math
from pathlib import Path
from typing import Iterable, List, Sequence, Tuple
import matplotlib.pyplot as plt
from matplotlib import patches
NUKE_SWAP_AXES = True
NUKE_INVERT_HORIZONTAL = True
NUKE_INVERT_VERTICAL = True
NUKE_HORIZONTAL_LABEL = "+Z (frame left) <- -> -Z (frame right)"
NUKE_VERTICAL_LABEL = "+X (frame down) v ^ -X (frame up)"
def euler_zxy_to_matrix(rx_rad: float, ry_rad: float, rz_rad: float) -> List[List[float]]:
sx, cx = math.sin(rx_rad), math.cos(rx_rad)
sy, cy = math.sin(ry_rad), math.cos(ry_rad)
sz, cz = math.sin(rz_rad), math.cos(rz_rad)
Rx = (
(1.0, 0.0, 0.0),
(0.0, cx, -sx),
(0.0, sx, cx),
)
Ry = (
(cy, 0.0, sy),
(0.0, 1.0, 0.0),
(-sy, 0.0, cy),
)
Rz = (
(cz, -sz, 0.0),
(sz, cz, 0.0),
(0.0, 0.0, 1.0),
)
def mat_mul(a, b):
return [
[
a[r][0] * b[0][c] + a[r][1] * b[1][c] + a[r][2] * b[2][c]
for c in range(3)
]
for r in range(3)
]
return mat_mul(Ry, mat_mul(Rx, Rz))
def load_track_samples(
translation_path: Path,
rotation_path: Path | None,
) -> List[dict]:
samples: List[dict] = []
def append_sample(tx: float, ty: float, tz: float, rx_deg: float | None, ry_deg: float | None, rz_deg: float | None):
forward = None
if rx_deg is not None and ry_deg is not None and rz_deg is not None:
R = euler_zxy_to_matrix(
math.radians(rx_deg),
math.radians(ry_deg),
math.radians(rz_deg),
)
# Camera in Nuke looks toward -Z (viewing down the negative Z axis)
forward = (-R[0][2], -R[1][2], -R[2][2])
samples.append({"translation": (tx, ty, tz), "forward": forward})
lines = [l.strip() for l in translation_path.read_text().splitlines() if l.strip()]
if rotation_path:
rot_lines = [l.strip() for l in rotation_path.read_text().splitlines() if l.strip()]
if len(lines) != len(rot_lines):
raise ValueError("Translation and rotation ASCI files must have the same number of lines.")
for t_line, r_line in zip(lines, rot_lines):
tx, ty, tz = map(float, t_line.split()[:3])
rx_deg, ry_deg, rz_deg = map(float, r_line.split()[:3])
append_sample(tx, ty, tz, rx_deg, ry_deg, rz_deg)
else:
for line in lines:
parts = line.split()
if len(parts) < 6:
raise ValueError(
"Combined ASCI expected 6 columns per line (Tx Rx Ty Ry Tz Rz). "
f"Got {len(parts)} columns: {line}"
)
tx = float(parts[0])
rx_deg = float(parts[1])
ty = float(parts[2])
ry_deg = float(parts[3])
tz = float(parts[4])
rz_deg = float(parts[5])
append_sample(tx, ty, tz, rx_deg, ry_deg, rz_deg)
return samples
def parse_pose_file(path: Path, limit: int | None) -> List[dict]:
rows: List[dict] = []
lines = [l.strip() for l in path.read_text().splitlines() if l.strip()]
for line in lines[1:]: # skip header
values = list(map(float, line.split()))
if len(values) != 19:
raise ValueError(f"Expected 19 columns per pose row, got {len(values)}")
rotation = [
[values[7], values[8], values[9]],
[values[11], values[12], values[13]],
[values[15], values[16], values[17]],
]
translation = (values[10], values[14], values[18])
rows.append({"R": rotation, "t": translation})
if limit is not None and len(rows) >= limit:
break
return rows
def scale_rows(rows: Iterable[Tuple[float, float, float]], scale: float) -> List[Tuple[float, float, float]]:
if scale == 1.0:
return list(rows)
return [(x * scale, y * scale, z * scale) for x, y, z in rows]
def split_axes(rows: Sequence[Tuple[float, float, float]]) -> Tuple[List[float], List[float], List[float]]:
xs, ys, zs = [], [], []
for x, y, z in rows:
xs.append(x)
ys.append(y)
zs.append(z)
return xs, ys, zs
def load_config(path: Path | None) -> dict:
if not path:
return {}
data = json.loads(path.read_text())
data["_config_dir"] = path.parent
return data
def resolve_path(config: dict, key: str) -> Path | None:
value = config.get(key)
if value is None:
return None
candidate = Path(value)
if not candidate.is_absolute() and "_config_dir" in config:
candidate = (config["_config_dir"] / candidate).resolve()
return candidate
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Top-view plot of camera path (Nuke-style top view).")
parser.add_argument("--config", type=Path, help="Optional JSON config (same format as converter).")
parser.add_argument("--translation", type=Path, help="Translation ASCI file (Tx Ty Tz).")
parser.add_argument("--rotation", type=Path, help="Rotation ASCI file (Rx Ry Rz) when exports are split.")
parser.add_argument(
"--pose",
type=Path,
help="Converted pose TXT (to overlay world-to-camera translations). Optional.",
)
parser.add_argument("--unit-scale", type=float, default=None, help="Scale to apply to translation vectors.")
parser.add_argument("--limit", type=int, help="Maximum number of frames to plot.")
parser.add_argument("--output", type=Path, help="Optional PNG to save instead of showing interactively.")
parser.add_argument("--show", action="store_true", help="Force interactive window even when saving.")
square = parser.add_argument_group("workspace overlays")
square.add_argument(
"--square-size",
type=float,
help="Explicit side length for centered square (meters). Overrides auto sizing.",
)
square.add_argument(
"--square-padding",
type=float,
default=0.15,
help="Extra fractional padding when auto-sizing the square (default 0.15).",
)
square.add_argument(
"--square-color",
default="#00b894",
help="Line color for the square/highlight overlays.",
)
square.add_argument(
"--highlight-quadrant",
choices=["++", "+-", "-+", "--"],
default="+-",
help="Quadrant to shade (horizontal sign first, default '+-').",
)
square.add_argument(
"--highlight-alpha",
type=float,
default=0.15,
help="Alpha used for quadrant shading.",
)
return parser
def main() -> None:
parser = build_parser()
args = parser.parse_args()
config = load_config(args.config)
def resolved_path(cli_value: Path | None, config_key: str) -> Path | None:
if cli_value:
return cli_value.resolve()
return resolve_path(config, config_key)
translation_path = resolved_path(args.translation, "input_translation")
rotation_path = resolved_path(args.rotation, "input_rotation")
if translation_path is None:
translation_path = resolved_path(None, "input")
rotation_path = None
if translation_path is None:
parser.error("Translation source missing (pass --translation or provide input/_translation in config).")
pose_path = resolved_path(args.pose, "output")
unit_scale = args.unit_scale if args.unit_scale is not None else config.get("unit_scale", 1.0)
frame_cfg = config.get("frame_range", {})
start_idx = frame_cfg.get("start", 0)
limit = args.limit
if limit is None:
limit = frame_cfg.get("max_frames")
if limit is None and frame_cfg.get("end") is not None:
limit = max(0, frame_cfg["end"] - start_idx)
track_samples = load_track_samples(translation_path, rotation_path)
track_samples = track_samples[start_idx:]
if limit is not None:
track_samples = track_samples[:limit]
for sample in track_samples:
tx, ty, tz = sample["translation"]
sample["translation"] = (tx * unit_scale, ty * unit_scale, tz * unit_scale)
translations_scaled = [sample["translation"] for sample in track_samples]
tx, _, tz = split_axes(translations_scaled)
def orient_axes(xs: Sequence[float], zs: Sequence[float]) -> Tuple[List[float], List[float]]:
horiz, vert = (list(zs), list(xs)) if NUKE_SWAP_AXES else (list(xs), list(zs))
if NUKE_INVERT_HORIZONTAL:
horiz = [-val for val in horiz]
if NUKE_INVERT_VERTICAL:
vert = [-val for val in vert]
return horiz, vert
def orient_point(x: float, z: float) -> Tuple[float, float]:
hx, hy = orient_axes([x], [z])
return hx[0], hy[0]
hx, hy = orient_axes(tx, tz)
track_screen_points = list(zip(hx, hy))
fig, ax = plt.subplots(figsize=(8, 6))
ax.axhline(0.0, color="#666666", linewidth=0.8, alpha=0.6)
ax.axvline(0.0, color="#666666", linewidth=0.8, alpha=0.6)
ax.scatter([0], [0], c="#222222", s=20, marker="x", label="Origin", zorder=5)
ax.scatter(hx, hy, s=12, c="tab:green", label="ASCI center (scaled)")
ax.plot(hx, hy, color="tab:green", alpha=0.3)
pose_rows: List[dict] | None = None
pose_screen_points: List[Tuple[float, float]] | None = None
if pose_path:
pose_rows = parse_pose_file(pose_path, None)
pose_rows = pose_rows[start_idx:]
if limit is not None:
pose_rows = pose_rows[:limit]
pose_translations = [row["t"] for row in pose_rows]
px, _, pz = split_axes(pose_translations)
phx, phy = orient_axes(px, pz)
pose_screen_points = list(zip(phx, phy))
ax.scatter(phx, phy, s=12, c="tab:purple", label="Pose t (world-to-camera)")
ax.plot(phx, phy, color="tab:purple", alpha=0.3)
ax.set_xlabel(NUKE_HORIZONTAL_LABEL)
ax.set_ylabel(NUKE_VERTICAL_LABEL)
if hx and hy:
if args.square_size is not None:
half = args.square_size / 2.0
else:
max_extent = max(max(abs(val) for val in hx), max(abs(val) for val in hy), 1e-6)
half = max_extent * (1.0 + args.square_padding)
ax.set_xlim(-half, half)
ax.set_ylim(-half, half)
square = patches.Rectangle(
(-half, -half),
half * 2,
half * 2,
linewidth=1.2,
edgecolor=args.square_color,
facecolor="none",
linestyle="-",
zorder=2,
)
ax.add_patch(square)
if args.highlight_quadrant:
x0 = 0.0 if args.highlight_quadrant[0] == "+" else -half
y0 = 0.0 if args.highlight_quadrant[1] == "+" else -half
width = half
height = half
highlight = patches.Rectangle(
(x0, y0),
width,
height,
facecolor=args.square_color,
alpha=args.highlight_alpha,
edgecolor="none",
zorder=1,
)
ax.add_patch(highlight)
pad = half * 0.04
label_kwargs = dict(color="#414141", fontsize=9, fontweight="semibold")
ax.text(-half + pad, pad, "+Z", ha="left", va="bottom", **label_kwargs)
ax.text(half - pad, pad, "-Z", ha="right", va="bottom", **label_kwargs)
ax.text(-pad, -half + pad, "+X", ha="right", va="top", **label_kwargs)
ax.text(-pad, half - pad, "-X", ha="right", va="bottom", **label_kwargs)
nuke_forward_samples = [s for s in track_samples if s["forward"] is not None]
if nuke_forward_samples:
arrow_scale = max(half * 0.2, 0.05)
head_width = half * 0.05
head_length = half * 0.08
for idx, label in ((0, "Start"), (-1, "End")):
base_x, base_y = track_screen_points[idx]
forward = track_samples[idx]["forward"]
dir_x, dir_y = orient_point(forward[0], forward[2])
ax.arrow(
base_x,
base_y,
dir_x * arrow_scale,
dir_y * arrow_scale,
head_width=head_width,
head_length=head_length,
length_includes_head=True,
color="tab:blue",
linewidth=1.5,
alpha=0.9,
zorder=4,
)
ax.text(
base_x,
base_y,
label,
color="tab:blue",
fontsize=8,
fontweight="bold",
ha="center",
va="center",
zorder=5,
)
ax.scatter([], [], marker="^", color="tab:blue", label="Nuke cam forward")
ax.set_title("Camera path (top view)")
ax.axis("equal")
ax.grid(True, linestyle="--", alpha=0.3)
ax.legend()
if args.output:
args.output.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(args.output, dpi=200, bbox_inches="tight")
print(f"Saved plot to {args.output}")
if not args.show:
plt.close(fig)
return
plt.show()
if __name__ == "__main__":
main()