File size: 5,661 Bytes
4ee0c8c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """Batch-visualize image + graph pairs for an episode.
Usage:
python visualize_graphs_batch.py \
--episode session_0408_162129/episode_00 \
--graphs session_0408_162129/episode_00/graph_per_side_image \
--output session_0408_162129/episode_00/graph_viz \
[--view side] [--stride 1] [--limit N]
"""
import argparse
import json
import re
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import networkx as nx
from PIL import Image
def load_side_graph_positions(episode: Path) -> dict:
side_graph_path = episode / "annotations" / "side_graph.json"
if not side_graph_path.exists():
side_graph_path = episode / "annotations" / "graph.json"
if not side_graph_path.exists():
return {}
with open(side_graph_path) as f:
data = json.load(f)
if "node_positions" in data:
return data["node_positions"]
if "side_graph" in data and isinstance(data["side_graph"], dict):
return data["side_graph"].get("node_positions", {})
return {}
def visualize_pair(frame_idx: int, rgb_path: Path, graph_dict: dict,
stored_pos: dict, out_path: Path):
G = nx.DiGraph()
for node in graph_dict["nodes"]:
G.add_node(node["id"], **node)
for edge in graph_dict["edges"]:
G.add_edge(edge["src"], edge["dst"], is_locked=edge["is_locked"])
pos = {}
for nid in G.nodes:
if nid in stored_pos:
x, y = stored_pos[nid]
pos[nid] = (x, -y)
elif nid == "robot":
pos[nid] = (450, 0)
missing = [n for n in G.nodes if n not in pos]
if missing:
if pos:
sub = nx.spring_layout(G, pos=pos, fixed=list(pos.keys()), seed=42)
else:
sub = nx.spring_layout(G, seed=42)
for n in missing:
pos[n] = sub[n]
fig, axes = plt.subplots(1, 2, figsize=(20, 8))
if rgb_path.exists():
axes[0].imshow(np.array(Image.open(rgb_path)))
else:
axes[0].text(0.5, 0.5, "RGB not found", ha="center", va="center")
axes[0].set_title(f"Frame {frame_idx} — RGB")
axes[0].axis("off")
ax = axes[1]
locked = [(e["src"], e["dst"]) for e in graph_dict["edges"] if e["is_locked"]]
unlocked = [(e["src"], e["dst"]) for e in graph_dict["edges"] if not e["is_locked"]]
node_colors, node_sizes = [], []
for nid in G.nodes:
nd = G.nodes[nid]
node_colors.append(nd.get("color", "#888"))
t = nd.get("type", "")
node_sizes.append(800 if t == "robot" else 1200 if t == "motherboard" else 600)
nx.draw_networkx_nodes(G, pos, ax=ax, node_color=node_colors,
node_size=node_sizes, edgecolors="black", linewidths=1.5)
if locked:
nx.draw_networkx_edges(G, pos, edgelist=locked, ax=ax, edge_color="#E74C3C",
width=2.0, alpha=0.8, arrows=True, arrowsize=15,
arrowstyle="-|>", connectionstyle="arc3,rad=0.1")
if unlocked:
nx.draw_networkx_edges(G, pos, edgelist=unlocked, ax=ax, edge_color="#2ECC71",
width=2.0, alpha=0.8, style="dashed", arrows=True,
arrowsize=15, arrowstyle="-|>",
connectionstyle="arc3,rad=0.1")
labels = {nid: f"{nid}{'' if G.nodes[nid].get('visible', True) else ' (hidden)'}"
for nid in G.nodes}
nx.draw_networkx_labels(G, pos, labels, ax=ax, font_size=7, font_weight="bold")
legend = [mpatches.Patch(color="#E74C3C", label="locked"),
mpatches.Patch(color="#2ECC71", label="unlocked")]
seen = {}
for n in graph_dict["nodes"]:
seen.setdefault(n.get("type", ""), n.get("color", "#888"))
for t, c in seen.items():
legend.append(mpatches.Patch(facecolor=c, edgecolor="black", label=t))
ax.legend(handles=legend, loc="upper left", fontsize=8, framealpha=0.9)
ax.set_title(f"Frame {frame_idx} — Graph ({len(locked)} locked, {len(unlocked)} unlocked)")
ax.axis("off")
plt.tight_layout()
plt.savefig(out_path, dpi=120, bbox_inches="tight")
plt.close(fig)
def main():
p = argparse.ArgumentParser()
p.add_argument("--episode", required=True)
p.add_argument("--graphs", required=True, help="Directory of frame_NNNNNN_graph.json")
p.add_argument("--output", required=True, help="Directory for visualization PNGs")
p.add_argument("--view", default="side")
p.add_argument("--stride", type=int, default=1)
p.add_argument("--limit", type=int, default=None)
args = p.parse_args()
episode = Path(args.episode)
graphs_dir = Path(args.graphs)
out_dir = Path(args.output)
out_dir.mkdir(parents=True, exist_ok=True)
rgb_dir = episode / args.view / "rgb"
stored_pos = load_side_graph_positions(episode)
graph_files = sorted(graphs_dir.glob("frame_*_graph.json"))
pairs = []
for gf in graph_files:
m = re.search(r"\d+", gf.stem)
if not m:
continue
idx = int(m.group())
pairs.append((idx, gf))
pairs = pairs[::args.stride]
if args.limit:
pairs = pairs[:args.limit]
for idx, gf in pairs:
with open(gf) as f:
graph_dict = json.load(f)
rgb_path = rgb_dir / f"frame_{idx:06d}.png"
out_path = out_dir / f"frame_{idx:06d}_viz.png"
visualize_pair(idx, rgb_path, graph_dict, stored_pos, out_path)
print(f"Saved {len(pairs)} visualizations to {out_dir}")
if __name__ == "__main__":
main()
|