import os import tempfile import shutil from typing import List, Optional, Callable, Dict, Tuple from PIL import Image, ImageColor import rembg from core import ffmpeg_utils # Available model descriptions AVAILABLE_MODELS = { "u2net": "U2-Net (Default - General Purpose)", "u2netp": "U2-Net Lite (Fast, Lightweight)", "u2net_human_seg": "U2-Net Human Seg (Optimized for People & Portraits)", "u2net_cloth_seg": "U2-Net Cloth Seg (Optimized for Clothes/Fashion)", "isnet-general-use": "IS-Net General Use (High Accuracy & Crisp Edges)", "birefnet-general": "BiRefNet General (State-of-the-Art High Precision)" } class BackgroundRemoverEngine: def __init__(self, model_name: str = "u2net"): self.model_name = model_name self.session = None self._init_session(model_name) def _init_session(self, model_name: str): self.model_name = model_name try: self.session = rembg.new_session(model_name) except Exception as e: # Fallback to standard u2net if specific model fails print(f"Warning: Failed to load model {model_name}: {e}. Falling back to u2net.") self.model_name = "u2net" self.session = rembg.new_session("u2net") def change_model(self, model_name: str): if model_name != self.model_name: self._init_session(model_name) def remove_background( self, image_input: Image.Image, alpha_matting: bool = False, af: int = 240, ab: int = 10, ae: int = 10, bg_color: Optional[str] = None ) -> Image.Image: """ Process PIL Image and return PIL Image with background removed or replaced. """ kwargs = {} if alpha_matting: kwargs.update({ "alpha_matting": True, "alpha_matting_foreground_threshold": af, "alpha_matting_background_threshold": ab, "alpha_matting_erode_size": ae }) output_img = rembg.remove(image_input, session=self.session, **kwargs) # Ensure RGBA mode if output_img.mode != "RGBA": output_img = output_img.convert("RGBA") # If solid background color requested (e.g. Green screen "#00FF00") if bg_color and bg_color.strip() and bg_color.lower() != "transparent": try: rgb = ImageColor.getrgb(bg_color) background = Image.new("RGBA", output_img.size, rgb + (255,)) background.paste(output_img, (0, 0), output_img) return background.convert("RGB") except Exception as e: print(f"Warning: Invalid background color '{bg_color}': {e}") return output_img def process_single_image( self, input_path: str, output_path: str, alpha_matting: bool = False, bg_color: Optional[str] = None ) -> str: """Process a single image file.""" if not os.path.exists(input_path): raise FileNotFoundError(f"Input file not found: {input_path}") img = Image.open(input_path) processed = self.remove_background( img, alpha_matting=alpha_matting, bg_color=bg_color ) os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) # Save based on file extension ext = os.path.splitext(output_path)[1].lower() if ext in [".jpg", ".jpeg"] and processed.mode == "RGBA": # Convert RGBA to RGB with white background for JPEG bg = Image.new("RGB", processed.size, (255, 255, 255)) bg.paste(processed, (0, 0), processed) bg.save(output_path, quality=95) else: processed.save(output_path) return output_path def process_batch_images( self, input_paths: List[str], output_dir: str, alpha_matting: bool = False, bg_color: Optional[str] = None, output_format: str = "png", progress_callback: Optional[Callable[[int, int, str], None]] = None ) -> List[str]: """Process a list of image paths or directory.""" os.makedirs(output_dir, exist_ok=True) results = [] total = len(input_paths) for i, in_path in enumerate(input_paths, 1): fname = os.path.splitext(os.path.basename(in_path))[0] out_ext = ".png" if output_format.lower() == "png" else f".{output_format.lower()}" out_path = os.path.join(output_dir, f"{fname}_nobg{out_ext}") if progress_callback: progress_callback(i, total, os.path.basename(in_path)) res = self.process_single_image( in_path, out_path, alpha_matting=alpha_matting, bg_color=bg_color ) results.append(res) return results def process_video( self, video_path: str, output_path: str, fps: Optional[float] = None, export_format: str = "mov_transparent", bg_color: Optional[str] = None, alpha_matting: bool = False, progress_callback: Optional[Callable[[int, int, str], None]] = None ) -> str: """ Extract video frames in lossless PNG format via ffmpeg, process frames with rembg, and reassemble into output video or PNG frame sequence. """ if not ffmpeg_utils.check_ffmpeg_installed(): raise RuntimeError("FFmpeg is not installed or not found on system PATH.") if not os.path.exists(video_path): raise FileNotFoundError(f"Video file not found: {video_path}") # Get video metadata vinfo = ffmpeg_utils.get_video_info(video_path) actual_fps = fps if (fps and fps > 0) else vinfo.get("fps", 30.0) temp_dir = tempfile.mkdtemp(prefix="rembg_video_") extracted_frames_dir = os.path.join(temp_dir, "extracted") processed_frames_dir = os.path.join(temp_dir, "processed") audio_file = os.path.join(temp_dir, "audio.aac") try: # 1. Extract frames losslessly if progress_callback: progress_callback(0, 100, "Extracting frames losslessly with FFmpeg...") num_frames = ffmpeg_utils.extract_frames_lossless( video_path, extracted_frames_dir, target_fps=actual_fps ) # 2. Extract audio track if present has_audio = False if vinfo.get("has_audio", False): if progress_callback: progress_callback(0, 100, "Extracting audio track...") has_audio = ffmpeg_utils.extract_audio(video_path, audio_file) # 3. Process frames with rembg os.makedirs(processed_frames_dir, exist_ok=True) frame_files = sorted([f for f in os.listdir(extracted_frames_dir) if f.endswith(".png")]) for idx, frame_name in enumerate(frame_files, 1): if progress_callback: progress_callback(idx, len(frame_files), f"Removing background from frame {idx}/{len(frame_files)}") in_f = os.path.join(extracted_frames_dir, frame_name) out_f = os.path.join(processed_frames_dir, frame_name) self.process_single_image( in_f, out_f, alpha_matting=alpha_matting, bg_color=bg_color ) # 4. Reassemble output video / PNG sequence if progress_callback: progress_callback(len(frame_files), len(frame_files), "Reassembling video...") final_output = ffmpeg_utils.reassemble_video( frames_dir=processed_frames_dir, output_video_path=output_path, fps=actual_fps, audio_path=audio_file if has_audio else None, export_format=export_format, bg_color=bg_color ) return final_output finally: # Cleanup temp directory shutil.rmtree(temp_dir, ignore_errors=True)