""" 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__) MEDIAPIPE_AVAILABLE = False try: import mediapipe as mp MEDIAPIPE_AVAILABLE = True except ImportError: pass def is_available() -> bool: return MEDIAPIPE_AVAILABLE def remove_background_on_export( input_path: str, output_path: str, replacement: str = "blur", replacement_value: str = "", ) -> str: """ 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', or 'image' replacement_value: hex color or image path (for color/image modes) Returns: output_path """ input_path = str(Path(input_path).resolve()) output_path = str(Path(output_path).resolve()) 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