로컬 깊이 비디오 변환기를 Gradio로 구축하기

import os import tempfile import subprocess import numpy as np import cv2 import torch import gradio as gr from PIL import Image from transformers import pipeline from scipy.ndimage import gaussian_filter1d import shutil import platform import json from pathlib import Path import warnings warnings.filterwarnings('ignore') # ============================================ # Configuration and Constants # ============================================ MODEL_SIZES = { "vits": "depth-anything/Depth-Anything-V2-Small", "vitb": "depth-anything/Depth-Anything-V2-Base", "vitl": "depth-anything/Depth-Anything-V2-Large", "vitg": "depth-anything/Depth-Anything-V2-Giant" } DEFAULT_OUTPUT_RES = 768 SUPPORTED_EXTENSIONS = ['.mp4', '.mov', '.avi', '.mkv'] # ============================================ # Device Detection # ============================================ def detect_device(): """Detect best available device for inference""" if platform.system() == 'Darwin': # macOS if torch.backends.mps.is_available(): return torch.device('mps') return torch.device('cpu') elif platform.system() == 'Windows': if torch.cuda.is_available(): return torch.device('cuda') return torch.device('cpu') else: if torch.cuda.is_available(): return torch.device('cuda') return torch.device('cpu') DEVICE = detect_device() # ============================================ # Depth Estimation Pipeline # ============================================ class DepthEstimator: def __init__(self, model_size='vitb'): self.model_id = MODEL_SIZES[model_size] self.device = DEVICE self.pipe = None self.load_model() def load_model(self): """Load the depth estimation pipeline""" try: self.pipe = pipeline( "depth-estimation", model=self.model_id, device=self.device ) except Exception as e: print(f"Error loading model: {e}") # Fallback to CPU if GPU loading fails self.device = torch.device('cpu') self.pipe = pipeline( "depth-estimation", model=self.model_id, device=self.device ) def estimate_depth(self, image): """Estimate depth map from image""" if self.pipe is None: raise RuntimeError("Model not loaded") # Convert to PIL if needed if isinstance(image, np.ndarray): image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) # Get depth prediction with torch.no_grad(): depth = self.pipe(image)["depth"] # Convert to numpy array depth_np = np.array(depth) # Normalize to 0-255 range depth_np = cv2.normalize(depth_np, None, 0, 255, cv2.NORM_MINMAX) depth_np = depth_np.astype(np.uint8) return depth_np # ============================================ # Video Processing Functions # ============================================ def apply_temporal_smoothing(depth_frames, sigma=2.0): """Apply temporal smoothing to reduce flicker""" if len(depth_frames) < 3: return depth_frames # Stack frames into 3D array stack = np.stack(depth_frames) # Apply Gaussian filter along temporal axis smoothed = gaussian_filter1d(stack, sigma=sigma, axis=0) # Convert back to list return [smoothed[i] for i in range(len(depth_frames))] def process_video(input_path, output_path, model_size='vitb', output_res=DEFAULT_OUTPUT_RES, invert=False, temporal_smoothing=True, preserve_audio=True): """Main video processing function""" # Initialize depth estimator estimator = DepthEstimator(model_size) # Open video cap = cv2.VideoCapture(input_path) if not cap.isOpened(): raise ValueError(f"Could not open video: {input_path}") # Get video properties fps = cap.get(cv2.CAP_PROP_FPS) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) # Calculate output dimensions aspect_ratio = width / height if output_res: if aspect_ratio > 1: out_width = output_res out_height = int(output_res / aspect_ratio) else: out_height = output_res out_width = int(output_res * aspect_ratio) # Ensure even dimensions out_width = out_width - (out_width % 2) out_height = out_height - (out_height % 2) else: out_width, out_height = width, height # Create temporary directory for frames temp_dir = tempfile.mkdtemp() depth_frames = [] print(f"Processing {total_frames} frames...") try: # Process each frame frame_count = 0 while True: ret, frame = cap.read() if not ret: break # Resize frame for processing if frame.shape[1] != out_width or frame.shape[0] != out_height: frame = cv2.resize(frame, (out_width, out_height)) # Estimate depth depth = estimator.estimate_depth(frame) # Invert if requested if invert: depth = 255 - depth depth_frames.append(depth) frame_count += 1 if frame_count % 10 == 0: print(f"Processed {frame_count}/{total_frames} frames") # Apply temporal smoothing if requested if temporal_smoothing and len(depth_frames) > 1: print("Applying temporal smoothing...") depth_frames = apply_temporal_smoothing(depth_frames) # Write output video print("Writing output video...") fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(output_path, fourcc, fps, (out_width, out_height), False) for frame in depth_frames: out.write(frame) out.release() # Preserve audio if requested if preserve_audio: print("Preserving original audio...") temp_video = output_path + '.temp.mp4' os.rename(output_path, temp_video) cmd = [ 'ffmpeg', '-y', '-i', temp_video, '-i', input_path, '-map', '0:v:0', '-map', '1:a:0', '-c:v', 'copy', '-c:a', 'aac', '-shortest', output_path ] result = subprocess.run(cmd, capture_output=True, text=True) if result.returncode != 0: print(f"Audio preservation failed: {result.stderr}") os.rename(temp_video, output_path) else: os.remove(temp_video) print(f"Processing complete! Output saved to: {output_path}") return output_path finally: cap.release() # Clean up temp directory shutil.rmtree(temp_dir, ignore_errors=True) # ============================================ # Gradio Interface # ============================================ def create_interface(): """Create the Gradio web interface""" def process_video_wrapper(video_file, model_size, output_res, invert, temporal_smoothing, preserve_audio): """Wrapper function for Gradio""" if video_file is None: return None, "Please upload a video file." # Get file extension ext = os.path.splitext(video_file.name)[1].lower() if ext not in SUPPORTED_EXTENSIONS: return None, f"Unsupported file format: {ext}. Please use {', '.join(SUPPORTED_EXTENSIONS)}" # Create output filename output_path = os.path.join( os.path.dirname(video_file.name), f"depth_{os.path.basename(video_file.name)}" ) try: # Process video result_path = process_video( video_file.name, output_path, model_size=model_size, output_res=output_res, invert=invert, temporal_smoothing=temporal_smoothing, preserve_audio=preserve_audio ) return result_path, f"Processing complete! Output saved to: {result_path}" except Exception as e: return None, f"Error processing video: {str(e)}" # Create Gradio interface with gr.Blocks(title="Depth Video Converter") as demo: gr.Markdown(""" # Depth Video Converter Convert regular videos into depth videos using Depth Anything V2. **Features:** - Converts each frame to grayscale depth map - Supports MP4/MOV uploads - Automatic hardware acceleration (CUDA/MPS/CPU) - Temporal smoothing to reduce flicker - Optional audio preservation """) with gr.Row(): with gr.Column(): video_input = gr.Video( label="Input Video", sources=["upload"], format="mp4" ) model_size = gr.Dropdown( choices=list(MODEL_SIZES.keys()), value="vitb", label="Model Size", info="Larger models are more accurate but slower" ) output_res = gr.Slider( minimum=256, maximum=2048, value=DEFAULT_OUTPUT_RES, step=64, label="Output Resolution", info="Maximum dimension of output video" ) invert = gr.Checkbox( label="Invert Depth Map", value=False, info="Invert black/white colors" ) temporal_smoothing = gr.Checkbox( label="Temporal Smoothing", value=True, info="Reduce flicker between frames" ) preserve_audio = gr.Checkbox( label="Preserve Original Audio", value=True, info="Keep original audio track" ) process_button = gr.Button("Convert Video", variant="primary") with gr.Column(): video_output = gr.Video( label="Output Video", interactive=False ) status_text = gr.Textbox( label="Status", interactive=False ) process_button.click( fn=process_video_wrapper, inputs=[ video_input, model_size, output_res, invert, temporal_smoothing, preserve_audio ], outputs=[video_output, status_text] ) return demo # ============================================ # Main Entry Point # ============================================ def main(): """Main function to launch the application""" print("=" * 60) print("Depth Video Converter") print("=" * 60) print(f"Platform: {platform.system()}") print(f"Device: {DEVICE}") print(f"PyTorch: {torch.__version__}") print(f"CUDA Available: {torch.cuda.is_available()}") if torch.cuda.is_available(): print(f"CUDA Device: {torch.cuda.get_device_name(0)}") print("=" * 60) # Create and launch interface demo = create_interface() demo.launch( server_name="127.0.0.1", server_port=7860, share=False ) if __name__ == "__main__": main()