wehub-resource-sync
356 lines
14 KiB
Python
356 lines
14 KiB
Python
# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Sana-WM action overlay (Genie-3 style) for camera-controlled videos.
|
|
|
|
Renders a small WASD key cluster + rotation joystick on top of each frame,
|
|
driven by the per-frame relative pose extracted from the camera trajectory.
|
|
|
|
Conventions (OpenCV camera frame):
|
|
+X right, +Y down, +Z forward.
|
|
W/S → translate along +Z / -Z (forward / back).
|
|
D/A → translate along +X / -X (right / left).
|
|
yaw > 0 (pan right) → joystick deflects right.
|
|
pitch > 0 (tilt down) → joystick deflects down.
|
|
|
|
The renderer pre-builds static layers (key tiles, joystick base) once;
|
|
per-frame work is just compositing them onto the input frame.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from dataclasses import dataclass
|
|
from typing import Sequence
|
|
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw, ImageFilter, ImageFont
|
|
from scipy.spatial.transform import Rotation
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pose → per-frame keys + normalized rotation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _pose_inverse(p: np.ndarray) -> np.ndarray:
|
|
R = p[:3, :3]
|
|
t = p[:3, 3]
|
|
inv = np.eye(4, dtype=p.dtype)
|
|
inv[:3, :3] = R.T
|
|
inv[:3, 3] = -R.T @ t
|
|
return inv
|
|
|
|
|
|
def _per_frame_deltas(c2w: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Return ``(N-1, 3)`` per-frame translation and YXZ-euler rotation (deg)."""
|
|
n = c2w.shape[0]
|
|
trans = np.zeros((n - 1, 3), dtype=np.float64)
|
|
rots = np.zeros((n - 1, 3), dtype=np.float64)
|
|
for i in range(n - 1):
|
|
rel = _pose_inverse(c2w[i]) @ c2w[i + 1]
|
|
trans[i] = rel[:3, 3]
|
|
rots[i] = Rotation.from_matrix(rel[:3, :3]).as_euler("YXZ", degrees=True)
|
|
return trans, rots
|
|
|
|
|
|
def _translation_keys(
|
|
trans: np.ndarray,
|
|
*,
|
|
floor_dx: float = 0.005,
|
|
floor_dz: float = 0.005,
|
|
frac_dx: float = 0.30,
|
|
frac_dz: float = 0.30,
|
|
) -> list[list[str]]:
|
|
"""Discretise per-frame translation into WASD key lists.
|
|
|
|
Thresholds are adaptive: ``thresh = max(floor, frac * p95(|delta|))``.
|
|
"""
|
|
p95 = np.percentile(np.abs(trans), 95.0, axis=0) if trans.size else np.zeros(3)
|
|
thr_dx = max(floor_dx, frac_dx * p95[0])
|
|
thr_dz = max(floor_dz, frac_dz * p95[2])
|
|
|
|
per_frame: list[list[str]] = []
|
|
for dx, _dy, dz in trans:
|
|
keys: list[str] = []
|
|
if abs(dz) > thr_dz:
|
|
keys.append("W" if dz > 0 else "S")
|
|
if abs(dx) > thr_dx:
|
|
keys.append("D" if dx > 0 else "A")
|
|
per_frame.append(keys)
|
|
per_frame.append(list(per_frame[-1]) if per_frame else [])
|
|
return per_frame
|
|
|
|
|
|
def _normalised_rotation(
|
|
rots: np.ndarray, *, floor_deg: float = 0.5, ema_alpha: float = 0.35
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Per-frame ``(yaw, pitch)`` in ``[-1, 1]`` with EMA smoothing."""
|
|
p95 = np.percentile(np.abs(rots), 95.0, axis=0) if rots.size else np.zeros(3)
|
|
yaw_scale = max(floor_deg, p95[0])
|
|
pitch_scale = max(floor_deg, p95[1])
|
|
n = rots.shape[0]
|
|
yaw = np.zeros(n + 1, dtype=np.float64)
|
|
pitch = np.zeros(n + 1, dtype=np.float64)
|
|
yaw_ema = pitch_ema = 0.0
|
|
for i in range(n):
|
|
y = float(np.clip(rots[i, 0] / yaw_scale, -1.0, 1.0))
|
|
p = float(np.clip(rots[i, 1] / pitch_scale, -1.0, 1.0))
|
|
yaw_ema = ema_alpha * y + (1.0 - ema_alpha) * yaw_ema
|
|
pitch_ema = ema_alpha * p + (1.0 - ema_alpha) * pitch_ema
|
|
yaw[i] = yaw_ema
|
|
pitch[i] = pitch_ema
|
|
if n > 0:
|
|
yaw[-1] = yaw[-2]
|
|
pitch[-1] = pitch[-2]
|
|
return yaw, pitch
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Drawing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_FONT_CANDIDATES: tuple[str, ...] = (
|
|
"/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
|
|
"/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
|
|
"/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
|
|
)
|
|
|
|
|
|
def _load_font(size: int) -> ImageFont.ImageFont:
|
|
for path in _FONT_CANDIDATES:
|
|
if os.path.exists(path):
|
|
try:
|
|
return ImageFont.truetype(path, size=size)
|
|
except Exception:
|
|
continue
|
|
return ImageFont.load_default()
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _Layout:
|
|
width: int
|
|
height: int
|
|
|
|
@property
|
|
def key_size(self) -> int:
|
|
return max(32, int(self.height * 0.08))
|
|
|
|
@property
|
|
def key_gap(self) -> int:
|
|
return max(4, int(self.key_size * 0.15))
|
|
|
|
@property
|
|
def key_radius(self) -> int:
|
|
return max(4, int(self.key_size * 0.2))
|
|
|
|
|
|
class ActionOverlayRenderer:
|
|
"""Renders the WASD-cluster + rotation-joystick overlay onto video frames."""
|
|
|
|
CORNER_CHOICES = ("bottom-left", "bottom-right", "top-left", "top-right")
|
|
|
|
def __init__(self, width: int, height: int):
|
|
self.layout = _Layout(int(width), int(height))
|
|
self.width, self.height = self.layout.width, self.layout.height
|
|
self.font = _load_font(int(self.layout.key_size * 0.5))
|
|
self._key_tiles = self._build_key_tiles()
|
|
|
|
def _build_key_tiles(self) -> dict[tuple[str, bool], Image.Image]:
|
|
sz, r = self.layout.key_size, self.layout.key_radius
|
|
tiles: dict[tuple[str, bool], Image.Image] = {}
|
|
for key in ("W", "A", "S", "D"):
|
|
for pressed in (False, True):
|
|
fill = (255, 255, 255, 200) if pressed else (0, 0, 0, 100)
|
|
outline = (255, 255, 255, 255) if pressed else (255, 255, 255, 60)
|
|
text_color = (0, 0, 0, 220) if pressed else (255, 255, 255, 180)
|
|
tile = Image.new("RGBA", (sz, sz), (0, 0, 0, 0))
|
|
d = ImageDraw.Draw(tile)
|
|
d.rounded_rectangle(
|
|
[0, 0, sz - 1, sz - 1],
|
|
radius=r,
|
|
fill=fill,
|
|
outline=outline,
|
|
width=max(1, int(sz * 0.03)),
|
|
)
|
|
bbox = d.textbbox((0, 0), key, font=self.font)
|
|
tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
|
|
d.text((sz / 2 - tw / 2, sz / 2 - th / 2 - 2), key, fill=text_color, font=self.font)
|
|
tiles[(key, pressed)] = tile
|
|
return tiles
|
|
|
|
def _draw_joystick(self, canvas: Image.Image, cx: int, cy: int, yaw: float, pitch: float) -> None:
|
|
yaw = float(np.clip(yaw, -1.0, 1.0))
|
|
pitch = float(np.clip(pitch, -1.0, 1.0))
|
|
radius = max(30, int((self.layout.key_size * 2 + self.layout.key_gap) * 0.47))
|
|
|
|
# Outer shadow / glassy plate.
|
|
shadow = Image.new("RGBA", (self.width, self.height), (0, 0, 0, 0))
|
|
ImageDraw.Draw(shadow).ellipse(
|
|
[cx - radius - 14, cy - radius - 14, cx + radius + 14, cy + radius + 14],
|
|
fill=(0, 0, 0, 88),
|
|
)
|
|
canvas.alpha_composite(shadow.filter(ImageFilter.GaussianBlur(max(8, int(radius * 0.16)))))
|
|
|
|
d = ImageDraw.Draw(canvas)
|
|
d.ellipse(
|
|
[cx - radius, cy - radius, cx + radius, cy + radius],
|
|
fill=(7, 9, 13, 104),
|
|
outline=(255, 255, 255, 95),
|
|
width=max(1, int(radius * 0.035)),
|
|
)
|
|
d.ellipse(
|
|
[cx - radius - 7, cy - radius - 7, cx + radius + 7, cy + radius + 7],
|
|
outline=(255, 255, 255, 42),
|
|
width=max(1, int(radius * 0.025)),
|
|
)
|
|
d.line([cx - radius * 0.63, cy, cx + radius * 0.63, cy], fill=(255, 255, 255, 56), width=1)
|
|
d.line([cx, cy - radius * 0.63, cx, cy + radius * 0.63], fill=(255, 255, 255, 56), width=1)
|
|
|
|
# Direction triangles + active arcs.
|
|
marker_offset = int(radius * 0.78)
|
|
marker_size = max(7, int(radius * 0.16))
|
|
self._draw_arrow(d, cx + marker_offset, cy, "right", yaw > 0.08, marker_size)
|
|
self._draw_arrow(d, cx - marker_offset, cy, "left", yaw < -0.08, marker_size)
|
|
self._draw_arrow(d, cx, cy - marker_offset, "up", pitch < -0.08, marker_size)
|
|
self._draw_arrow(d, cx, cy + marker_offset, "down", pitch > 0.08, marker_size)
|
|
|
|
# Joystick knob + travel line.
|
|
max_offset = radius * 0.48
|
|
kx = int(cx + yaw * max_offset)
|
|
ky = int(cy + pitch * max_offset)
|
|
glow = Image.new("RGBA", (self.width, self.height), (0, 0, 0, 0))
|
|
gd = ImageDraw.Draw(glow)
|
|
gd.line([cx, cy, kx, ky], fill=(255, 255, 255, 78), width=max(3, int(radius * 0.055)))
|
|
gd.ellipse(
|
|
[kx - radius * 0.22, ky - radius * 0.22, kx + radius * 0.22, ky + radius * 0.22], fill=(255, 255, 255, 110)
|
|
)
|
|
canvas.alpha_composite(glow.filter(ImageFilter.GaussianBlur(max(5, int(radius * 0.09)))))
|
|
|
|
d = ImageDraw.Draw(canvas)
|
|
d.line([cx, cy, kx, ky], fill=(255, 255, 255, 120), width=max(1, int(radius * 0.025)))
|
|
kr = max(7, int(radius * 0.13))
|
|
d.ellipse(
|
|
[kx - kr, ky - kr, kx + kr, ky + kr], fill=(255, 255, 255, 230), outline=(255, 255, 255, 255), width=1
|
|
)
|
|
ir = max(3, int(kr * 0.36))
|
|
d.ellipse([kx - ir, ky - ir, kx + ir, ky + ir], fill=(20, 24, 30, 170))
|
|
|
|
@staticmethod
|
|
def _draw_arrow(d: ImageDraw.ImageDraw, cx: int, cy: int, direction: str, active: bool, size: int) -> None:
|
|
if direction == "right":
|
|
pts = [(cx - size * 0.55, cy - size), (cx + size * 0.65, cy), (cx - size * 0.55, cy + size)]
|
|
elif direction == "left":
|
|
pts = [(cx + size * 0.55, cy - size), (cx - size * 0.65, cy), (cx + size * 0.55, cy + size)]
|
|
elif direction == "up":
|
|
pts = [(cx - size, cy + size * 0.55), (cx, cy - size * 0.65), (cx + size, cy + size * 0.55)]
|
|
else:
|
|
pts = [(cx - size, cy - size * 0.55), (cx, cy + size * 0.65), (cx + size, cy - size * 0.55)]
|
|
d.polygon(pts, fill=(255, 255, 255, 210) if active else (255, 255, 255, 72))
|
|
|
|
def render_panel(
|
|
self,
|
|
pressed_keys: Sequence[str],
|
|
yaw: float,
|
|
pitch: float,
|
|
corner: str = "bottom-left",
|
|
) -> Image.Image:
|
|
"""Return an RGBA overlay of size ``(width, height)``."""
|
|
canvas = Image.new("RGBA", (self.width, self.height), (0, 0, 0, 0))
|
|
sz, gap = self.layout.key_size, self.layout.key_gap
|
|
margin = int(self.height * 0.05)
|
|
cluster_w = sz * 3 + gap * 2
|
|
cluster_h = sz * 2 + gap
|
|
joy_radius = max(30, int(cluster_h * 0.47))
|
|
joy_gap = int(sz * 1.0)
|
|
|
|
if corner == "bottom-right":
|
|
sx = self.width - margin - cluster_w
|
|
sy = self.height - margin - cluster_h
|
|
jcx = sx - joy_gap - joy_radius
|
|
elif corner == "top-left":
|
|
sx, sy = margin, margin
|
|
jcx = sx + cluster_w + joy_gap + joy_radius
|
|
elif corner == "top-right":
|
|
sx = self.width - margin - cluster_w
|
|
sy = margin
|
|
jcx = sx - joy_gap - joy_radius
|
|
else: # bottom-left default
|
|
sx, sy = margin, self.height - margin - cluster_h
|
|
jcx = sx + cluster_w + joy_gap + joy_radius
|
|
jcy = sy + cluster_h // 2
|
|
|
|
# Soft panel shadow.
|
|
shadow = Image.new("RGBA", (self.width, self.height), (0, 0, 0, 0))
|
|
ImageDraw.Draw(shadow).rounded_rectangle(
|
|
[sx - 8, sy - 8, sx + cluster_w + 8, sy + cluster_h + 8],
|
|
radius=max(10, int(self.layout.key_radius * 1.35)),
|
|
fill=(0, 0, 0, 74),
|
|
)
|
|
canvas.alpha_composite(shadow.filter(ImageFilter.GaussianBlur(max(7, int(sz * 0.18)))))
|
|
|
|
positions = {
|
|
"W": (sx + sz + gap, sy),
|
|
"A": (sx, sy + sz + gap),
|
|
"S": (sx + sz + gap, sy + sz + gap),
|
|
"D": (sx + (sz + gap) * 2, sy + sz + gap),
|
|
}
|
|
for key in ("W", "A", "S", "D"):
|
|
canvas.alpha_composite(self._key_tiles[(key, key in pressed_keys)], dest=positions[key])
|
|
self._draw_joystick(canvas, jcx, jcy, yaw, pitch)
|
|
return canvas
|
|
|
|
|
|
def apply_overlay(
|
|
video_hwc: np.ndarray,
|
|
c2w: np.ndarray,
|
|
*,
|
|
corner: str = "bottom-left",
|
|
) -> np.ndarray:
|
|
"""Composite the action overlay onto each frame.
|
|
|
|
Args:
|
|
video_hwc: ``(T, H, W, 3)`` uint8 video.
|
|
c2w: ``(T_pose, 4, 4)`` camera-to-world poses driving the overlay.
|
|
Truncated to ``video_hwc.shape[0]`` frames if longer.
|
|
corner: Panel placement.
|
|
|
|
Returns:
|
|
``(T, H, W, 3)`` uint8 array with the overlay composited.
|
|
"""
|
|
T, H, W = video_hwc.shape[:3]
|
|
n_poses = min(int(c2w.shape[0]), T)
|
|
poses = c2w[:n_poses].astype(np.float32)
|
|
|
|
trans, rots = _per_frame_deltas(poses)
|
|
keys = _translation_keys(trans)
|
|
yaw, pitch = _normalised_rotation(rots)
|
|
|
|
renderer = ActionOverlayRenderer(width=W, height=H)
|
|
out = np.empty_like(video_hwc)
|
|
for t in range(T):
|
|
i = min(t, len(keys) - 1)
|
|
panel = renderer.render_panel(
|
|
pressed_keys=keys[i],
|
|
yaw=float(yaw[i]),
|
|
pitch=float(pitch[i]),
|
|
corner=corner,
|
|
)
|
|
frame = Image.fromarray(video_hwc[t]).convert("RGBA")
|
|
frame.alpha_composite(panel)
|
|
out[t] = np.asarray(frame.convert("RGB"), dtype=np.uint8)
|
|
return out
|