Green screen pipeline: new tools, AnimatedBackground, caption burn fixes
Add green_screen_processor (auto-detect + rembg fallback) and green_screen_composite (4 layout presets with alpha compositing) tools to automate the full keying-to-composite pipeline. Remotion: add AnimatedBackground with gradient mesh and floating orbs to Explainer, fix caption burn tool (remove entry point arg, auto-detect dimensions, extend TalkingHead duration to 300s). Update scene-director, compose-director, and asset-director skill docs with green screen workflow steps and component constraints.
This commit is contained in:
calesthio
@@ -0,0 +1,430 @@
|
||||
"""Green screen composite tool for talking-head pipeline.
|
||||
|
||||
Composites a keyed speaker (dark/solid background) over a Remotion
|
||||
background video with layout presets. Supports news anchor, full behind,
|
||||
picture-in-picture, and split layouts.
|
||||
|
||||
Uses PIL/numpy for frame-level alpha compositing and FFmpeg for
|
||||
frame extraction, encoding, and audio muxing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ResumeSupport,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class GreenScreenComposite(BaseTool):
|
||||
name = "green_screen_composite"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "video_post"
|
||||
provider = "ffmpeg"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["cmd:ffmpeg", "python:numpy", "python:PIL"]
|
||||
install_instructions = (
|
||||
"Install FFmpeg: https://ffmpeg.org/download.html — "
|
||||
"pip install numpy Pillow"
|
||||
)
|
||||
agent_skills = ["ffmpeg"]
|
||||
|
||||
capabilities = [
|
||||
"green_screen_composite",
|
||||
"speaker_overlay",
|
||||
"layout_preset",
|
||||
"alpha_composite",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["speaker_path", "background_path", "output_path"],
|
||||
"properties": {
|
||||
"speaker_path": {
|
||||
"type": "string",
|
||||
"description": "Path to keyed speaker video (dark bg, from green_screen_processor)",
|
||||
},
|
||||
"background_path": {
|
||||
"type": "string",
|
||||
"description": "Path to Remotion background video",
|
||||
},
|
||||
"output_path": {
|
||||
"type": "string",
|
||||
"description": "Output composite video path",
|
||||
},
|
||||
"original_audio_path": {
|
||||
"type": "string",
|
||||
"description": "Path to original footage to extract audio from",
|
||||
},
|
||||
"layout": {
|
||||
"type": "string",
|
||||
"enum": ["news_anchor", "full_behind", "pip", "split"],
|
||||
"default": "news_anchor",
|
||||
"description": (
|
||||
"news_anchor=speaker bottom-center over shifted bg, "
|
||||
"full_behind=speaker full-frame on bg, "
|
||||
"pip=speaker 30% bottom-right, "
|
||||
"split=speaker left 50% bg right 50%"
|
||||
),
|
||||
},
|
||||
"speaker_scale": {
|
||||
"type": "number",
|
||||
"default": 0.65,
|
||||
"description": "Scale factor for speaker layer",
|
||||
},
|
||||
"bg_shift_up": {
|
||||
"type": "integer",
|
||||
"default": 300,
|
||||
"description": "Pixels to shift background content upward",
|
||||
},
|
||||
"bg_color_hex": {
|
||||
"type": "string",
|
||||
"default": "#0E172A",
|
||||
"description": "The keyed speaker's background color for alpha creation",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=4, ram_mb=4096, vram_mb=0, disk_mb=8000, network_required=False
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=1, retryable_errors=["FFmpeg error"])
|
||||
resume_support = ResumeSupport.FROM_START
|
||||
idempotency_key_fields = [
|
||||
"speaker_path", "background_path", "layout",
|
||||
"speaker_scale", "bg_shift_up", "bg_color_hex",
|
||||
]
|
||||
side_effects = ["writes composite video to output_path"]
|
||||
user_visible_verification = [
|
||||
"Watch output — speaker should be cleanly composited without color fringing",
|
||||
"Check layout positioning matches the chosen preset",
|
||||
"Verify audio is synced if original_audio_path was provided",
|
||||
]
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
speaker_path = Path(inputs["speaker_path"])
|
||||
background_path = Path(inputs["background_path"])
|
||||
output_path = Path(inputs["output_path"])
|
||||
original_audio_path = inputs.get("original_audio_path")
|
||||
layout = inputs.get("layout", "news_anchor")
|
||||
speaker_scale = inputs.get("speaker_scale", 0.65)
|
||||
bg_shift_up = inputs.get("bg_shift_up", 300)
|
||||
bg_color_hex = inputs.get("bg_color_hex", "#0E172A")
|
||||
|
||||
if not speaker_path.exists():
|
||||
return ToolResult(success=False, error=f"Speaker video not found: {speaker_path}")
|
||||
if not background_path.exists():
|
||||
return ToolResult(success=False, error=f"Background video not found: {background_path}")
|
||||
if original_audio_path and not Path(original_audio_path).exists():
|
||||
return ToolResult(success=False, error=f"Audio source not found: {original_audio_path}")
|
||||
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
start = time.time()
|
||||
|
||||
# Parse bg color
|
||||
bg_color = self._parse_hex_color(bg_color_hex)
|
||||
|
||||
# Step 1: Probe both videos
|
||||
speaker_info = self._probe_video(speaker_path)
|
||||
bg_info = self._probe_video(background_path)
|
||||
|
||||
if not speaker_info or not bg_info:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Failed to probe one or both input videos",
|
||||
)
|
||||
|
||||
# Step 2: Use the LOWER fps (typically 15fps from speaker)
|
||||
target_fps = min(speaker_info["fps"], bg_info["fps"])
|
||||
if target_fps <= 0:
|
||||
target_fps = 15.0
|
||||
|
||||
# Determine output dimensions from background
|
||||
out_w = bg_info["width"]
|
||||
out_h = bg_info["height"]
|
||||
|
||||
# Use shorter duration
|
||||
duration = min(speaker_info["duration"], bg_info["duration"])
|
||||
|
||||
# Step 3: Extract frames from both videos
|
||||
temp_dir = output_path.parent / ".greenscreen_composite_tmp"
|
||||
speaker_frames_dir = temp_dir / "speaker"
|
||||
bg_frames_dir = temp_dir / "bg"
|
||||
comp_frames_dir = temp_dir / "composite"
|
||||
|
||||
for d in [speaker_frames_dir, bg_frames_dir, comp_frames_dir]:
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
self._extract_frames(speaker_path, speaker_frames_dir, target_fps)
|
||||
self._extract_frames(background_path, bg_frames_dir, target_fps)
|
||||
|
||||
# Get sorted frame lists
|
||||
speaker_frames = sorted(speaker_frames_dir.glob("*.png"))
|
||||
bg_frames = sorted(bg_frames_dir.glob("*.png"))
|
||||
|
||||
if not speaker_frames or not bg_frames:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Frame extraction produced no frames",
|
||||
)
|
||||
|
||||
frame_count = min(len(speaker_frames), len(bg_frames))
|
||||
log_interval = max(1, frame_count // 10)
|
||||
|
||||
# Step 4: Composite each frame pair
|
||||
for i in range(frame_count):
|
||||
if i % log_interval == 0:
|
||||
print(f"[green_screen_composite] Compositing frame {i + 1}/{frame_count}")
|
||||
|
||||
speaker_img = Image.open(speaker_frames[i]).convert("RGB")
|
||||
bg_img = Image.open(bg_frames[i]).convert("RGB")
|
||||
|
||||
comp = self._composite_frame(
|
||||
speaker_img, bg_img, bg_color,
|
||||
layout=layout,
|
||||
speaker_scale=speaker_scale,
|
||||
bg_shift_up=bg_shift_up,
|
||||
out_w=out_w,
|
||||
out_h=out_h,
|
||||
)
|
||||
comp.save(comp_frames_dir / f"frame_{i:06d}.png")
|
||||
|
||||
print(f"[green_screen_composite] All {frame_count} frames composited")
|
||||
|
||||
# Step 5: Encode composite frames to video
|
||||
no_audio_path = output_path if not original_audio_path else temp_dir / "no_audio.mp4"
|
||||
self._encode_frames(comp_frames_dir, no_audio_path, target_fps, out_w, out_h)
|
||||
|
||||
# Step 6: Mux audio if provided
|
||||
if original_audio_path:
|
||||
self._mux_audio(no_audio_path, Path(original_audio_path), output_path, duration)
|
||||
|
||||
if not output_path.exists() or output_path.stat().st_size == 0:
|
||||
return ToolResult(success=False, error="Output video was not created")
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"output": str(output_path),
|
||||
"layout": layout,
|
||||
"fps": target_fps,
|
||||
"frame_count": frame_count,
|
||||
"duration": round(duration, 2),
|
||||
"dimensions": f"{out_w}x{out_h}",
|
||||
"speaker_scale": speaker_scale,
|
||||
"has_audio": bool(original_audio_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Composite failed: {e}")
|
||||
finally:
|
||||
# Step 7: Clean up temp directories
|
||||
self._cleanup_temp(temp_dir)
|
||||
|
||||
def _parse_hex_color(self, hex_str: str) -> np.ndarray:
|
||||
"""Parse a hex color string like '#0E172A' to an RGB numpy array."""
|
||||
hex_str = hex_str.lstrip("#")
|
||||
r = int(hex_str[0:2], 16)
|
||||
g = int(hex_str[2:4], 16)
|
||||
b = int(hex_str[4:6], 16)
|
||||
return np.array([r, g, b])
|
||||
|
||||
def _probe_video(self, path: Path) -> dict[str, Any] | None:
|
||||
"""Probe a video for fps, duration, and dimensions."""
|
||||
cmd = [
|
||||
"ffprobe", "-v", "quiet",
|
||||
"-print_format", "json",
|
||||
"-show_format", "-show_streams",
|
||||
str(path),
|
||||
]
|
||||
try:
|
||||
result = self.run_command(cmd, timeout=30)
|
||||
data = json.loads(result.stdout)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
# Find video stream
|
||||
video_stream = None
|
||||
for stream in data.get("streams", []):
|
||||
if stream.get("codec_type") == "video":
|
||||
video_stream = stream
|
||||
break
|
||||
|
||||
if not video_stream:
|
||||
return None
|
||||
|
||||
# Parse fps from r_frame_rate (e.g., "30/1" or "15000/1001")
|
||||
fps_str = video_stream.get("r_frame_rate", "30/1")
|
||||
try:
|
||||
num, den = fps_str.split("/")
|
||||
fps = float(num) / float(den)
|
||||
except (ValueError, ZeroDivisionError):
|
||||
fps = 30.0
|
||||
|
||||
duration = float(data.get("format", {}).get("duration", 0))
|
||||
|
||||
return {
|
||||
"fps": fps,
|
||||
"duration": duration,
|
||||
"width": int(video_stream.get("width", 1920)),
|
||||
"height": int(video_stream.get("height", 1080)),
|
||||
}
|
||||
|
||||
def _extract_frames(self, video_path: Path, output_dir: Path, fps: float) -> None:
|
||||
"""Extract frames from a video at the given fps."""
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(video_path),
|
||||
"-vf", f"fps={fps}",
|
||||
str(output_dir / "frame_%06d.png"),
|
||||
]
|
||||
self.run_command(cmd, timeout=600)
|
||||
|
||||
def _composite_frame(
|
||||
self,
|
||||
speaker_img: Image.Image,
|
||||
bg_img: Image.Image,
|
||||
bg_color: np.ndarray,
|
||||
*,
|
||||
layout: str,
|
||||
speaker_scale: float,
|
||||
bg_shift_up: int,
|
||||
out_w: int,
|
||||
out_h: int,
|
||||
) -> Image.Image:
|
||||
"""Composite a single speaker frame over a background frame using the given layout."""
|
||||
# Create alpha mask from speaker frame
|
||||
speaker_arr = np.array(speaker_img).astype(float)
|
||||
dist = np.sqrt(np.sum((speaker_arr - bg_color.astype(float)) ** 2, axis=2))
|
||||
threshold = 35
|
||||
alpha = np.clip((dist - threshold) * 8, 0, 255).astype(np.uint8)
|
||||
|
||||
speaker_rgba = Image.new("RGBA", speaker_img.size)
|
||||
speaker_rgba.paste(speaker_img, (0, 0))
|
||||
speaker_rgba.putalpha(Image.fromarray(alpha))
|
||||
|
||||
# Prepare background canvas at output size
|
||||
canvas = Image.new("RGBA", (out_w, out_h), (0, 0, 0, 255))
|
||||
|
||||
if layout == "news_anchor":
|
||||
# Background shifted up so graphics appear above speaker's head
|
||||
bg_resized = bg_img.resize((out_w, out_h), Image.LANCZOS).convert("RGBA")
|
||||
# Shift background up: paste it higher so bottom content scrolls up
|
||||
shifted_bg = Image.new("RGBA", (out_w, out_h), (0, 0, 0, 255))
|
||||
shifted_bg.paste(bg_resized, (0, -bg_shift_up))
|
||||
canvas = shifted_bg
|
||||
|
||||
# Scale speaker and place at bottom center
|
||||
sp_w = int(speaker_rgba.width * speaker_scale)
|
||||
sp_h = int(speaker_rgba.height * speaker_scale)
|
||||
speaker_scaled = speaker_rgba.resize((sp_w, sp_h), Image.LANCZOS)
|
||||
x = (out_w - sp_w) // 2
|
||||
y = out_h - sp_h
|
||||
canvas.paste(speaker_scaled, (x, y), speaker_scaled)
|
||||
|
||||
elif layout == "full_behind":
|
||||
# Speaker full-frame on background, no scaling, no shifting
|
||||
bg_resized = bg_img.resize((out_w, out_h), Image.LANCZOS).convert("RGBA")
|
||||
canvas = bg_resized
|
||||
|
||||
# Resize speaker to match output
|
||||
speaker_full = speaker_rgba.resize((out_w, out_h), Image.LANCZOS)
|
||||
canvas.paste(speaker_full, (0, 0), speaker_full)
|
||||
|
||||
elif layout == "pip":
|
||||
# Background full-frame, speaker 30% in bottom-right
|
||||
bg_resized = bg_img.resize((out_w, out_h), Image.LANCZOS).convert("RGBA")
|
||||
canvas = bg_resized
|
||||
|
||||
pip_scale = 0.30
|
||||
sp_w = int(out_w * pip_scale)
|
||||
sp_h = int(out_h * pip_scale)
|
||||
speaker_pip = speaker_rgba.resize((sp_w, sp_h), Image.LANCZOS)
|
||||
margin = 20
|
||||
x = out_w - sp_w - margin
|
||||
y = out_h - sp_h - margin
|
||||
canvas.paste(speaker_pip, (x, y), speaker_pip)
|
||||
|
||||
elif layout == "split":
|
||||
# Speaker on left 50%, background on right 50%
|
||||
half_w = out_w // 2
|
||||
|
||||
# Left side: speaker resized to fill left half
|
||||
speaker_left = speaker_rgba.resize((half_w, out_h), Image.LANCZOS)
|
||||
# Right side: background cropped/resized to fill right half
|
||||
bg_right = bg_img.resize((half_w, out_h), Image.LANCZOS).convert("RGBA")
|
||||
|
||||
canvas.paste(speaker_left, (0, 0), speaker_left)
|
||||
canvas.paste(bg_right, (half_w, 0), bg_right)
|
||||
|
||||
# Convert to RGB for output
|
||||
return canvas.convert("RGB")
|
||||
|
||||
def _encode_frames(
|
||||
self, frames_dir: Path, output_path: Path, fps: float, width: int, height: int
|
||||
) -> None:
|
||||
"""Encode PNG frames to an MP4 video."""
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-framerate", str(fps),
|
||||
"-i", str(frames_dir / "frame_%06d.png"),
|
||||
"-c:v", "libx264", "-crf", "18", "-preset", "fast",
|
||||
"-pix_fmt", "yuv420p",
|
||||
"-vf", f"scale={width}:{height}",
|
||||
str(output_path),
|
||||
]
|
||||
self.run_command(cmd, timeout=600)
|
||||
|
||||
def _mux_audio(
|
||||
self, video_path: Path, audio_source: Path, output_path: Path, duration: float
|
||||
) -> None:
|
||||
"""Mux audio from the original source into the composite video."""
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(video_path),
|
||||
"-i", str(audio_source),
|
||||
"-t", f"{duration:.3f}",
|
||||
"-c:v", "copy",
|
||||
"-c:a", "aac", "-b:a", "192k",
|
||||
"-map", "0:v:0", "-map", "1:a:0",
|
||||
"-shortest",
|
||||
str(output_path),
|
||||
]
|
||||
self.run_command(cmd, timeout=300)
|
||||
|
||||
def _cleanup_temp(self, temp_dir: Path) -> None:
|
||||
"""Remove temporary frame directories."""
|
||||
if temp_dir.exists():
|
||||
try:
|
||||
shutil.rmtree(temp_dir)
|
||||
except OSError:
|
||||
# Best-effort cleanup; log but don't fail
|
||||
print(f"[green_screen_composite] Warning: could not fully clean {temp_dir}")
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
return 120.0
|
||||
@@ -0,0 +1,622 @@
|
||||
"""Green screen keying processor.
|
||||
|
||||
Removes green/blue screen backgrounds from footage using either FFmpeg
|
||||
chromakey filtering or rembg AI segmentation. Supports automatic method
|
||||
detection by analyzing frame color histograms.
|
||||
|
||||
Methods:
|
||||
- auto: Analyze frames to pick the best method (chromakey vs rembg)
|
||||
- chromakey: FFmpeg chromakey filter (fast, works well on clean screens)
|
||||
- rembg: AI background removal via rembg/u2net (slower, handles any bg)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import shutil
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ResumeSupport,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class GreenScreenProcessor(BaseTool):
|
||||
name = "green_screen_processor"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "video_post"
|
||||
provider = "ffmpeg"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["cmd:ffmpeg"]
|
||||
install_instructions = (
|
||||
"Install FFmpeg: https://ffmpeg.org/download.html "
|
||||
"For rembg method: pip install rembg[gpu] onnxruntime"
|
||||
)
|
||||
agent_skills = ["ffmpeg"]
|
||||
|
||||
capabilities = [
|
||||
"green_screen_keying",
|
||||
"chromakey",
|
||||
"background_removal",
|
||||
"rembg_segmentation",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["input_path", "output_path"],
|
||||
"properties": {
|
||||
"input_path": {
|
||||
"type": "string",
|
||||
"description": "Path to raw green screen footage",
|
||||
},
|
||||
"output_path": {
|
||||
"type": "string",
|
||||
"description": "Path for keyed output video",
|
||||
},
|
||||
"method": {
|
||||
"type": "string",
|
||||
"enum": ["auto", "chromakey", "rembg"],
|
||||
"default": "auto",
|
||||
"description": "Keying method: auto detects best approach, chromakey uses FFmpeg, rembg uses AI segmentation",
|
||||
},
|
||||
"fps": {
|
||||
"type": "integer",
|
||||
"default": 15,
|
||||
"description": "Output frames per second",
|
||||
},
|
||||
"bg_color": {
|
||||
"type": "string",
|
||||
"default": "#0E172A",
|
||||
"description": "Hex color for output background",
|
||||
},
|
||||
"max_frames": {
|
||||
"type": "integer",
|
||||
"default": 0,
|
||||
"description": "Limit frames to process (0 = all)",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=4, ram_mb=4096, vram_mb=0, disk_mb=8000, network_required=False
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=1, retryable_errors=["FFmpeg error"])
|
||||
resume_support = ResumeSupport.FROM_START
|
||||
idempotency_key_fields = [
|
||||
"input_path", "method", "fps", "bg_color", "max_frames",
|
||||
]
|
||||
side_effects = ["writes keyed video to output_path"]
|
||||
user_visible_verification = [
|
||||
"Check output for green fringing around subject edges",
|
||||
"Verify background is cleanly replaced with target color",
|
||||
"Look for flickering or inconsistent keying between frames",
|
||||
]
|
||||
|
||||
# Platform-specific null device
|
||||
_null_device = "NUL" if platform.system() == "Windows" else "/dev/null"
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
|
||||
output_path = Path(inputs["output_path"])
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
method = inputs.get("method", "auto")
|
||||
fps = inputs.get("fps", 15)
|
||||
bg_color = inputs.get("bg_color", "#0E172A")
|
||||
max_frames = inputs.get("max_frames", 0)
|
||||
start = time.time()
|
||||
|
||||
# Step 1: Probe input video
|
||||
probe = self._probe_video(input_path)
|
||||
if not probe:
|
||||
return ToolResult(success=False, error="Failed to probe input video")
|
||||
|
||||
duration = probe["duration"]
|
||||
width = probe["width"]
|
||||
height = probe["height"]
|
||||
src_fps = probe["fps"]
|
||||
|
||||
# Step 2: Determine method
|
||||
if method == "auto":
|
||||
method = self._auto_detect_method(input_path, duration, width, height)
|
||||
|
||||
# Step 3: Set up temp directory for frame processing
|
||||
temp_dir = output_path.parent / f".gs_tmp_{int(time.time())}"
|
||||
temp_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
# Step 4: Extract frames at target fps
|
||||
frames_dir = temp_dir / "frames"
|
||||
frames_dir.mkdir(exist_ok=True)
|
||||
frame_count = self._extract_frames(
|
||||
input_path, frames_dir, fps, max_frames
|
||||
)
|
||||
if frame_count == 0:
|
||||
return ToolResult(
|
||||
success=False, error="No frames extracted from input"
|
||||
)
|
||||
|
||||
# Step 5: Process frames
|
||||
processed_dir = temp_dir / "processed"
|
||||
processed_dir.mkdir(exist_ok=True)
|
||||
|
||||
if method == "chromakey":
|
||||
ok = self._process_chromakey(
|
||||
frames_dir, processed_dir, bg_color, frame_count
|
||||
)
|
||||
else:
|
||||
ok = self._process_rembg(
|
||||
frames_dir, processed_dir, bg_color, frame_count
|
||||
)
|
||||
|
||||
if not ok:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Frame processing failed with method={method}",
|
||||
)
|
||||
|
||||
# Step 6: Reconstruct video from processed frames
|
||||
self._reconstruct_video(processed_dir, output_path, fps, width, height)
|
||||
|
||||
if not output_path.exists() or output_path.stat().st_size == 0:
|
||||
return ToolResult(
|
||||
success=False, error="Output video was not created"
|
||||
)
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"method_used": method,
|
||||
"frame_count": frame_count,
|
||||
"duration": round(duration, 2),
|
||||
"output_path": str(output_path),
|
||||
"resolution": f"{width}x{height}",
|
||||
"fps": fps,
|
||||
"bg_color": bg_color,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Green screen processing failed: {e}")
|
||||
finally:
|
||||
# Clean up temp directory
|
||||
self._cleanup_dir(temp_dir)
|
||||
|
||||
def _probe_video(self, input_path: Path) -> dict[str, Any] | None:
|
||||
"""Probe video for duration, dimensions, and fps."""
|
||||
cmd = [
|
||||
"ffprobe", "-v", "quiet",
|
||||
"-show_entries", "format=duration:stream=width,height,r_frame_rate",
|
||||
"-select_streams", "v:0",
|
||||
"-of", "json",
|
||||
str(input_path),
|
||||
]
|
||||
try:
|
||||
result = self.run_command(cmd, timeout=30)
|
||||
data = json.loads(result.stdout)
|
||||
|
||||
duration = float(data.get("format", {}).get("duration", 0))
|
||||
|
||||
stream = data.get("streams", [{}])[0]
|
||||
width = int(stream.get("width", 0))
|
||||
height = int(stream.get("height", 0))
|
||||
|
||||
# Parse r_frame_rate like "30/1" or "30000/1001"
|
||||
fps_str = stream.get("r_frame_rate", "30/1")
|
||||
if "/" in fps_str:
|
||||
num, den = fps_str.split("/")
|
||||
fps_val = float(num) / float(den) if float(den) != 0 else 30.0
|
||||
else:
|
||||
fps_val = float(fps_str)
|
||||
|
||||
return {
|
||||
"duration": duration,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"fps": fps_val,
|
||||
}
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def _auto_detect_method(
|
||||
self, input_path: Path, duration: float, width: int, height: int
|
||||
) -> str:
|
||||
"""Analyze sample frames to decide between chromakey and rembg.
|
||||
|
||||
Extracts 5 evenly-spaced frames, checks color histograms for
|
||||
green/blue screen presence, then tests chromakey quality on a sample.
|
||||
"""
|
||||
temp_dir = input_path.parent / f".gs_detect_{int(time.time())}"
|
||||
temp_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
# Extract 5 sample frames evenly spaced
|
||||
interval = max(duration / 6, 0.1)
|
||||
sample_paths = []
|
||||
for i in range(5):
|
||||
ts = interval * (i + 1)
|
||||
out = temp_dir / f"sample_{i}.png"
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-ss", f"{ts:.3f}",
|
||||
"-i", str(input_path),
|
||||
"-frames:v", "1",
|
||||
str(out),
|
||||
]
|
||||
try:
|
||||
self.run_command(cmd, timeout=30)
|
||||
if out.exists():
|
||||
sample_paths.append(out)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
if not sample_paths:
|
||||
return "rembg" # fallback if we can't extract samples
|
||||
|
||||
# Analyze color histograms for green/blue screen presence
|
||||
has_green_screen = self._detect_green_screen_histogram(sample_paths)
|
||||
|
||||
if not has_green_screen:
|
||||
# No obvious green/blue screen detected, use rembg
|
||||
return "rembg"
|
||||
|
||||
# Test chromakey on a sample frame and check quality
|
||||
test_frame = sample_paths[len(sample_paths) // 2]
|
||||
chromakey_quality = self._test_chromakey_quality(test_frame, temp_dir)
|
||||
|
||||
if chromakey_quality > 80:
|
||||
return "chromakey"
|
||||
else:
|
||||
return "rembg"
|
||||
|
||||
finally:
|
||||
self._cleanup_dir(temp_dir)
|
||||
|
||||
def _detect_green_screen_histogram(self, sample_paths: list[Path]) -> bool:
|
||||
"""Analyze frames for dominant green or blue channel presence.
|
||||
|
||||
Uses FFmpeg signalstats to measure average hue. A strong green
|
||||
screen typically has a large area of similar green/blue hue.
|
||||
"""
|
||||
green_votes = 0
|
||||
for sample in sample_paths:
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(sample),
|
||||
"-vf", "signalstats=stat=tout+vrep+brng,metadata=mode=print",
|
||||
"-frames:v", "1",
|
||||
"-f", "null", self._null_device,
|
||||
]
|
||||
try:
|
||||
result = self.run_command(cmd, timeout=15)
|
||||
# Check stderr for color stats
|
||||
output = result.stderr or ""
|
||||
|
||||
# Alternative: use FFmpeg to count green-ish pixels
|
||||
# Run a simpler hue check with colorchannelmixer
|
||||
cmd2 = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(sample),
|
||||
"-vf", (
|
||||
"split[a][b];"
|
||||
"[a]colorchannelmixer=rr=0:gg=1:bb=0,"
|
||||
"threshold=threshold=0.3:similarity=0.3[mask];"
|
||||
"[mask]blackframe=amount=0:threshold=32"
|
||||
),
|
||||
"-frames:v", "1",
|
||||
"-f", "null", self._null_device,
|
||||
]
|
||||
# This is complex; use a simpler approach: check raw pixels
|
||||
# via a green-range filter
|
||||
cmd_green = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(sample),
|
||||
"-vf", (
|
||||
"colorkey=color=0x00FF00:similarity=0.4:blend=0.0,"
|
||||
"alphaextract,"
|
||||
"blackframe=amount=0:threshold=128"
|
||||
),
|
||||
"-frames:v", "1",
|
||||
"-f", "null", self._null_device,
|
||||
]
|
||||
try:
|
||||
result2 = self.run_command(cmd_green, timeout=15)
|
||||
stderr = result2.stderr or ""
|
||||
# blackframe reports percentage of black pixels
|
||||
# If many pixels became transparent (black in alpha), there's green
|
||||
if "pblack:" in stderr:
|
||||
import re
|
||||
pblack_matches = re.findall(r"pblack:(\d+)", stderr)
|
||||
if pblack_matches:
|
||||
pblack = int(pblack_matches[0])
|
||||
if pblack >= 20:
|
||||
green_votes += 1
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
# If majority of frames show green screen
|
||||
return green_votes >= len(sample_paths) // 2
|
||||
|
||||
def _test_chromakey_quality(self, test_frame: Path, temp_dir: Path) -> float:
|
||||
"""Run chromakey on a test frame and estimate quality percentage.
|
||||
|
||||
Returns a score 0-100 indicating what percentage of the expected
|
||||
background was successfully keyed out.
|
||||
"""
|
||||
keyed_out = temp_dir / "chromakey_test.png"
|
||||
|
||||
# Apply chromakey and output with alpha
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(test_frame),
|
||||
"-vf", "chromakey=color=0x00FF00:similarity=0.3:blend=0.08",
|
||||
str(keyed_out),
|
||||
]
|
||||
try:
|
||||
self.run_command(cmd, timeout=15)
|
||||
except Exception:
|
||||
return 0.0
|
||||
|
||||
if not keyed_out.exists():
|
||||
return 0.0
|
||||
|
||||
# Count transparent pixels via alphaextract + blackframe
|
||||
cmd2 = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(keyed_out),
|
||||
"-vf", "alphaextract,blackframe=amount=0:threshold=32",
|
||||
"-frames:v", "1",
|
||||
"-f", "null", self._null_device,
|
||||
]
|
||||
try:
|
||||
result = self.run_command(cmd2, timeout=15)
|
||||
stderr = result.stderr or ""
|
||||
import re
|
||||
pblack_matches = re.findall(r"pblack:(\d+)", stderr)
|
||||
if pblack_matches:
|
||||
# pblack = percentage of black pixels in alpha = transparent pixels
|
||||
return float(pblack_matches[0])
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return 0.0
|
||||
|
||||
def _extract_frames(
|
||||
self, input_path: Path, frames_dir: Path, fps: int, max_frames: int
|
||||
) -> int:
|
||||
"""Extract frames from video at target fps."""
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(input_path),
|
||||
"-vf", f"fps={fps}",
|
||||
str(frames_dir / "frame_%06d.png"),
|
||||
]
|
||||
|
||||
if max_frames > 0:
|
||||
cmd.insert(-1, "-frames:v")
|
||||
cmd.insert(-1, str(max_frames))
|
||||
|
||||
try:
|
||||
self.run_command(cmd, timeout=600)
|
||||
except Exception as e:
|
||||
# ffmpeg may return non-zero but still produce frames
|
||||
pass
|
||||
|
||||
# Count extracted frames
|
||||
frame_files = sorted(frames_dir.glob("frame_*.png"))
|
||||
count = len(frame_files)
|
||||
|
||||
if count > 0:
|
||||
# Log progress for large frame counts
|
||||
if count > 100:
|
||||
print(f"[green_screen_processor] Extracted {count} frames")
|
||||
|
||||
return count
|
||||
|
||||
def _process_chromakey(
|
||||
self,
|
||||
frames_dir: Path,
|
||||
processed_dir: Path,
|
||||
bg_color: str,
|
||||
frame_count: int,
|
||||
) -> bool:
|
||||
"""Process frames using FFmpeg chromakey filter.
|
||||
|
||||
Applies chromakey to remove green, then composites onto bg_color.
|
||||
"""
|
||||
bg_hex = bg_color.lstrip("#")
|
||||
# Convert hex to FFmpeg color format
|
||||
ffmpeg_bg = f"0x{bg_hex}"
|
||||
|
||||
frame_files = sorted(frames_dir.glob("frame_*.png"))
|
||||
processed = 0
|
||||
|
||||
for i, frame in enumerate(frame_files):
|
||||
out_path = processed_dir / frame.name
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-f", "lavfi", "-i", f"color=c={ffmpeg_bg}:size=1x1",
|
||||
"-i", str(frame),
|
||||
"-filter_complex",
|
||||
(
|
||||
f"[0:v]scale=iw:ih[bg];"
|
||||
f"[1:v]chromakey=color=0x00FF00:similarity=0.3:blend=0.08[fg];"
|
||||
f"[bg][fg]overlay=0:0"
|
||||
),
|
||||
"-frames:v", "1",
|
||||
str(out_path),
|
||||
]
|
||||
try:
|
||||
self.run_command(cmd, timeout=30)
|
||||
if out_path.exists():
|
||||
processed += 1
|
||||
except Exception:
|
||||
# Try with the frame size explicitly to fix scale
|
||||
try:
|
||||
cmd_retry = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(frame),
|
||||
"-vf",
|
||||
f"chromakey=color=0x00FF00:similarity=0.3:blend=0.08,"
|
||||
f"split[fg][alpha];"
|
||||
f"[alpha]alphaextract[a];"
|
||||
f"color=c={ffmpeg_bg}[bg];"
|
||||
f"[bg][fg][a]maskedmerge",
|
||||
"-frames:v", "1",
|
||||
str(out_path),
|
||||
]
|
||||
# Simpler fallback: just apply chromakey without compositing
|
||||
cmd_simple = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(frame),
|
||||
"-vf", f"chromakey=color=0x00FF00:similarity=0.3:blend=0.08",
|
||||
str(out_path),
|
||||
]
|
||||
self.run_command(cmd_simple, timeout=30)
|
||||
if out_path.exists():
|
||||
processed += 1
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
if frame_count > 100 and (i + 1) % 50 == 0:
|
||||
print(
|
||||
f"[green_screen_processor] Chromakey: {i + 1}/{frame_count} frames"
|
||||
)
|
||||
|
||||
return processed > 0
|
||||
|
||||
def _process_rembg(
|
||||
self,
|
||||
frames_dir: Path,
|
||||
processed_dir: Path,
|
||||
bg_color: str,
|
||||
frame_count: int,
|
||||
) -> bool:
|
||||
"""Process frames using rembg AI segmentation.
|
||||
|
||||
Removes background with u2net_human_seg model, then composites
|
||||
the subject onto bg_color background.
|
||||
"""
|
||||
try:
|
||||
import rembg
|
||||
from PIL import Image
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
# Parse bg_color hex to RGB
|
||||
bg_hex = bg_color.lstrip("#")
|
||||
bg_r = int(bg_hex[0:2], 16)
|
||||
bg_g = int(bg_hex[2:4], 16)
|
||||
bg_b = int(bg_hex[4:6], 16)
|
||||
|
||||
session = rembg.new_session("u2net_human_seg")
|
||||
|
||||
frame_files = sorted(frames_dir.glob("frame_*.png"))
|
||||
processed = 0
|
||||
|
||||
for i, frame in enumerate(frame_files):
|
||||
try:
|
||||
img = Image.open(frame).convert("RGB")
|
||||
import numpy as np
|
||||
|
||||
# Remove background (returns RGBA)
|
||||
result = rembg.remove(
|
||||
np.array(img),
|
||||
session=session,
|
||||
)
|
||||
result_img = Image.fromarray(result)
|
||||
|
||||
# Composite onto bg_color background
|
||||
bg = Image.new("RGBA", result_img.size, (bg_r, bg_g, bg_b, 255))
|
||||
bg.paste(result_img, (0, 0), result_img)
|
||||
|
||||
# Save as RGB
|
||||
out_path = processed_dir / frame.name
|
||||
bg.convert("RGB").save(out_path)
|
||||
processed += 1
|
||||
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
if frame_count > 100 and (i + 1) % 50 == 0:
|
||||
print(
|
||||
f"[green_screen_processor] rembg: {i + 1}/{frame_count} frames"
|
||||
)
|
||||
|
||||
return processed > 0
|
||||
|
||||
def _reconstruct_video(
|
||||
self,
|
||||
frames_dir: Path,
|
||||
output_path: Path,
|
||||
fps: int,
|
||||
width: int,
|
||||
height: int,
|
||||
) -> None:
|
||||
"""Reconstruct video from processed frames using FFmpeg."""
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-framerate", str(fps),
|
||||
"-i", str(frames_dir / "frame_%06d.png"),
|
||||
"-vf", f"scale={width}:{height}:flags=lanczos",
|
||||
"-c:v", "libx264",
|
||||
"-crf", "18",
|
||||
"-preset", "fast",
|
||||
"-pix_fmt", "yuv420p",
|
||||
str(output_path),
|
||||
]
|
||||
self.run_command(cmd, timeout=600)
|
||||
|
||||
@staticmethod
|
||||
def _cleanup_dir(dir_path: Path) -> None:
|
||||
"""Recursively remove a temp directory."""
|
||||
if not dir_path.exists():
|
||||
return
|
||||
try:
|
||||
shutil.rmtree(dir_path)
|
||||
except OSError:
|
||||
# Best-effort cleanup; individual file removal as fallback
|
||||
for f in dir_path.rglob("*"):
|
||||
try:
|
||||
if f.is_file():
|
||||
f.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
dir_path.rmdir()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
method = inputs.get("method", "auto")
|
||||
if method == "rembg":
|
||||
return 120.0
|
||||
elif method == "chromakey":
|
||||
return 30.0
|
||||
return 60.0 # auto
|
||||
@@ -106,6 +106,19 @@ class RemotionCaptionBurn(BaseTool):
|
||||
"correct replacement. Example: {\"cloud\": \"Claude\"}."
|
||||
),
|
||||
},
|
||||
"overlays": {
|
||||
"type": "array",
|
||||
"description": (
|
||||
"Array of overlay objects to render on top of the video. "
|
||||
"Each overlay has: type (text_card, stat_card, callout, "
|
||||
"comparison, bar_chart, line_chart, pie_chart, kpi_grid, "
|
||||
"hero_title, section_title, stat_reveal), in_seconds, "
|
||||
"out_seconds, position (lower_third, upper_third, "
|
||||
"left_panel, right_panel, full_overlay), and component-"
|
||||
"specific props (text, stat, chartData, etc.). "
|
||||
"See asset_manifest overlays from the asset-director."
|
||||
),
|
||||
},
|
||||
"force_ffmpeg": {
|
||||
"type": "boolean",
|
||||
"default": False,
|
||||
@@ -245,6 +258,7 @@ class RemotionCaptionBurn(BaseTool):
|
||||
words_per_page: int,
|
||||
font_size: int,
|
||||
highlight_color: str,
|
||||
overlays: list[dict] | None = None,
|
||||
) -> ToolResult:
|
||||
root = self._find_remotion_root()
|
||||
if root is None:
|
||||
@@ -262,6 +276,19 @@ class RemotionCaptionBurn(BaseTool):
|
||||
duration_s = float(dur_out.strip().split("\n")[0])
|
||||
total_frames = math.ceil(duration_s * 30)
|
||||
|
||||
# Detect video dimensions
|
||||
dim_cmd = [
|
||||
"ffprobe", "-v", "error",
|
||||
"-select_streams", "v:0",
|
||||
"-show_entries", "stream=width,height",
|
||||
"-of", "csv=p=0:s=x",
|
||||
input_path,
|
||||
]
|
||||
dim_result = self.run_command(dim_cmd)
|
||||
dim_parts = dim_result.stdout.strip().split("x")
|
||||
width = int(dim_parts[0])
|
||||
height = int(dim_parts[1])
|
||||
|
||||
# Copy video to Remotion public folder
|
||||
pub_dir = root / "public" / "talking-head"
|
||||
pub_dir.mkdir(parents=True, exist_ok=True)
|
||||
@@ -273,6 +300,7 @@ class RemotionCaptionBurn(BaseTool):
|
||||
props = {
|
||||
"videoSrc": f"public/talking-head/{video_filename}",
|
||||
"captions": captions,
|
||||
"overlays": overlays or [],
|
||||
"wordsPerPage": words_per_page,
|
||||
"fontSize": font_size,
|
||||
"highlightColor": highlight_color,
|
||||
@@ -287,12 +315,12 @@ class RemotionCaptionBurn(BaseTool):
|
||||
npx_bin = "npx.cmd" if sys.platform == "win32" else "npx"
|
||||
render_cmd = [
|
||||
npx_bin, "remotion", "render",
|
||||
"src/index.tsx", "TalkingHead",
|
||||
"TalkingHead",
|
||||
f"--props={props_file.relative_to(root)}",
|
||||
"--width=1080", "--height=1920", "--fps=30",
|
||||
f"--width={width}", f"--height={height}", "--fps=30",
|
||||
f"--frames=0-{total_frames - 1}",
|
||||
"--codec=h264", "--crf=18",
|
||||
str(Path(output_path).resolve()),
|
||||
f"--output={str(Path(output_path).resolve())}",
|
||||
]
|
||||
self.run_command(render_cmd, cwd=str(root))
|
||||
|
||||
@@ -307,6 +335,7 @@ class RemotionCaptionBurn(BaseTool):
|
||||
"duration_seconds": round(duration_s, 2),
|
||||
"total_frames": total_frames,
|
||||
"caption_count": len(captions),
|
||||
"overlay_count": len(overlays or []),
|
||||
"words_per_page": words_per_page,
|
||||
},
|
||||
artifacts=[output_path],
|
||||
@@ -429,11 +458,14 @@ class RemotionCaptionBurn(BaseTool):
|
||||
if not captions:
|
||||
return ToolResult(success=False, error="No caption words extracted.")
|
||||
|
||||
overlays = inputs.get("overlays")
|
||||
|
||||
# Choose render method
|
||||
if not force_ffmpeg and self._remotion_available():
|
||||
result = self._render_remotion(
|
||||
input_path, output_path, captions,
|
||||
words_per_page, font_size, highlight_color,
|
||||
overlays=overlays,
|
||||
)
|
||||
else:
|
||||
result = self._render_ffmpeg(input_path, output_path, captions)
|
||||
|
||||
Reference in New Issue
Block a user