feat(rembg): add RemoveBG Studio - native background removal app
C++ WebView2 host app driving a Python/rembg engine as a JSON-stdio sidecar, with live before/after preview for images and video frames. Also includes the original CLI (removebg.py) built on the same core/ engine. The WebView2 SDK is fetched on demand via native/third_party/fetch-webview2.ps1 (wired into the Makefile) rather than vendored, since it's ~15MB of prebuilt/generated Microsoft SDK content. A Makefile provides make sync/build/run/clean as the standard entry points.
This commit is contained in:
@@ -0,0 +1,239 @@
|
||||
import os
|
||||
import re
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
from typing import Dict, Any, Optional, Callable
|
||||
|
||||
def check_ffmpeg_installed() -> bool:
|
||||
"""Check if ffmpeg is available on the system PATH."""
|
||||
return shutil.which("ffmpeg") is not None
|
||||
|
||||
def check_ffprobe_installed() -> bool:
|
||||
"""Check if ffprobe is available on the system PATH."""
|
||||
return shutil.which("ffprobe") is not None
|
||||
|
||||
def get_video_info(video_path: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Get video metadata using ffprobe or ffmpeg.
|
||||
Returns dictionary with fps, width, height, duration, frame_count, has_audio.
|
||||
"""
|
||||
info = {
|
||||
"fps": 30.0,
|
||||
"width": 1920,
|
||||
"height": 1080,
|
||||
"duration": 0.0,
|
||||
"frame_count": 0,
|
||||
"has_audio": False
|
||||
}
|
||||
|
||||
if check_ffprobe_installed():
|
||||
cmd = [
|
||||
"ffprobe",
|
||||
"-v", "quiet",
|
||||
"-print_format", "json",
|
||||
"-show_format",
|
||||
"-show_streams",
|
||||
video_path
|
||||
]
|
||||
try:
|
||||
result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True)
|
||||
data = json.loads(result.stdout)
|
||||
|
||||
for stream in data.get("streams", []):
|
||||
if stream.get("codec_type") == "video" and "r_frame_rate" in stream:
|
||||
# Evaluate fps string like "30/1" or "30000/1001"
|
||||
r_fps = stream["r_frame_rate"]
|
||||
if "/" in r_fps:
|
||||
num, den = r_fps.split("/")
|
||||
if float(den) > 0:
|
||||
info["fps"] = float(num) / float(den)
|
||||
else:
|
||||
info["fps"] = float(r_fps)
|
||||
|
||||
info["width"] = int(stream.get("width", 1920))
|
||||
info["height"] = int(stream.get("height", 1080))
|
||||
|
||||
if "nb_frames" in stream and stream["nb_frames"].isdigit():
|
||||
info["frame_count"] = int(stream["nb_frames"])
|
||||
elif stream.get("codec_type") == "audio":
|
||||
info["has_audio"] = True
|
||||
|
||||
if "format" in data and "duration" in data["format"]:
|
||||
info["duration"] = float(data["format"]["duration"])
|
||||
if info["frame_count"] == 0 and info["fps"] > 0:
|
||||
info["frame_count"] = int(info["duration"] * info["fps"])
|
||||
return info
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Fallback parsing via ffmpeg stderr
|
||||
cmd = ["ffmpeg", "-i", video_path]
|
||||
res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||||
stderr = res.stderr
|
||||
|
||||
fps_match = re.search(r"(\d+(?:\.\d+)?)\s+fps", stderr)
|
||||
if fps_match:
|
||||
info["fps"] = float(fps_match.group(1))
|
||||
|
||||
res_match = re.search(r"(\d{2,5})x(\d{2,5})", stderr)
|
||||
if res_match:
|
||||
info["width"] = int(res_match.group(1))
|
||||
info["height"] = int(res_match.group(2))
|
||||
|
||||
dur_match = re.search(r"Duration:\s*(\d+):(\d+):(\d+(?:\.\d+)?)", stderr)
|
||||
if dur_match:
|
||||
h, m, s = float(dur_match.group(1)), float(dur_match.group(2)), float(dur_match.group(3))
|
||||
info["duration"] = h * 3600 + m * 60 + s
|
||||
info["frame_count"] = int(info["duration"] * info["fps"])
|
||||
|
||||
if "Audio:" in stderr:
|
||||
info["has_audio"] = True
|
||||
|
||||
return info
|
||||
|
||||
def extract_frames_lossless(
|
||||
video_path: str,
|
||||
output_dir: str,
|
||||
target_fps: Optional[float] = None,
|
||||
progress_callback: Optional[Callable[[str], None]] = None
|
||||
) -> int:
|
||||
"""
|
||||
Extract frames from a video into PNG images (100% Lossless) inside output_dir.
|
||||
Returns total number of extracted frames.
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
out_pattern = os.path.join(output_dir, "frame_%06d.png")
|
||||
|
||||
cmd = ["ffmpeg", "-y", "-i", video_path]
|
||||
|
||||
if target_fps and target_fps > 0:
|
||||
cmd.extend(["-vf", f"fps={target_fps}"])
|
||||
|
||||
# Lossless PNG output codec (-c:v png -pred mixed)
|
||||
cmd.extend(["-c:v", "png", "-pred", "mixed", out_pattern])
|
||||
|
||||
if progress_callback:
|
||||
progress_callback("Extracting video frames in lossless PNG format...")
|
||||
|
||||
process = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||||
if process.returncode != 0:
|
||||
raise RuntimeError(f"FFmpeg frame extraction failed: {process.stderr}")
|
||||
|
||||
frames = [f for f in os.listdir(output_dir) if f.startswith("frame_") and f.endswith(".png")]
|
||||
frames.sort()
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(f"Successfully extracted {len(frames)} lossless frames.")
|
||||
|
||||
return len(frames)
|
||||
|
||||
def extract_single_frame(video_path: str, output_path: str, time_offset: float = 0.5) -> str:
|
||||
"""Grab one representative frame from a video (for preview purposes)."""
|
||||
cmd = ["ffmpeg", "-y", "-ss", str(time_offset), "-i", video_path, "-frames:v", "1", output_path]
|
||||
res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||||
if res.returncode != 0 or not os.path.exists(output_path):
|
||||
raise RuntimeError(f"FFmpeg single-frame extraction failed: {res.stderr}")
|
||||
return output_path
|
||||
|
||||
def extract_audio(video_path: str, output_audio_path: str) -> bool:
|
||||
"""Extract audio stream from video file if present."""
|
||||
cmd = ["ffmpeg", "-y", "-i", video_path, "-vn", "-acodec", "aac", "-b:a", "192k", output_audio_path]
|
||||
res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||||
return res.returncode == 0 and os.path.exists(output_audio_path) and os.path.getsize(output_audio_path) > 0
|
||||
|
||||
def reassemble_video(
|
||||
frames_dir: str,
|
||||
output_video_path: str,
|
||||
fps: float = 30.0,
|
||||
audio_path: Optional[str] = None,
|
||||
export_format: str = "mov_transparent",
|
||||
bg_color: Optional[str] = None,
|
||||
progress_callback: Optional[Callable[[str], None]] = None
|
||||
) -> str:
|
||||
"""
|
||||
Reassemble processed frame sequence back into a video file or frame folder.
|
||||
|
||||
export_format options:
|
||||
- 'mov_transparent': QuickTime MOV with ProRes 4444 + Alpha Transparency
|
||||
- 'webm_transparent': WebM with VP9 + Alpha Transparency
|
||||
- 'mp4_solid': MP4 (H.264) with solid background color (e.g. Green screen #00FF00)
|
||||
- 'png_sequence': Copies/Leaves processed PNG frames in output directory
|
||||
"""
|
||||
if export_format == "png_sequence":
|
||||
out_dir = output_video_path
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
for fname in os.listdir(frames_dir):
|
||||
if fname.endswith(".png"):
|
||||
shutil.copy2(os.path.join(frames_dir, fname), os.path.join(out_dir, fname))
|
||||
if progress_callback:
|
||||
progress_callback(f"PNG sequence exported to {out_dir}")
|
||||
return out_dir
|
||||
|
||||
input_pattern = os.path.join(frames_dir, "frame_%06d.png")
|
||||
|
||||
# Ensure parent output directory exists
|
||||
os.makedirs(os.path.dirname(os.path.abspath(output_video_path)), exist_ok=True)
|
||||
|
||||
cmd = ["ffmpeg", "-y", "-framerate", str(fps), "-i", input_pattern]
|
||||
|
||||
if audio_path and os.path.exists(audio_path):
|
||||
cmd.extend(["-i", audio_path, "-c:a", "aac", "-shortest"])
|
||||
|
||||
if export_format == "mov_transparent":
|
||||
# ProRes 4444 with Alpha Channel (Lossless/visually lossless transparent video)
|
||||
cmd.extend([
|
||||
"-c:v", "prores_ks",
|
||||
"-profile:v", "4",
|
||||
"-pix_fmt", "yuva444p10le",
|
||||
output_video_path
|
||||
])
|
||||
elif export_format == "webm_transparent":
|
||||
# VP9 WebM with Alpha Channel
|
||||
cmd.extend([
|
||||
"-c:v", "libvpx-vp9",
|
||||
"-pix_fmt", "yuva420p",
|
||||
"-b:v", "4M",
|
||||
output_video_path
|
||||
])
|
||||
elif export_format == "mp4_solid":
|
||||
# Standard MP4 H.264
|
||||
cmd.extend([
|
||||
"-c:v", "libx264",
|
||||
"-pix_fmt", "yuv420p",
|
||||
"-crf", "18",
|
||||
output_video_path
|
||||
])
|
||||
else:
|
||||
# Default MP4 fallback
|
||||
cmd.extend([
|
||||
"-c:v", "libx264",
|
||||
"-pix_fmt", "yuv420p",
|
||||
output_video_path
|
||||
])
|
||||
|
||||
if progress_callback:
|
||||
progress_callback("Reassembling frames into output video format...")
|
||||
|
||||
res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||||
if res.returncode != 0:
|
||||
# Fallback to standard libx264 if prores_ks or libvpx-vp9 isn't available
|
||||
if export_format in ["mov_transparent", "webm_transparent"]:
|
||||
if progress_callback:
|
||||
progress_callback(f"Codec warning: trying fallback MP4 export for {output_video_path}...")
|
||||
fallback_cmd = [
|
||||
"ffmpeg", "-y", "-framerate", str(fps), "-i", input_pattern
|
||||
]
|
||||
if audio_path and os.path.exists(audio_path):
|
||||
fallback_cmd.extend(["-i", audio_path, "-c:a", "aac", "-shortest"])
|
||||
fallback_cmd.extend(["-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "18", output_video_path])
|
||||
res_fb = subprocess.run(fallback_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||||
if res_fb.returncode != 0:
|
||||
raise RuntimeError(f"FFmpeg video reassembly failed: {res_fb.stderr}")
|
||||
else:
|
||||
raise RuntimeError(f"FFmpeg video reassembly failed: {res.stderr}")
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(f"Video created successfully: {output_video_path}")
|
||||
|
||||
return output_video_path
|
||||
Reference in New Issue
Block a user