203 lines
7.5 KiB
Python
203 lines
7.5 KiB
Python
from __future__ import annotations
|
|
|
|
# Keyframe interaction and interpolation semantics are adapted from
|
|
# ComfyUI-CrossViewWarp by cseti007 at upstream commit
|
|
# 7c758557c827dfc9a65bab0823b494344d653985 (Apache-2.0). See NOTICE.
|
|
|
|
import json
|
|
import logging
|
|
import math
|
|
from typing import Any, Dict, List, Tuple
|
|
|
|
import numpy as np
|
|
|
|
from comfy_extras.sam3d_body.utils import apply_camera_override
|
|
|
|
|
|
def _subject_distance(frame: List[Dict[str, Any]]) -> float:
|
|
centers = []
|
|
for person in frame:
|
|
cam_t = person.get("pred_cam_t")
|
|
if cam_t is None:
|
|
continue
|
|
center = np.asarray(cam_t, dtype=np.float32).reshape(3)
|
|
vertices = person.get("pred_vertices")
|
|
if vertices is not None:
|
|
center = center + np.asarray(vertices, dtype=np.float32).reshape(-1, 3).mean(axis=0)
|
|
centers.append(center)
|
|
if not centers:
|
|
return 1.0
|
|
return max(0.1, float(np.linalg.norm(np.mean(np.stack(centers), axis=0))))
|
|
|
|
|
|
def _orbit_camera(azimuth: float, elevation: float, distance: float, fov: float, base_distance: float) -> Dict[str, Any]:
|
|
az = math.radians(float(azimuth))
|
|
el = math.radians(float(elevation))
|
|
radius = max(0.01, float(distance)) * base_distance
|
|
camera = {
|
|
"position": {
|
|
"x": radius * math.cos(el) * math.sin(az),
|
|
"y": radius * math.sin(el),
|
|
"z": radius * math.cos(el) * math.cos(az),
|
|
},
|
|
"target": {"x": 0.0, "y": 0.0, "z": 0.0},
|
|
"zoom": 1.0,
|
|
"cameraType": "perspective",
|
|
}
|
|
if fov > 0.0:
|
|
camera["fov"] = float(fov)
|
|
return camera
|
|
|
|
|
|
def _wrap_degrees(value: float) -> float:
|
|
return ((value + 180.0) % 360.0) - 180.0
|
|
|
|
|
|
def _ease(value: float, mode: str) -> float:
|
|
if mode == "ease_in_out":
|
|
return 0.5 - 0.5 * math.cos(math.pi * value)
|
|
if mode == "ease_in":
|
|
return value * value
|
|
if mode == "ease_out":
|
|
return 1.0 - (1.0 - value) * (1.0 - value)
|
|
return value
|
|
|
|
|
|
def _catmull(p0: float, p1: float, p2: float, p3: float, value: float) -> float:
|
|
return 0.5 * (
|
|
2.0 * p1
|
|
+ (-p0 + p2) * value
|
|
+ (2.0 * p0 - 5.0 * p1 + 4.0 * p2 - p3) * value * value
|
|
+ (-p0 + 3.0 * p1 - 3.0 * p2 + p3) * value * value * value
|
|
)
|
|
|
|
|
|
def _segment_value(values: List[float], segment: int, value: float, smooth: bool) -> float:
|
|
p1, p2 = values[segment], values[segment + 1]
|
|
if not smooth or len(values) < 3:
|
|
return p1 + (p2 - p1) * value
|
|
p0 = values[segment - 1] if segment > 0 else p1 + (p1 - p2)
|
|
p3 = values[segment + 2] if segment + 2 < len(values) else p2 + (p2 - p1)
|
|
return _catmull(p0, p1, p2, p3, value)
|
|
|
|
|
|
def parse_keyframes(raw: str, frame_count: int) -> List[Tuple[int, float, float, float]]:
|
|
text = "" if raw is None else str(raw).strip()
|
|
if not text or text in ("[]", "null"):
|
|
return []
|
|
try:
|
|
data = json.loads(text)
|
|
except ValueError as exc:
|
|
raise ValueError(
|
|
"SAM3D BodyMod: keyframes is not valid JSON. Expected "
|
|
'[{"f":1,"az":-30,"el":20,"dist":1.0},{"f":49,"az":45,"el":10,"dist":1.2}]'
|
|
) from exc
|
|
if not isinstance(data, list):
|
|
raise ValueError("SAM3D BodyMod: keyframes must be a JSON list")
|
|
|
|
out = []
|
|
for index, keyframe in enumerate(data):
|
|
if not isinstance(keyframe, dict):
|
|
raise ValueError(f"SAM3D BodyMod: keyframe #{index} is not an object")
|
|
try:
|
|
frame = int(round(float(keyframe["f"])))
|
|
azimuth = float(keyframe["az"])
|
|
elevation = float(keyframe["el"])
|
|
distance = float(keyframe["dist"])
|
|
except (KeyError, TypeError, ValueError):
|
|
raise ValueError(
|
|
f"SAM3D BodyMod: keyframe #{index} needs numeric f, az, el and dist values"
|
|
) from None
|
|
if frame < 1 or frame > frame_count:
|
|
raise ValueError(
|
|
f"SAM3D BodyMod: keyframe #{index} uses frame {frame}, but this pose has frames 1-{frame_count}"
|
|
)
|
|
if not -90.0 <= elevation <= 90.0 or not 0.2 <= distance <= 3.0:
|
|
raise ValueError(
|
|
f"SAM3D BodyMod: keyframe #{index} elevation must be -90..90 and distance 0.2..3.0"
|
|
)
|
|
out.append((frame, azimuth, elevation, distance))
|
|
out.sort(key=lambda item: item[0])
|
|
if len({item[0] for item in out}) != len(out):
|
|
raise ValueError("SAM3D BodyMod: two camera keyframes share the same frame number")
|
|
return out
|
|
|
|
|
|
def _prepare_path(keyframes: List[Tuple[int, float, float, float]]) -> Dict[str, List[float]]:
|
|
azimuths = [float(keyframes[0][1])]
|
|
for keyframe in keyframes[1:]:
|
|
azimuths.append(azimuths[-1] + _wrap_degrees(float(keyframe[1]) - azimuths[-1]))
|
|
return {
|
|
"frames": [item[0] for item in keyframes],
|
|
"azimuths": azimuths,
|
|
"elevations": [item[2] for item in keyframes],
|
|
"distances": [item[3] for item in keyframes],
|
|
}
|
|
|
|
|
|
def sample_path(path: Dict[str, List[float]], frame: int, easing: str, smooth: bool) -> Tuple[float, float, float]:
|
|
frames = path["frames"]
|
|
if frame <= frames[0]:
|
|
return path["azimuths"][0], path["elevations"][0], path["distances"][0]
|
|
if frame >= frames[-1]:
|
|
return path["azimuths"][-1], path["elevations"][-1], path["distances"][-1]
|
|
segment = next(index for index in range(len(frames) - 1) if frames[index] <= frame <= frames[index + 1])
|
|
value = _ease((frame - frames[segment]) / float(frames[segment + 1] - frames[segment]), easing)
|
|
azimuth = _wrap_degrees(_segment_value(path["azimuths"], segment, value, smooth))
|
|
elevation = float(np.clip(_segment_value(path["elevations"], segment, value, smooth), -90.0, 90.0))
|
|
distance = float(np.clip(_segment_value(path["distances"], segment, value, smooth), 0.2, 3.0))
|
|
return azimuth, elevation, distance
|
|
|
|
|
|
def drive_camera(
|
|
pose_data: Dict[str, Any],
|
|
azimuth: float,
|
|
elevation: float,
|
|
distance: float,
|
|
fov: float,
|
|
use_keyframes: bool,
|
|
frame_count: int,
|
|
keyframes: str,
|
|
interp_motion: str,
|
|
interpolation: str,
|
|
) -> Dict[str, Any]:
|
|
frames = pose_data.get("frames", [])
|
|
if not frames:
|
|
return pose_data
|
|
H, W = pose_data["image_size"]
|
|
count = len(frames)
|
|
if use_keyframes and frame_count and int(frame_count) != count:
|
|
logging.warning(
|
|
"SAM3D BodyMod: camera frame_count is %d but pose data has %d frames; the pose data length is used",
|
|
int(frame_count), count,
|
|
)
|
|
parsed = parse_keyframes(keyframes, count) if use_keyframes else []
|
|
path = _prepare_path(parsed) if len(parsed) >= 2 else None
|
|
if len(parsed) == 1:
|
|
azimuth, elevation, distance = parsed[0][1:]
|
|
|
|
out_frames = []
|
|
for index, frame in enumerate(frames):
|
|
if path is not None:
|
|
current = sample_path(path, index + 1, interp_motion, interpolation == "smooth")
|
|
else:
|
|
current = float(azimuth), float(elevation), float(distance)
|
|
camera = _orbit_camera(*current, float(fov), _subject_distance(frame))
|
|
one_frame = dict(pose_data)
|
|
one_frame["frames"] = [frame]
|
|
out_frames.append(apply_camera_override(one_frame, camera, int(H), int(W))["frames"][0])
|
|
|
|
out = dict(pose_data)
|
|
out["frames"] = out_frames
|
|
out["_bodymod_camera_drive"] = {
|
|
"azimuth": float(azimuth),
|
|
"elevation": float(elevation),
|
|
"distance": float(distance),
|
|
"fov": float(fov),
|
|
"use_keyframes": bool(use_keyframes),
|
|
"keyframes": str(keyframes),
|
|
"interp_motion": str(interp_motion),
|
|
"interpolation": str(interpolation),
|
|
}
|
|
return out
|