返回 JoyAI-Echo
action_overlay.py
根目录 / echo_wm / helpers / action_overlay.py
1 """Optional Genie-style action HUD for generated WM videos."""
2
3 from __future__ import annotations
4
5 import os
6 import shutil
7 import subprocess
8 from dataclasses import dataclass
9 from pathlib import Path
10 from typing import Sequence
11
12 import cv2
13 import numpy as np
14 import torch
15 from PIL import Image, ImageDraw, ImageFilter, ImageFont
16 from scipy.spatial.transform import Rotation
17
18
19 def _ffmpeg_binary() -> str:
20 """Resolve an encoder-capable ffmpeg without relying on a shell PATH."""
21 configured = os.environ.get("FFMPEG_BIN")
22 if configured:
23 return configured
24 # The cluster's /usr/local/bin/ffmpeg is NVENC-only and lacks libx264;
25 # /usr/bin/ffmpeg is the broadly playable software-H.264 build.
26 if Path("/usr/bin/ffmpeg").is_file():
27 return "/usr/bin/ffmpeg"
28 return "ffmpeg"
29
30
31 def _pose_inverse(p: np.ndarray) -> np.ndarray:
32 R = p[:3, :3]
33 t = p[:3, 3]
34 inv = np.eye(4, dtype=p.dtype)
35 inv[:3, :3] = R.T
36 inv[:3, 3] = -R.T @ t
37 return inv
38
39
40 def _per_frame_deltas(c2w: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
41 """Return ``(N-1, 3)`` per-frame translation and YXZ-euler rotation (deg)."""
42 n = c2w.shape[0]
43 trans = np.zeros((n - 1, 3), dtype=np.float64)
44 rots = np.zeros((n - 1, 3), dtype=np.float64)
45 for i in range(n - 1):
46 rel = _pose_inverse(c2w[i]) @ c2w[i + 1]
47 trans[i] = rel[:3, 3]
48 rots[i] = Rotation.from_matrix(rel[:3, :3]).as_euler("YXZ", degrees=True)
49 return trans, rots
50
51
52 def _translation_keys(
53 trans: np.ndarray,
54 *,
55 floor_dx: float = 0.005,
56 floor_dz: float = 0.005,
57 frac_dx: float = 0.30,
58 frac_dz: float = 0.30,
59 ) -> list[list[str]]:
60 """Discretise per-frame translation into WASD key lists.
61
62 Thresholds are adaptive: ``thresh = max(floor, frac * p95(|delta|))``.
63 """
64 p95 = np.percentile(np.abs(trans), 95.0, axis=0) if trans.size else np.zeros(3)
65 thr_dx = max(floor_dx, frac_dx * p95[0])
66 thr_dz = max(floor_dz, frac_dz * p95[2])
67
68 per_frame: list[list[str]] = []
69 for dx, _dy, dz in trans:
70 keys: list[str] = []
71 if abs(dz) > thr_dz:
72 keys.append("W" if dz > 0 else "S")
73 if abs(dx) > thr_dx:
74 keys.append("D" if dx > 0 else "A")
75 per_frame.append(keys)
76 per_frame.append(list(per_frame[-1]) if per_frame else [])
77 return per_frame
78
79
80 def _normalised_rotation(
81 rots: np.ndarray, *, floor_deg: float = 0.5, ema_alpha: float = 0.35
82 ) -> tuple[np.ndarray, np.ndarray]:
83 """Per-frame ``(yaw, pitch)`` in ``[-1, 1]`` with EMA smoothing."""
84 p95 = np.percentile(np.abs(rots), 95.0, axis=0) if rots.size else np.zeros(3)
85 yaw_scale = max(floor_deg, p95[0])
86 pitch_scale = max(floor_deg, p95[1])
87 n = rots.shape[0]
88 yaw = np.zeros(n + 1, dtype=np.float64)
89 pitch = np.zeros(n + 1, dtype=np.float64)
90 yaw_ema = pitch_ema = 0.0
91 for i in range(n):
92 y = float(np.clip(rots[i, 0] / yaw_scale, -1.0, 1.0))
93 p = float(np.clip(rots[i, 1] / pitch_scale, -1.0, 1.0))
94 yaw_ema = ema_alpha * y + (1.0 - ema_alpha) * yaw_ema
95 pitch_ema = ema_alpha * p + (1.0 - ema_alpha) * pitch_ema
96 yaw[i] = yaw_ema
97 pitch[i] = pitch_ema
98 if n > 0:
99 yaw[-1] = yaw[-2]
100 pitch[-1] = pitch[-2]
101 return yaw, pitch
102
103
104 _FONT_CANDIDATES: tuple[str, ...] = (
105 "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
106 "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
107 "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
108 )
109
110
111 def _load_font(size: int) -> ImageFont.ImageFont:
112 for path in _FONT_CANDIDATES:
113 if os.path.exists(path):
114 try:
115 return ImageFont.truetype(path, size=size)
116 except Exception:
117 continue
118 return ImageFont.load_default()
119
120
121 @dataclass(frozen=True)
122 class _Layout:
123 width: int
124 height: int
125
126 @property
127 def key_size(self) -> int:
128 return max(32, int(self.height * 0.08))
129
130 @property
131 def key_gap(self) -> int:
132 return max(4, int(self.key_size * 0.15))
133
134 @property
135 def key_radius(self) -> int:
136 return max(4, int(self.key_size * 0.2))
137
138
139 class ActionOverlayRenderer:
140 """Renders the WASD-cluster + rotation-joystick overlay onto video frames."""
141
142 CORNER_CHOICES = ("bottom-left", "bottom-right", "top-left", "top-right")
143
144 def __init__(self, width: int, height: int):
145 self.layout = _Layout(int(width), int(height))
146 self.width, self.height = self.layout.width, self.layout.height
147 self.font = _load_font(int(self.layout.key_size * 0.5))
148 self._key_tiles = self._build_key_tiles()
149
150 def _build_key_tiles(self) -> dict[tuple[str, bool], Image.Image]:
151 sz, r = self.layout.key_size, self.layout.key_radius
152 tiles: dict[tuple[str, bool], Image.Image] = {}
153 for key in ("W", "A", "S", "D"):
154 for pressed in (False, True):
155 fill = (255, 255, 255, 200) if pressed else (0, 0, 0, 100)
156 outline = (255, 255, 255, 255) if pressed else (255, 255, 255, 60)
157 text_color = (0, 0, 0, 220) if pressed else (255, 255, 255, 180)
158 tile = Image.new("RGBA", (sz, sz), (0, 0, 0, 0))
159 d = ImageDraw.Draw(tile)
160 d.rounded_rectangle(
161 [0, 0, sz - 1, sz - 1],
162 radius=r,
163 fill=fill,
164 outline=outline,
165 width=max(1, int(sz * 0.03)),
166 )
167 bbox = d.textbbox((0, 0), key, font=self.font)
168 tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
169 d.text((sz / 2 - tw / 2, sz / 2 - th / 2 - 2), key, fill=text_color, font=self.font)
170 tiles[(key, pressed)] = tile
171 return tiles
172
173 def _draw_joystick(self, canvas: Image.Image, cx: int, cy: int, yaw: float, pitch: float) -> None:
174 yaw = float(np.clip(yaw, -1.0, 1.0))
175 pitch = float(np.clip(pitch, -1.0, 1.0))
176 radius = max(30, int((self.layout.key_size * 2 + self.layout.key_gap) * 0.47))
177
178 shadow = Image.new("RGBA", (self.width, self.height), (0, 0, 0, 0))
179 ImageDraw.Draw(shadow).ellipse(
180 [cx - radius - 14, cy - radius - 14, cx + radius + 14, cy + radius + 14],
181 fill=(0, 0, 0, 88),
182 )
183 canvas.alpha_composite(shadow.filter(ImageFilter.GaussianBlur(max(8, int(radius * 0.16)))))
184
185 d = ImageDraw.Draw(canvas)
186 d.ellipse(
187 [cx - radius, cy - radius, cx + radius, cy + radius],
188 fill=(7, 9, 13, 104),
189 outline=(255, 255, 255, 95),
190 width=max(1, int(radius * 0.035)),
191 )
192 d.ellipse(
193 [cx - radius - 7, cy - radius - 7, cx + radius + 7, cy + radius + 7],
194 outline=(255, 255, 255, 42),
195 width=max(1, int(radius * 0.025)),
196 )
197 d.line([cx - radius * 0.63, cy, cx + radius * 0.63, cy], fill=(255, 255, 255, 56), width=1)
198 d.line([cx, cy - radius * 0.63, cx, cy + radius * 0.63], fill=(255, 255, 255, 56), width=1)
199
200 marker_offset = int(radius * 0.78)
201 marker_size = max(7, int(radius * 0.16))
202 self._draw_arrow(d, cx + marker_offset, cy, "right", yaw > 0.08, marker_size)
203 self._draw_arrow(d, cx - marker_offset, cy, "left", yaw < -0.08, marker_size)
204 self._draw_arrow(d, cx, cy - marker_offset, "up", pitch > 0.08, marker_size)
205 self._draw_arrow(d, cx, cy + marker_offset, "down", pitch < -0.08, marker_size)
206
207 max_offset = radius * 0.48
208 kx = int(cx + yaw * max_offset)
209 ky = int(cy - pitch * max_offset)
210 glow = Image.new("RGBA", (self.width, self.height), (0, 0, 0, 0))
211 gd = ImageDraw.Draw(glow)
212 gd.line([cx, cy, kx, ky], fill=(255, 255, 255, 78), width=max(3, int(radius * 0.055)))
213 gd.ellipse(
214 [kx - radius * 0.22, ky - radius * 0.22, kx + radius * 0.22, ky + radius * 0.22], fill=(255, 255, 255, 110)
215 )
216 canvas.alpha_composite(glow.filter(ImageFilter.GaussianBlur(max(5, int(radius * 0.09)))))
217
218 d = ImageDraw.Draw(canvas)
219 d.line([cx, cy, kx, ky], fill=(255, 255, 255, 120), width=max(1, int(radius * 0.025)))
220 kr = max(7, int(radius * 0.13))
221 d.ellipse(
222 [kx - kr, ky - kr, kx + kr, ky + kr], fill=(255, 255, 255, 230), outline=(255, 255, 255, 255), width=1
223 )
224 ir = max(3, int(kr * 0.36))
225 d.ellipse([kx - ir, ky - ir, kx + ir, ky + ir], fill=(20, 24, 30, 170))
226
227 @staticmethod
228 def _draw_arrow(d: ImageDraw.ImageDraw, cx: int, cy: int, direction: str, active: bool, size: int) -> None:
229 if direction == "right":
230 pts = [(cx - size * 0.55, cy - size), (cx + size * 0.65, cy), (cx - size * 0.55, cy + size)]
231 elif direction == "left":
232 pts = [(cx + size * 0.55, cy - size), (cx - size * 0.65, cy), (cx + size * 0.55, cy + size)]
233 elif direction == "up":
234 pts = [(cx - size, cy + size * 0.55), (cx, cy - size * 0.65), (cx + size, cy + size * 0.55)]
235 else:
236 pts = [(cx - size, cy - size * 0.55), (cx, cy + size * 0.65), (cx + size, cy - size * 0.55)]
237 d.polygon(pts, fill=(255, 255, 255, 210) if active else (255, 255, 255, 72))
238
239 def render_panel(
240 self,
241 pressed_keys: Sequence[str],
242 yaw: float,
243 pitch: float,
244 corner: str = "bottom-left",
245 ) -> Image.Image:
246 """Return an RGBA overlay of size ``(width, height)``."""
247 canvas = Image.new("RGBA", (self.width, self.height), (0, 0, 0, 0))
248 sz, gap = self.layout.key_size, self.layout.key_gap
249 margin = int(self.height * 0.05)
250 cluster_w = sz * 3 + gap * 2
251 cluster_h = sz * 2 + gap
252 joy_radius = max(30, int(cluster_h * 0.47))
253 joy_gap = int(sz * 1.0)
254
255 if corner == "bottom-right":
256 sx = self.width - margin - cluster_w
257 sy = self.height - margin - cluster_h
258 jcx = sx - joy_gap - joy_radius
259 elif corner == "top-left":
260 sx, sy = margin, margin
261 jcx = sx + cluster_w + joy_gap + joy_radius
262 elif corner == "top-right":
263 sx = self.width - margin - cluster_w
264 sy = margin
265 jcx = sx - joy_gap - joy_radius
266 else: # bottom-left default
267 sx, sy = margin, self.height - margin - cluster_h
268 jcx = sx + cluster_w + joy_gap + joy_radius
269 jcy = sy + cluster_h // 2
270
271 shadow = Image.new("RGBA", (self.width, self.height), (0, 0, 0, 0))
272 ImageDraw.Draw(shadow).rounded_rectangle(
273 [sx - 8, sy - 8, sx + cluster_w + 8, sy + cluster_h + 8],
274 radius=max(10, int(self.layout.key_radius * 1.35)),
275 fill=(0, 0, 0, 74),
276 )
277 canvas.alpha_composite(shadow.filter(ImageFilter.GaussianBlur(max(7, int(sz * 0.18)))))
278
279 positions = {
280 "W": (sx + sz + gap, sy),
281 "A": (sx, sy + sz + gap),
282 "S": (sx + sz + gap, sy + sz + gap),
283 "D": (sx + (sz + gap) * 2, sy + sz + gap),
284 }
285 for key in ("W", "A", "S", "D"):
286 canvas.alpha_composite(self._key_tiles[(key, key in pressed_keys)], dest=positions[key])
287 self._draw_joystick(canvas, jcx, jcy, yaw, pitch)
288 return canvas
289
290
291 def apply_overlay(
292 video_hwc: np.ndarray,
293 c2w: np.ndarray,
294 *,
295 corner: str = "bottom-left",
296 ) -> np.ndarray:
297 """Composite the camera-pose-driven action overlay onto each frame.
298
299 Args:
300 video_hwc: ``(T, H, W, 3)`` uint8 video.
301 c2w: ``(T_pose, 4, 4)`` camera-to-world poses driving the overlay.
302 Truncated to ``video_hwc.shape[0]`` frames if longer.
303 corner: Panel placement.
304
305 Returns:
306 ``(T, H, W, 3)`` uint8 array with the overlay composited.
307 """
308 T, H, W = video_hwc.shape[:3]
309 n_poses = min(int(c2w.shape[0]), T)
310 poses = c2w[:n_poses].astype(np.float32)
311
312 trans, rots = _per_frame_deltas(poses)
313 keys = _translation_keys(trans)
314 yaw, pitch = _normalised_rotation(rots)
315
316 renderer = ActionOverlayRenderer(width=W, height=H)
317 out = np.empty_like(video_hwc)
318 for t in range(T):
319 i = min(t, len(keys) - 1)
320 panel = renderer.render_panel(
321 pressed_keys=keys[i],
322 yaw=float(yaw[i]),
323 pitch=float(pitch[i]),
324 corner=corner,
325 )
326 frame = Image.fromarray(video_hwc[t]).convert("RGBA")
327 frame.alpha_composite(panel)
328 out[t] = np.asarray(frame.convert("RGB"), dtype=np.uint8)
329 return out
330
331
332 def overlay_genie_on_video(
333 video_path: Path,
334 c2w: np.ndarray | torch.Tensor,
335 output_path: Path | None = None,
336 *,
337 corner: str = "bottom-left",
338 ) -> Path:
339 """File-level wrapper around ``apply_overlay`` (Genie-3 camera-pose overlay).
340
341 Reads an mp4, composites the WASD-cluster + rotation-joystick overlay
342 DERIVED from the camera trajectory ``c2w`` onto every frame, and writes the
343 result back (re-encoded with libx264, original audio preserved). Mirrors the
344 I/O contract of ``overlay_action_on_video`` so the trainer can swap between
345 the real-action overlay and this camera-driven one transparently.
346
347 Args:
348 video_path: existing mp4 (replaced in-place unless output_path given).
349 c2w: ``(T_pose, 4, 4)`` camera-to-world poses (e.g. action.pt c2ws_raw).
350 output_path: write here instead of replacing input. Defaults to in-place.
351 corner: overlay panel placement.
352
353 Returns:
354 Path of the written mp4 (equals input path when output_path is None).
355 """
356 video_path = Path(video_path)
357 if output_path is None:
358 output_path = video_path
359 output_path = Path(output_path)
360
361 c2w_np = c2w.detach().float().cpu().numpy() if isinstance(c2w, torch.Tensor) else np.asarray(c2w)
362
363 cap = cv2.VideoCapture(str(video_path))
364 if not cap.isOpened():
365 raise RuntimeError(f"cannot open {video_path}")
366 src_fps = cap.get(cv2.CAP_PROP_FPS) or 24.0
367 n_video = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
368 W = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
369 H = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
370
371 n_poses = min(int(c2w_np.shape[0]), n_video)
372 poses = c2w_np[:n_poses].astype(np.float32)
373 trans, rots = _per_frame_deltas(poses)
374 keys = _translation_keys(trans)
375 yaw, pitch = _normalised_rotation(rots)
376 renderer = ActionOverlayRenderer(width=W, height=H)
377
378 silent_tmp = output_path.with_suffix(".silent.tmp.mp4")
379 fourcc = cv2.VideoWriter_fourcc(*"mp4v")
380 writer = cv2.VideoWriter(str(silent_tmp), fourcc, src_fps, (W, H))
381 written = 0
382 while True:
383 ok, frame = cap.read()
384 if not ok:
385 break
386 i = min(written, len(keys) - 1)
387 panel = renderer.render_panel(
388 pressed_keys=keys[i],
389 yaw=float(yaw[i]),
390 pitch=float(pitch[i]),
391 corner=corner,
392 )
393 frame_rgba = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).convert("RGBA")
394 frame_rgba.alpha_composite(panel)
395 writer.write(cv2.cvtColor(np.asarray(frame_rgba.convert("RGB")), cv2.COLOR_RGB2BGR))
396 written += 1
397 cap.release()
398 writer.release()
399 if written == 0:
400 silent_tmp.unlink(missing_ok=True)
401 raise RuntimeError(f"no frames decoded from {video_path}")
402
403 # Re-encode as broadly playable H.264 + preserve original audio if any.
404 has_audio = subprocess.run(
405 [os.environ.get("FFPROBE_BIN", "/usr/bin/ffprobe"), "-v", "error", "-select_streams", "a:0",
406 "-show_entries", "stream=index", "-of", "csv=p=0", str(video_path)],
407 capture_output=True,
408 ).stdout.strip()
409
410 final_tmp = output_path.with_suffix(".final.tmp.mp4")
411 ffmpeg = _ffmpeg_binary()
412
413 # Prefer hardware H.264 on the cluster, then software H.264. MPEG-4 Part 2
414 # is only a last resort because VSCode/browser players often reject it.
415 encoder_candidates = [
416 ("h264_nvenc", ["-preset", "p5", "-cq", "19"]),
417 ("libx264", ["-preset", "fast", "-crf", "20"]),
418 ("mpeg4", []),
419 ]
420
421 def build_command(codec: str, codec_args: list[str]) -> list[str]:
422 command = [ffmpeg, "-y", "-loglevel", "error", "-i", str(silent_tmp)]
423 if has_audio:
424 command.extend(["-i", str(video_path), "-map", "0:v", "-map", "1:a"])
425 command.extend(["-c:v", codec, *codec_args, "-pix_fmt", "yuv420p"])
426 if has_audio:
427 command.extend([
428 "-c:a", "aac", "-b:a", "192k",
429 "-t", f"{written / src_fps:.9f}",
430 ])
431 command.append(str(final_tmp))
432 return command
433
434 cmd = build_command("mpeg4", [])
435 if has_audio:
436 pass
437 res = None
438 errors: list[str] = []
439 for codec, codec_args in encoder_candidates:
440 final_tmp.unlink(missing_ok=True)
441 cmd = build_command(codec, codec_args)
442 candidate = subprocess.run(cmd, capture_output=True)
443 if candidate.returncode == 0:
444 res = candidate
445 break
446 errors.append(f"{codec}: {candidate.stderr.decode(errors='replace')[-300:]}")
447 silent_tmp.unlink(missing_ok=True)
448 if res is None or res.returncode != 0:
449 final_tmp.unlink(missing_ok=True)
450 raise RuntimeError(
451 f"ffmpeg genie overlay failed for {video_path}:\n" + "\n".join(errors)
452 )
453
454 shutil.move(str(final_tmp), str(output_path))
455 return output_path
456
456 lines PYTHON