Create app.py

app.py

@@ -0,0 +1,134 @@
+import os
+import cv2
+import gradio as gr
+import torch
+import numpy as np
+import onnxruntime as ort
+import requests
+import time
+import ffmpeg
+import asyncio
+import shutil
+from rembg import remove, new_session
+from concurrent.futures import ThreadPoolExecutor
+
+# Ensure model directory exists
+MODEL_DIR = os.path.expanduser("~/.u2net/")
+MODEL_PATH = os.path.join(MODEL_DIR, "u2net.onnx")
+CACHE_MODEL_PATH = "cached_model.onnx"
+
+if not os.path.exists(CACHE_MODEL_PATH):
+    print("Checking model cache...")
+    if not os.path.exists(MODEL_PATH):
+        print("Downloading U2Net model...")
+        os.makedirs(MODEL_DIR, exist_ok=True)
+        response = requests.get("https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx", stream=True)
+        with open(MODEL_PATH, "wb") as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+        print("Download complete!")
+    shutil.copy(MODEL_PATH, CACHE_MODEL_PATH)
+    print("Model cached for future use.")
+else:
+    print("Cached model found. Skipping download.")
+
+# Enable GPU Execution if available
+execution_providers = ["CUDAExecutionProvider"] if torch.cuda.is_available() else ["CPUExecutionProvider"]
+
+# Lazy-load rembg session to avoid blocking
+session = None
+def get_rembg_session():
+    global session
+    if session is None:
+        session = new_session("u2net", providers=execution_providers)
+    return session
+
+print("Model successfully loaded and ready!")
+
+def extract_audio(video_path, audio_path, progress_callback):
+    progress_callback.update(0.05, "🔴 Extracting Audio")
+    ffmpeg.input(video_path).output(audio_path, codec="aac").run(overwrite_output=True, quiet=True)
+    progress_callback.update(0.15, "🔴 Audio Extracted")
+
+def merge_audio(video_path, audio_path, output_path, progress_callback):
+    progress_callback.update(0.85, "🟢 Preparing to Merge Audio")
+    temp_video = "temp_video.mp4"
+    ffmpeg.input(video_path).output(temp_video, vcodec="copy", acodec="aac").run(overwrite_output=True, quiet=True)
+    progress_callback.update(0.90, "🟢 Merging Audio & Finalizing")
+    ffmpeg.input(temp_video).input(audio_path).output(
+        output_path, vcodec="copy", acodec="aac", shortest=True, **{"map": "0:v:0", "map": "1:a:0"}
+    ).run(overwrite_output=True, quiet=False)
+    progress_callback.update(1.00, "✅ Processing Complete")
+
+def process_video(video_file, background_file, progress=gr.Progress(track_tqdm=True)):
+    if not video_file:
+        return "Error: No video file provided."
+    if not background_file:
+        return "Error: No background image provided."
+    
+    foreground_path = video_file.name
+    background_path = background_file.name
+    temp_video_path = "temp_output.mp4"
+    final_output_path = "output_with_audio.mp4"
+    audio_path = "temp_audio.aac"
+    
+    cap_fg = cv2.VideoCapture(foreground_path)
+    if not cap_fg.isOpened():
+        return "Error: Unable to open foreground video."
+    
+    progress.update(0.05, "🔴 Extracting Audio")
+    extract_audio(foreground_path, audio_path, progress)
+    
+    background_img = cv2.imread(background_path)
+    if background_img is None:
+        return "Error: Unable to open background image."
+    
+    frame_width = int(cap_fg.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap_fg.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = int(cap_fg.get(cv2.CAP_PROP_FPS))
+    total_frames = int(cap_fg.get(cv2.CAP_PROP_FRAME_COUNT))
+    
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    out = cv2.VideoWriter(temp_video_path, fourcc, fps, (frame_width, frame_height))
+    
+    background_img = cv2.resize(background_img, (frame_width, frame_height))
+    session = get_rembg_session()
+    
+    for frame_idx in range(total_frames):
+        ret_fg, frame_fg = cap_fg.read()
+        if not ret_fg:
+            break
+        
+        progress.update((frame_idx + 1) / total_frames, f"🔵 Processing Frame {frame_idx + 1}/{total_frames}")
+        
+        frame_no_bg = remove(frame_fg, session=session)
+        frame_no_bg = cv2.cvtColor(frame_no_bg, cv2.COLOR_RGBA2BGRA)
+        
+        alpha_channel = frame_no_bg[:, :, 3]
+        mask_inv = cv2.bitwise_not(alpha_channel)
+        
+        fg_part = cv2.bitwise_and(frame_fg, frame_fg, mask=alpha_channel)
+        bg_part = cv2.bitwise_and(background_img, background_img, mask=mask_inv)
+        final_frame = cv2.add(bg_part, fg_part)
+        out.write(final_frame)
+    
+    cap_fg.release()
+    out.release()
+    
+    progress.update(0.85, "🟢 Merging Audio")
+    merge_audio(temp_video_path, audio_path, final_output_path, progress)
+    progress.update(1.00, "✅ Processing Complete")
+    
+    return final_output_path
+
+iface = gr.Interface(
+    fn=process_video,
+    inputs=[gr.File(type="filepath"), gr.File(type="filepath")],
+    outputs=gr.File(type="filepath"),
+    title="🎬 AI Video Background Remover - Optimized",
+    description="Upload a video and a background image. The AI replaces the background while keeping the audio.",
+    live=False
+)
+
+if __name__ == "__main__":
+    iface.launch()