434 lines
15 KiB
Python
434 lines
15 KiB
Python
"""
|
||
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
|