Construa um Conversor Local de Vídeo de Profundidade com Gradio

import os import sys import tempfile import subprocess import numpy as np import cv2 import torch import gradio as gr from PIL import Image from pathlib import Path import shutil import json import time from typing import Optional, Tuple # Check for available acceleration def get_device(): """Determine the best available device for inference.""" if torch.cuda.is_available(): return "cuda" elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): return "mps" else: return "cpu" DEVICE = get_device() print(f"Using device: {DEVICE}") # Try to import Depth Anything V2 try: from depth_anything_v2.dpt import DepthAnythingV2 DEPTH_ANYTHING_AVAILABLE = True except ImportError: DEPTH_ANYTHING_AVAILABLE = False print("Depth Anything V2 not found. Will download on first use.") # Model configurations MODEL_CONFIGS = { "vits": {"encoder": "vits", "features": 64, "out_channels": [48, 96, 192, 384]}, "vitb": {"encoder": "vitb", "features": 128, "out_channels": [96, 192, 384, 768]}, "vitl": {"encoder": "vitl", "features": 256, "out_channels": [256, 512, 1024, 1024]}, "vitg": {"encoder": "vitg", "features": 384, "out_channels": [1536, 1536, 1536, 1536]}, } MODEL_URLS = { "vits": "https://huggingface.co/depth-anything/Depth-Anything-V2-Small/resolve/main/depth_anything_v2_vits.pth", "vitb": "https://huggingface.co/depth-anything/Depth-Anything-V2-Base/resolve/main/depth_anything_v2_vitb.pth", "vitl": "https://huggingface.co/depth-anything/Depth-Anything-V2-Large/resolve/main/depth_anything_v2_vitl.pth", "vitg": "https://huggingface.co/depth-anything/Depth-Anything-V2-Giant/resolve/main/depth_anything_v2_vitg.pth", } class DepthVideoConverter: def __init__(self): self.models = {} self.current_model = None self.model_size = "vits" self.model_dir = Path("models") self.model_dir.mkdir(exist_ok=True) def download_model(self, model_size: str) -> str: """Download model weights if not present.""" model_path = self.model_dir / f"depth_anything_v2_{model_size}.pth" if not model_path.exists(): print(f"Downloading model {model_size}...") import urllib.request url = MODEL_URLS[model_size] urllib.request.urlretrieve(url, model_path) print(f"Model downloaded to {model_path}") return str(model_path) def load_model(self, model_size: str): """Load the specified model.""" if model_size == self.current_model and model_size in self.models: return self.models[model_size] if not DEPTH_ANYTHING_AVAILABLE: raise ImportError("Depth Anything V2 is not installed. Please install it first.") model_path = self.download_model(model_size) config = MODEL_CONFIGS[model_size] model = DepthAnythingV2(**config) state_dict = torch.load(model_path, map_location="cpu") model.load_state_dict(state_dict) model.to(DEVICE).eval() self.models[model_size] = model self.current_model = model_size return model def process_frame(self, frame: np.ndarray, model, invert: bool = False) -> np.ndarray: """Process a single frame to depth map.""" # Convert BGR to RGB rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) image = Image.fromarray(rgb) # Resize for model input (multiple of 14) w, h = image.size new_w = (w // 14) * 14 new_h = (h // 14) * 14 image_resized = image.resize((new_w, new_h), Image.LANCZOS) # Convert to tensor image_array = np.array(image_resized).astype(np.float32) / 255.0 image_tensor = torch.from_numpy(image_array).permute(2, 0, 1).unsqueeze(0) # Inference with torch.no_grad(): depth = model(image_tensor.to(DEVICE)) depth = torch.nn.functional.interpolate( depth.unsqueeze(1), size=(h, w), mode="bilinear", align_corners=False, ).squeeze(1) depth = depth.squeeze().cpu().numpy() # Normalize to 0-255 depth_min = depth.min() depth_max = depth.max() depth_normalized = (depth - depth_min) / (depth_max - depth_min + 1e-8) depth_uint8 = (depth_normalized * 255).astype(np.uint8) if invert: depth_uint8 = 255 - depth_uint8 return depth_uint8 def temporal_smooth(self, depth_frames: list, strength: float = 0.5) -> list: """Apply temporal smoothing to reduce flicker.""" if len(depth_frames) < 2: return depth_frames smoothed = [depth_frames[0]] for i in range(1, len(depth_frames)): prev = smoothed[-1].astype(np.float32) curr = depth_frames[i].astype(np.float32) blended = prev * strength + curr * (1 - strength) smoothed.append(blended.astype(np.uint8)) return smoothed def extract_audio(self, input_path: str, output_path: str): """Extract audio from input video.""" cmd = [ "ffmpeg", "-i", input_path, "-vn", "-acodec", "copy", output_path, "-y" ] subprocess.run(cmd, capture_output=True) def merge_audio(self, video_path: str, audio_path: str, output_path: str): """Merge audio back into video.""" cmd = [ "ffmpeg", "-i", video_path, "-i", audio_path, "-c:v", "copy", "-c:a", "aac", "-map", "0:v:0", "-map", "1:a:0", "-shortest", output_path, "-y" ] subprocess.run(cmd, capture_output=True) def convert_video( self, input_path: str, model_size: str = "vits", output_resolution: str = "original", invert: bool = False, temporal_smoothing: float = 0.0, preserve_audio: bool = True, progress_callback=None ) -> str: """Main conversion function.""" # Load model model = self.load_model(model_size) # Open video cap = cv2.VideoCapture(input_path) if not cap.isOpened(): raise ValueError("Could not open video file") # Get video properties 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)) # Determine output resolution if output_resolution == "original": out_width, out_height = width, height else: # Parse resolution like "1280x720" parts = output_resolution.lower().split("x") out_width = int(parts[0]) out_height = int(parts[1]) # Create temp directory temp_dir = tempfile.mkdtemp() temp_video = os.path.join(temp_dir, "temp_video.mp4") temp_audio = os.path.join(temp_dir, "temp_audio.m4a") # Video writer fourcc = cv2.VideoWriter_fourcc(*"mp4v") out = cv2.VideoWriter(temp_video, fourcc, fps, (out_width, out_height), False) # Process frames depth_frames = [] frame_count = 0 while True: ret, frame = cap.read() if not ret: break # Process frame depth = self.process_frame(frame, model, invert) # Resize if needed if (out_width, out_height) != (width, height): depth = cv2.resize(depth, (out_width, out_height), interpolation=cv2.INTER_LINEAR) depth_frames.append(depth) frame_count += 1 if progress_callback: progress_callback(frame_count / total_frames) cap.release() # Apply temporal smoothing if temporal_smoothing > 0: depth_frames = self.temporal_smooth(depth_frames, temporal_smoothing) # Write frames for depth in depth_frames: out.write(depth) out.release() # Handle audio output_path = os.path.join(temp_dir, "output.mp4") if preserve_audio: # Extract audio from original self.extract_audio(input_path, temp_audio) # Merge audio with depth video self.merge_audio(temp_video, temp_audio, output_path) else: # Just copy the video shutil.copy(temp_video, output_path) # Clean up temp files os.remove(temp_video) if os.path.exists(temp_audio): os.remove(temp_audio) # Move to final location final_path = os.path.join(os.path.dirname(input_path), f"depth_{os.path.basename(input_path)}") shutil.move(output_path, final_path) shutil.rmtree(temp_dir) return final_path # Global converter instance converter = DepthVideoConverter() def convert_video_interface( video_file, model_size, output_resolution, invert, temporal_smoothing, preserve_audio, progress=gr.Progress() ): """Gradio interface function.""" if video_file is None: raise gr.Error("Please upload a video file") progress(0, desc="Starting conversion...") try: result = converter.convert_video( video_file, model_size=model_size, output_resolution=output_resolution, invert=invert, temporal_smoothing=temporal_smoothing, preserve_audio=preserve_audio, progress_callback=lambda p: progress(p, desc=f"Processing frames... {p*100:.1f}%") ) progress(1, desc="Done!") return result except Exception as e: raise gr.Error(f"Conversion failed: {str(e)}") # Create Gradio interface def create_interface(): with gr.Blocks(title="Depth Video Converter") as demo: gr.Markdown(""" # Depth Video Converter Convert regular videos to grayscale depth videos using Depth Anything V2. **Features:** - Local processing (no cloud upload) - CUDA acceleration on Windows - Apple Silicon (MPS) acceleration on Mac - Temporal smoothing to reduce flicker - Optional audio preservation """) with gr.Row(): with gr.Column(scale=1): video_input = gr.Video( label="Input Video", sources=["upload"], format="mp4" ) model_size = gr.Dropdown( choices=["vits", "vitb", "vitl", "vitg"], value="vits", label="Model Size", info="Smaller = faster, Larger = more accurate" ) output_resolution = gr.Dropdown( choices=["original", "640x480", "1280x720", "1920x1080"], value="original", label="Output Resolution" ) invert = gr.Checkbox( label="Invert Depth (black/white)", value=False ) temporal_smoothing = gr.Slider( minimum=0.0, maximum=0.9, value=0.0, step=0.1, label="Temporal Smoothing", info="Higher = less flicker but more ghosting" ) preserve_audio = gr.Checkbox( label="Preserve Original Audio", value=True ) convert_btn = gr.Button("Convert Video", variant="primary") with gr.Column(scale=1): video_output = gr.Video( label="Output Depth Video", format="mp4" ) convert_btn.click( fn=convert_video_interface, inputs=[ video_input, model_size, output_resolution, invert, temporal_smoothing, preserve_audio ], outputs=video_output ) gr.Markdown(""" ### Installation Instructions **Windows:** ```bash pip install -r requirements.txt python app.py ``` **macOS:** ```bash pip install -r requirements.txt python app.py ``` ### Requirements - Python 3.8+ - FFmpeg installed and in PATH - NVIDIA GPU (optional, for CUDA acceleration) - Apple Silicon Mac (optional, for MPS acceleration) """) return demo if __name__ == "__main__": demo = create_interface() demo.launch(server_name="127.0.0.1", server_port=7860)