LyAhn
fa0170daa9
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.
240 lines
9.2 KiB
Python
240 lines
9.2 KiB
Python
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
|