[Quy Anh] «Elliot» Nguyen
108 lines
3.0 KiB
Python
108 lines
3.0 KiB
Python
import sys
|
|
import numpy as np
|
|
import pandas as pd
|
|
import matplotlib.pyplot as plt
|
|
from matplotlib.animation import FuncAnimation
|
|
|
|
if len(sys.argv) < 2:
|
|
sys.exit(
|
|
"Use it like the following: python visualization.py positions.csv (out.mp4)"
|
|
)
|
|
|
|
df = pd.read_csv(sys.argv[1])
|
|
frames = sorted(df["time"].unique())
|
|
|
|
# Replace p.e with direction-of-next-hop: where the ped actually moves to next frame.
|
|
# Falls back to p.e for the last frame and for teleport jumps.
|
|
df = df.sort_values(["id", "time"]).reset_index(drop=True)
|
|
df["dx"] = df.groupby("id")["x"].shift(-1) - df["x"]
|
|
df["dy"] = df.groupby("id")["y"].shift(-1) - df["y"]
|
|
mag = np.sqrt(df["dx"] ** 2 + df["dy"] ** 2)
|
|
TELEPORT_THRESHOLD = (
|
|
5.0 # any per-step jump larger than this is the teleport, not motion
|
|
)
|
|
fallback = (mag > TELEPORT_THRESHOLD) | mag.isna() | (mag < 1e-9)
|
|
df.loc[fallback, "dx"] = df.loc[fallback, "p.e.x"]
|
|
df.loc[fallback, "dy"] = df.loc[fallback, "p.e.y"]
|
|
mag = np.sqrt(df["dx"] ** 2 + df["dy"] ** 2).replace(0, 1.0)
|
|
df["dx"] = df["dx"] / mag
|
|
df["dy"] = df["dy"] / mag
|
|
|
|
fig, ax = plt.subplots(figsize=(10, 2.2))
|
|
ax.set_xlim(0, 50)
|
|
ax.set_ylim(0, 10)
|
|
ax.set_aspect("equal")
|
|
ax.set_facecolor("white")
|
|
ax.set_xticks([])
|
|
ax.set_yticks([])
|
|
for s in ("top", "right", "left", "bottom"):
|
|
ax.spines[s].set_color("gray")
|
|
|
|
s_white = ax.scatter([], [], s=30, facecolors="gray", edgecolor="gray", zorder=2)
|
|
s_black = ax.scatter([], [], s=30, facecolors="black", edgecolor="black", zorder=2)
|
|
title = ax.set_title("")
|
|
|
|
# Keep a reference to quiver objects so we can remove them each frame
|
|
quivers = []
|
|
|
|
|
|
def update(t):
|
|
sub = df[df["time"] == t]
|
|
w = sub[sub["color"] == 0]
|
|
b = sub[sub["color"] == 1]
|
|
|
|
s_white.set_offsets(w[["x", "y"]].values)
|
|
s_black.set_offsets(b[["x", "y"]].values)
|
|
|
|
# Remove old quiver arrows
|
|
for q in quivers:
|
|
q.remove()
|
|
quivers.clear()
|
|
|
|
# Draw fresh arrows each frame
|
|
arrow_scale = 1.5 # Tune this to change arrow length visually
|
|
if len(w):
|
|
q_white = ax.quiver(
|
|
w["x"],
|
|
w["y"],
|
|
w["dx"],
|
|
w["dy"],
|
|
color="gray",
|
|
angles="xy",
|
|
scale_units="xy",
|
|
scale=1 / arrow_scale,
|
|
width=0.003,
|
|
headwidth=4,
|
|
headlength=5,
|
|
zorder=3,
|
|
)
|
|
quivers.append(q_white)
|
|
if len(b):
|
|
q_black = ax.quiver(
|
|
b["x"],
|
|
b["y"],
|
|
b["dx"],
|
|
b["dy"],
|
|
color="black",
|
|
angles="xy",
|
|
scale_units="xy",
|
|
scale=1 / arrow_scale,
|
|
width=0.003,
|
|
headwidth=4,
|
|
headlength=5,
|
|
zorder=3,
|
|
)
|
|
quivers.append(q_black)
|
|
|
|
title.set_text(f"t = {t}")
|
|
return s_white, s_black, title
|
|
|
|
|
|
anim = FuncAnimation(fig, update, frames=frames, interval=50, blit=False)
|
|
|
|
if len(sys.argv) >= 3:
|
|
anim.save(sys.argv[2], writer="ffmpeg", fps=25)
|
|
print(f"Saved {sys.argv[2]}")
|
|
else:
|
|
plt.show()
|