reorder
This commit is contained in:
mike
433
tour_comfy/orbit_module.py
Normal file
433
tour_comfy/orbit_module.py
Normal file
@@ -0,0 +1,433 @@
|
||||
"""
|
||||
orbit_module.py — 2.5D actor orbit preview via depth-card parallax.
|
||||
|
||||
Pipeline:
|
||||
1. load actor image — use provided path directly (selection is caller's responsibility)
|
||||
2. create_depth_map — fake depth from alpha mask distance transform
|
||||
3. find_bg_plate — static background for compositing (original for nobg; blurred for opaque)
|
||||
4. render_orbit — per-frame: parallax-warp actor, composite over static bg
|
||||
5. save_orbit_output — write RGBA PNGs + MP4
|
||||
|
||||
The "orbit" illusion requires a static reference (background plate).
|
||||
Without it the viewer has nothing to anchor on and it reads as a side-slide.
|
||||
|
||||
Usage:
|
||||
from orbit_module import run_orbit_pipeline
|
||||
result = run_orbit_pipeline(image_path, output_dir)
|
||||
|
||||
CLI: see orbit_poc.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import math
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from scipy.ndimage import distance_transform_edt
|
||||
|
||||
__all__ = [
|
||||
"create_depth_map",
|
||||
"render_orbit_frame",
|
||||
"render_orbit",
|
||||
"save_orbit_output",
|
||||
"run_orbit_pipeline",
|
||||
]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Image loading helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _load_rgba(path: str) -> np.ndarray:
|
||||
"""Load any image as RGBA uint8 H×W×4."""
|
||||
img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
|
||||
if img is None:
|
||||
raise FileNotFoundError(f"Cannot read image: {path}")
|
||||
if img.ndim == 2:
|
||||
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGRA)
|
||||
elif img.shape[2] == 3:
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
|
||||
return cv2.cvtColor(img, cv2.COLOR_BGRA2RGBA)
|
||||
|
||||
|
||||
def _has_real_alpha(rgba: np.ndarray) -> bool:
|
||||
"""True if the image contains meaningful transparency (not just all-255)."""
|
||||
alpha = rgba[:, :, 3]
|
||||
transparent_pct = float((alpha < 32).mean())
|
||||
return transparent_pct > 0.05 # >5% transparent pixels = real alpha
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Background plate
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _find_original_for_nobg(actor_path: str) -> str | None:
|
||||
"""
|
||||
Given a .nobg.png sidecar path, find the original opaque image.
|
||||
e.g. foo.nobg.png → foo.png or foo.jpg
|
||||
"""
|
||||
root, _ = os.path.splitext(actor_path)
|
||||
if not root.endswith(".nobg"):
|
||||
return None
|
||||
base = root[: -len(".nobg")]
|
||||
for ext in (".png", ".jpg", ".jpeg", ".webp"):
|
||||
p = base + ext
|
||||
if os.path.exists(p):
|
||||
return p
|
||||
return None
|
||||
|
||||
|
||||
def _make_bg_plate(actor_rgba: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Build a static background plate for opaque images (no nobg available).
|
||||
|
||||
Strategy: blur the source image with a large kernel. The blurred copy stays
|
||||
fixed while the sharp actor layer shifts — creates subtle depth separation.
|
||||
"""
|
||||
H, W = actor_rgba.shape[:2]
|
||||
# Large blur: simulates out-of-focus background
|
||||
blurred_rgb = cv2.GaussianBlur(actor_rgba[:, :, :3], (0, 0), max(H, W) * 0.04)
|
||||
plate = np.dstack([blurred_rgb, np.full((H, W), 255, dtype=np.uint8)])
|
||||
return plate.astype(np.uint8)
|
||||
|
||||
|
||||
def get_bg_plate(actor_path: str, actor_rgba: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Return the background plate for the orbit:
|
||||
- .nobg.png: load the matching original opaque image (best result)
|
||||
- Opaque image: return blurred copy (subtle but functional)
|
||||
Plate is resized to match actor_rgba dimensions.
|
||||
"""
|
||||
H, W = actor_rgba.shape[:2]
|
||||
|
||||
orig = _find_original_for_nobg(actor_path)
|
||||
if orig:
|
||||
bg = _load_rgba(orig)
|
||||
if bg.shape[:2] != (H, W):
|
||||
bg = cv2.resize(bg, (W, H), interpolation=cv2.INTER_AREA)
|
||||
bg[:, :, 3] = 255 # ensure fully opaque background
|
||||
return bg
|
||||
|
||||
return _make_bg_plate(actor_rgba)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Depth map
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def create_depth_map(image_rgba: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Float32 H×W depth in [0,1]. 1 = closest (subject centre), 0 = far/background.
|
||||
|
||||
Uses the alpha mask distance transform:
|
||||
- For transparent-bg images: EDT of the foreground mask (subject body)
|
||||
- For opaque images: EDT from image edges (assumes subject is centred)
|
||||
|
||||
Power-law shaping (^0.5) keeps the gradient gradual near the centre.
|
||||
"""
|
||||
alpha = image_rgba[:, :, 3]
|
||||
mask = (alpha > 32).astype(np.uint8)
|
||||
|
||||
if mask.sum() == 0:
|
||||
return np.zeros(alpha.shape, dtype=np.float32)
|
||||
|
||||
dist = distance_transform_edt(mask).astype(np.float32)
|
||||
max_d = dist.max()
|
||||
if max_d > 0:
|
||||
dist /= max_d
|
||||
return np.sqrt(dist)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Orbit rendering
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _alpha_composite(fg: np.ndarray, bg: np.ndarray) -> np.ndarray:
|
||||
"""Alpha-composite RGBA fg over RGBA bg. Returns RGBA uint8."""
|
||||
a = fg[:, :, 3:4].astype(np.float32) / 255.0
|
||||
out_rgb = fg[:, :, :3].astype(np.float32) * a + bg[:, :, :3].astype(np.float32) * (1.0 - a)
|
||||
out_a = fg[:, :, 3:4].astype(np.float32) + bg[:, :, 3:4].astype(np.float32) * (1.0 - a)
|
||||
return np.dstack([out_rgb.clip(0, 255), out_a.clip(0, 255)]).astype(np.uint8)
|
||||
|
||||
|
||||
def render_orbit_frame(
|
||||
actor_rgba: np.ndarray,
|
||||
depth: np.ndarray,
|
||||
theta: float,
|
||||
parallax_strength: float = 0.08,
|
||||
bg_rgba: np.ndarray | None = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Swing-mode frame: depth-based parallax shift only.
|
||||
Closer pixels (depth≈1) shift more than far pixels (depth≈0).
|
||||
Composited over static bg plate for perceivable depth.
|
||||
Returns RGBA uint8 H×W×4.
|
||||
"""
|
||||
H, W = actor_rgba.shape[:2]
|
||||
|
||||
shift_x = depth * (W * parallax_strength * math.sin(theta))
|
||||
shift_y = depth * (H * parallax_strength * 0.03 * -math.cos(theta))
|
||||
|
||||
yc, xc = np.mgrid[0:H, 0:W].astype(np.float32)
|
||||
map_x = (xc - shift_x).astype(np.float32)
|
||||
map_y = (yc - shift_y).astype(np.float32)
|
||||
|
||||
bgra = cv2.cvtColor(actor_rgba, cv2.COLOR_RGBA2BGRA)
|
||||
warped_bgra = cv2.remap(bgra, map_x, map_y,
|
||||
interpolation=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=(0, 0, 0, 0))
|
||||
warped = cv2.cvtColor(warped_bgra, cv2.COLOR_BGRA2RGBA)
|
||||
|
||||
if bg_rgba is None:
|
||||
return warped
|
||||
return _alpha_composite(warped, bg_rgba)
|
||||
|
||||
|
||||
def _perspective_card_frame(
|
||||
bgra_src: np.ndarray,
|
||||
theta: float,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Orbit-mode frame: simulate a flat card rotating around its vertical axis.
|
||||
|
||||
- Squishes width by |cos(θ)|, centred on canvas
|
||||
- Mirrors source for the back half (cos < 0) so the "back" is visible
|
||||
- Returns BGRA with transparent borders (ready for bg composite)
|
||||
|
||||
At θ=0° → full width, unmirrored (front face)
|
||||
At θ=90° → hair-thin line
|
||||
At θ=180°→ full width, mirrored (back face)
|
||||
"""
|
||||
H, W = bgra_src.shape[:2]
|
||||
cos_t = math.cos(theta)
|
||||
compress = abs(cos_t)
|
||||
is_back = cos_t < 0
|
||||
|
||||
if compress < 0.025:
|
||||
# Near-edge-on: return a one-pixel-wide vertical strip to avoid singularity
|
||||
out = np.zeros_like(bgra_src)
|
||||
mid = W // 2
|
||||
out[:, mid:mid+1] = bgra_src[:, mid:mid+1]
|
||||
return out
|
||||
|
||||
new_w = max(int(round(W * compress)), 2)
|
||||
x0 = (W - new_w) // 2
|
||||
x1 = x0 + new_w
|
||||
|
||||
# Source corners: mirror left↔right for the back face
|
||||
if is_back:
|
||||
src = np.float32([[W, 0], [0, 0], [0, H], [W, H]])
|
||||
else:
|
||||
src = np.float32([[0, 0], [W, 0], [W, H], [0, H]])
|
||||
|
||||
dst = np.float32([[x0, 0], [x1, 0], [x1, H], [x0, H]])
|
||||
|
||||
M = cv2.getPerspectiveTransform(src, dst)
|
||||
return cv2.warpPerspective(bgra_src, M, (W, H),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=(0, 0, 0, 0))
|
||||
|
||||
|
||||
def render_orbit(
|
||||
actor_rgba: np.ndarray,
|
||||
depth: np.ndarray,
|
||||
n_frames: int = 36,
|
||||
parallax_strength: float = 0.08,
|
||||
mode: str = "swing",
|
||||
max_angle_deg: float = 35.0,
|
||||
bg_rgba: np.ndarray | None = None,
|
||||
) -> list:
|
||||
"""
|
||||
Render all orbit frames.
|
||||
|
||||
mode='swing' — sinusoidal ±max_angle_deg depth-parallax, loops cleanly
|
||||
mode='orbit' — full 360° perspective card rotation (compress + mirror)
|
||||
|
||||
Returns list of RGBA uint8 frames.
|
||||
"""
|
||||
if mode == "swing":
|
||||
max_rad = math.radians(max_angle_deg)
|
||||
angles = [max_rad * math.sin(2 * math.pi * i / n_frames) for i in range(n_frames)]
|
||||
return [render_orbit_frame(actor_rgba, depth, theta, parallax_strength, bg_rgba)
|
||||
for theta in angles]
|
||||
|
||||
elif mode == "orbit":
|
||||
# Don't use the photo bg plate — the transparent areas should stay transparent
|
||||
# so the perspective compression (card getting thin at 90°, mirrored at 180°)
|
||||
# is clearly visible. Solid bg is added at MP4 write time.
|
||||
angles = [2 * math.pi * i / n_frames for i in range(n_frames)]
|
||||
bgra_src = cv2.cvtColor(actor_rgba, cv2.COLOR_RGBA2BGRA)
|
||||
return [
|
||||
cv2.cvtColor(_perspective_card_frame(bgra_src, theta), cv2.COLOR_BGRA2RGBA)
|
||||
for theta in angles
|
||||
]
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown mode: {mode!r}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Output saving
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _composite_over_solid(frame_rgba: np.ndarray, bg: tuple = (18, 18, 18)) -> np.ndarray:
|
||||
"""Alpha-composite RGBA over a solid colour; return BGR uint8 for ffmpeg."""
|
||||
rgb = frame_rgba[:, :, :3].astype(np.float32)
|
||||
a = frame_rgba[:, :, 3:4].astype(np.float32) / 255.0
|
||||
bg_f = np.array(bg, dtype=np.float32)
|
||||
out = (rgb * a + bg_f * (1.0 - a)).clip(0, 255).astype(np.uint8)
|
||||
return cv2.cvtColor(out, cv2.COLOR_RGB2BGR)
|
||||
|
||||
|
||||
def save_orbit_output(
|
||||
frames: list,
|
||||
output_dir: str,
|
||||
fps: int = 24,
|
||||
bg_color: tuple = (18, 18, 18),
|
||||
) -> dict:
|
||||
"""
|
||||
Write orbit_frames/frame_NNN.png (RGBA) and orbit_preview.mp4.
|
||||
Returns dict with paths.
|
||||
"""
|
||||
frames_dir = os.path.join(output_dir, "orbit_frames")
|
||||
os.makedirs(frames_dir, exist_ok=True)
|
||||
|
||||
frame_paths = []
|
||||
for i, frame in enumerate(frames):
|
||||
path = os.path.join(frames_dir, f"frame_{i:03d}.png")
|
||||
cv2.imwrite(path, cv2.cvtColor(frame, cv2.COLOR_RGBA2BGRA))
|
||||
frame_paths.append(path)
|
||||
|
||||
video_path = os.path.join(output_dir, "orbit_preview.mp4")
|
||||
_frames_to_mp4(frames, video_path, fps=fps, bg_color=bg_color)
|
||||
|
||||
return {
|
||||
"frames_dir": frames_dir,
|
||||
"n_frames": len(frames),
|
||||
"video_path": video_path,
|
||||
"frame_paths": frame_paths,
|
||||
}
|
||||
|
||||
|
||||
def _frames_to_mp4(
|
||||
frames: list, output_path: str, fps: int = 24, bg_color: tuple = (18, 18, 18)
|
||||
) -> None:
|
||||
"""Composite frames over solid bg, write MP4 via ffmpeg."""
|
||||
if not frames:
|
||||
return
|
||||
with tempfile.TemporaryDirectory(prefix="orbit_mp4_") as tmpdir:
|
||||
for i, frame in enumerate(frames):
|
||||
bgr = _composite_over_solid(frame, bg_color)
|
||||
cv2.imwrite(
|
||||
os.path.join(tmpdir, f"frame_{i:04d}.jpg"), bgr,
|
||||
[cv2.IMWRITE_JPEG_QUALITY, 95],
|
||||
)
|
||||
H, W = frames[0].shape[:2]
|
||||
W2, H2 = W - (W % 2), H - (H % 2)
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-framerate", str(fps),
|
||||
"-i", os.path.join(tmpdir, "frame_%04d.jpg"),
|
||||
"-vf", f"crop={W2}:{H2}:0:0",
|
||||
"-c:v", "libx264", "-pix_fmt", "yuv420p",
|
||||
"-crf", "18", "-movflags", "+faststart",
|
||||
output_path,
|
||||
]
|
||||
r = subprocess.run(cmd, capture_output=True, text=True)
|
||||
if r.returncode != 0:
|
||||
raise RuntimeError(f"ffmpeg failed: {r.stderr[-600:]}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Debug helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _save_debug(actor_rgba, actor_path, bg_rgba, debug_dir):
|
||||
os.makedirs(debug_dir, exist_ok=True)
|
||||
if os.path.exists(actor_path):
|
||||
shutil.copy2(actor_path, os.path.join(debug_dir, "selected_frame.png"))
|
||||
cv2.imwrite(os.path.join(debug_dir, "actor_rgba.png"),
|
||||
cv2.cvtColor(actor_rgba, cv2.COLOR_RGBA2BGRA))
|
||||
cv2.imwrite(os.path.join(debug_dir, "mask.png"), actor_rgba[:, :, 3])
|
||||
if bg_rgba is not None:
|
||||
cv2.imwrite(os.path.join(debug_dir, "bg_plate.png"),
|
||||
cv2.cvtColor(bg_rgba, cv2.COLOR_RGBA2BGRA))
|
||||
|
||||
|
||||
def _save_depth_debug(depth, debug_dir):
|
||||
os.makedirs(debug_dir, exist_ok=True)
|
||||
d8 = (depth * 255).astype(np.uint8)
|
||||
cv2.imwrite(os.path.join(debug_dir, "depth.png"), d8)
|
||||
cv2.imwrite(os.path.join(debug_dir, "depth_colorized.png"),
|
||||
cv2.applyColorMap(d8, cv2.COLORMAP_MAGMA))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Full pipeline
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def run_orbit_pipeline(
|
||||
image_path: str,
|
||||
output_dir: str,
|
||||
n_frames: int = 36,
|
||||
parallax_strength: float = 0.08,
|
||||
mode: str = "swing",
|
||||
fps: int = 24,
|
||||
max_angle_deg: float = 35.0,
|
||||
debug: bool = True,
|
||||
) -> dict:
|
||||
"""
|
||||
Full pipeline: load → bg-plate → depth → render → save.
|
||||
|
||||
image_path: the specific image to orbit (caller selects; no sharpness heuristic)
|
||||
Returns dict: actor_path, frames_dir, video_path, n_frames, debug_dir, has_alpha, has_bg
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
debug_dir = os.path.join(output_dir, "debug")
|
||||
|
||||
# 1. Load actor — prefer nobg sidecar for cleaner depth
|
||||
actor_path = image_path
|
||||
root, _ = os.path.splitext(image_path)
|
||||
nobg_candidate = root + ".nobg.png"
|
||||
if not root.endswith(".nobg") and os.path.exists(nobg_candidate):
|
||||
actor_path = nobg_candidate
|
||||
|
||||
actor_rgba = _load_rgba(actor_path)
|
||||
has_alpha = _has_real_alpha(actor_rgba)
|
||||
|
||||
# 2. Background plate (static reference — essential for perceivable depth)
|
||||
bg_rgba = get_bg_plate(actor_path, actor_rgba)
|
||||
has_bg = bg_rgba is not None
|
||||
|
||||
if debug:
|
||||
_save_debug(actor_rgba, actor_path, bg_rgba, debug_dir)
|
||||
|
||||
# 3. Depth map
|
||||
depth = create_depth_map(actor_rgba)
|
||||
|
||||
if debug:
|
||||
_save_depth_debug(depth, debug_dir)
|
||||
|
||||
# 4. Render
|
||||
frames = render_orbit(
|
||||
actor_rgba, depth,
|
||||
n_frames=n_frames,
|
||||
parallax_strength=parallax_strength,
|
||||
mode=mode,
|
||||
max_angle_deg=max_angle_deg,
|
||||
bg_rgba=bg_rgba,
|
||||
)
|
||||
|
||||
# 5. Save
|
||||
result = save_orbit_output(frames, output_dir, fps=fps)
|
||||
result.update({
|
||||
"actor_path": actor_path,
|
||||
"debug_dir": debug_dir,
|
||||
"has_alpha": has_alpha,
|
||||
"has_bg": has_bg,
|
||||
})
|
||||
return result
|
||||
Reference in New Issue
Block a user