implemented the lower priority features; haven't tested them
This commit is contained in:
@ -1,18 +1,17 @@
|
||||
"""
|
||||
AI background removal (Phase 5 - future).
|
||||
Uses MediaPipe or Robust Video Matting for person segmentation.
|
||||
Export-only -- no real-time preview.
|
||||
AI background removal using MediaPipe for person segmentation.
|
||||
Applied during export as a post-processing step — no real-time preview.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import subprocess
|
||||
import tempfile
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Placeholder for Phase 5 implementation
|
||||
# Will use mediapipe or rvm for segmentation at export time
|
||||
|
||||
MEDIAPIPE_AVAILABLE = False
|
||||
RVM_AVAILABLE = False
|
||||
|
||||
try:
|
||||
import mediapipe as mp
|
||||
@ -20,14 +19,9 @@ try:
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
pass # rvm import would go here
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
def is_available() -> bool:
|
||||
return MEDIAPIPE_AVAILABLE or RVM_AVAILABLE
|
||||
return MEDIAPIPE_AVAILABLE
|
||||
|
||||
|
||||
def remove_background_on_export(
|
||||
@ -37,23 +31,189 @@ def remove_background_on_export(
|
||||
replacement_value: str = "",
|
||||
) -> str:
|
||||
"""
|
||||
Process video frame-by-frame to remove/replace background.
|
||||
Only runs during export (not real-time).
|
||||
Process video frame-by-frame using FFmpeg chromakey fallback,
|
||||
or MediaPipe-based segmentation if available.
|
||||
|
||||
Args:
|
||||
input_path: source video
|
||||
output_path: destination
|
||||
replacement: 'blur', 'color', 'image', or 'video'
|
||||
replacement_value: hex color, image path, or video path
|
||||
replacement: 'blur', 'color', or 'image'
|
||||
replacement_value: hex color or image path (for color/image modes)
|
||||
|
||||
Returns:
|
||||
output_path
|
||||
"""
|
||||
if not is_available():
|
||||
raise RuntimeError(
|
||||
"Background removal requires mediapipe or robust-video-matting. "
|
||||
"Install with: pip install mediapipe"
|
||||
)
|
||||
input_path = str(Path(input_path).resolve())
|
||||
output_path = str(Path(output_path).resolve())
|
||||
|
||||
# Phase 5 implementation will go here
|
||||
raise NotImplementedError("Background removal is planned for Phase 5")
|
||||
if MEDIAPIPE_AVAILABLE:
|
||||
return _remove_with_mediapipe(input_path, output_path, replacement, replacement_value)
|
||||
else:
|
||||
return _remove_with_ffmpeg_portrait(input_path, output_path, replacement, replacement_value)
|
||||
|
||||
|
||||
def _remove_with_mediapipe(
|
||||
input_path: str,
|
||||
output_path: str,
|
||||
replacement: str = "blur",
|
||||
replacement_value: str = "",
|
||||
) -> str:
|
||||
"""Use MediaPipe Selfie Segmentation + FFmpeg for background removal.
|
||||
|
||||
Extracts frames, applies segmentation, composites replacement background.
|
||||
"""
|
||||
try:
|
||||
import cv2
|
||||
import numpy as np
|
||||
import mediapipe as mp
|
||||
|
||||
mp_selfie_segmentation = mp.solutions.selfie_segmentation
|
||||
|
||||
# Determine background color/image
|
||||
if replacement == "color":
|
||||
color_hex = replacement_value or "#00FF00"
|
||||
color_hex = color_hex.lstrip("#")
|
||||
bg_color = tuple(int(color_hex[i:i+2], 16) for i in (0, 2, 4))
|
||||
bg_color = bg_color[::-1] # RGB -> BGR
|
||||
elif replacement == "image":
|
||||
bg_image = cv2.imread(replacement_value) if replacement_value else None
|
||||
if bg_image is None:
|
||||
bg_color = (0, 255, 0)
|
||||
bg_image = None
|
||||
else:
|
||||
# Blur background (default)
|
||||
bg_color = None
|
||||
|
||||
# Open video
|
||||
cap = cv2.VideoCapture(input_path)
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
|
||||
# Temp directory for processed frames
|
||||
temp_dir = tempfile.mkdtemp(prefix="aive_bgrem_")
|
||||
frame_dir = os.path.join(temp_dir, "frames")
|
||||
os.makedirs(frame_dir, exist_ok=True)
|
||||
|
||||
with mp_selfie_segmentation.SelfieSegmentation(model_selection=0) as segmenter:
|
||||
frame_idx = 0
|
||||
while cap.isOpened():
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
|
||||
# Convert to RGB for MediaPipe
|
||||
rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
result = segmenter.process(rgb)
|
||||
mask = result.segmentation_mask
|
||||
|
||||
# Threshold the mask
|
||||
condition = mask > 0.5
|
||||
|
||||
if replacement == "blur":
|
||||
# Apply strong blur to background
|
||||
blurred = cv2.GaussianBlur(frame, (99, 99), 0)
|
||||
output_frame = np.where(condition[..., None], frame, blurred)
|
||||
elif replacement == "color":
|
||||
bg = np.full(frame.shape, bg_color, dtype=np.uint8)
|
||||
output_frame = np.where(condition[..., None], frame, bg)
|
||||
elif replacement == "image" and bg_image is not None:
|
||||
bg_resized = cv2.resize(bg_image, (width, height))
|
||||
output_frame = np.where(condition[..., None], frame, bg_resized)
|
||||
else:
|
||||
output_frame = frame
|
||||
|
||||
out_path = os.path.join(frame_dir, f"frame_{frame_idx:06d}.png")
|
||||
cv2.imwrite(out_path, output_frame)
|
||||
frame_idx += 1
|
||||
|
||||
if frame_idx % 100 == 0:
|
||||
logger.info("Background removal: %d/%d frames", frame_idx, total_frames)
|
||||
|
||||
cap.release()
|
||||
|
||||
# Encode frames back to video using FFmpeg
|
||||
import subprocess as _sp
|
||||
ffmpeg = "ffmpeg"
|
||||
cmd = [
|
||||
ffmpeg, "-y",
|
||||
"-framerate", str(fps),
|
||||
"-i", os.path.join(frame_dir, "frame_%06d.png"),
|
||||
"-i", input_path,
|
||||
"-map", "0:v:0",
|
||||
"-map", "1:a:0?",
|
||||
"-c:v", "libx264", "-preset", "medium", "-crf", "18",
|
||||
"-c:a", "aac", "-b:a", "192k",
|
||||
"-shortest",
|
||||
"-pix_fmt", "yuv420p",
|
||||
output_path,
|
||||
]
|
||||
result = _sp.run(cmd, capture_output=True, text=True)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"FFmpeg frame encode failed: {result.stderr[-500:]}")
|
||||
|
||||
# Cleanup
|
||||
for f in os.listdir(frame_dir):
|
||||
try:
|
||||
os.remove(os.path.join(frame_dir, f))
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
os.rmdir(frame_dir)
|
||||
os.rmdir(temp_dir)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
logger.info("MediaPipe background removal completed -> %s", output_path)
|
||||
return output_path
|
||||
|
||||
except ImportError:
|
||||
logger.warning("mediapipe/cv2 not available, falling back to FFmpeg portrait mode")
|
||||
return _remove_with_ffmpeg_portrait(input_path, output_path, replacement, replacement_value)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"MediaPipe background removal failed: {e}")
|
||||
|
||||
|
||||
def _remove_with_ffmpeg_portrait(
|
||||
input_path: str,
|
||||
output_path: str,
|
||||
replacement: str = "blur",
|
||||
replacement_value: str = "",
|
||||
) -> str:
|
||||
"""Fallback: use FFmpeg's colorkey + chromakey for basic background removal.
|
||||
|
||||
This is a crude approximation. For best results, install mediapipe + opencv-python.
|
||||
"""
|
||||
ffmpeg = "ffmpeg"
|
||||
|
||||
# Use a simple chromakey-based approach with a neutral background
|
||||
# This won't work well for most real videos but provides a fallback
|
||||
if replacement == "color":
|
||||
color = replacement_value or "00FF00"
|
||||
filter_complex = f"colorkey=0x{color}:0.3:0.1,chromakey=0x{color}:0.3:0.1"
|
||||
elif replacement == "blur":
|
||||
filter_complex = "gblur=sigma=20:enable='gt(scene,0.01)'"
|
||||
else:
|
||||
filter_complex = "null"
|
||||
|
||||
if filter_complex == "null":
|
||||
# No-op, copy input to output
|
||||
cmd = [ffmpeg, "-y", "-i", input_path, "-c", "copy", output_path]
|
||||
else:
|
||||
cmd = [
|
||||
ffmpeg, "-y",
|
||||
"-i", input_path,
|
||||
"-vf", filter_complex,
|
||||
"-c:v", "libx264", "-preset", "medium", "-crf", "18",
|
||||
"-c:a", "aac", "-b:a", "192k",
|
||||
"-movflags", "+faststart",
|
||||
output_path,
|
||||
]
|
||||
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"FFmpeg background removal failed: {result.stderr[-500:]}")
|
||||
|
||||
logger.info("FFmpeg portait background removal completed -> %s", output_path)
|
||||
return output_path
|
||||
|
||||
Reference in New Issue
Block a user