import gradio as gr import torch import os import tempfile import shutil import time import ffmpeg import numpy as np from PIL import Image from concurrent.futures import ThreadPoolExecutor import moviepy.editor as mp from infer import lotus # Import the depth model inference function # Set device to use the L40s GPU device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Add the preprocess_video function to limit video resolution and frame rate def preprocess_video(video_path, target_fps=24, max_resolution=(1920, 1080)): """Preprocess the video to resize and reduce its frame rate.""" video = mp.VideoFileClip(video_path) # Resize video if it's larger than the target resolution if video.size[0] > max_resolution[0] or video.size[1] > max_resolution[1]: video = video.resize(newsize=max_resolution) # Limit FPS video = video.set_fps(target_fps) return video def process_frame(frame, seed=0): """Process a single frame through the depth model and return depth map.""" try: # Convert frame to PIL Image image = Image.fromarray(frame) # Save temporary image (lotus requires a file path) with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp: image.save(tmp.name) # Process through the depth model (lotus) _, output_d = lotus(tmp.name, 'depth', seed, device) # Clean up temp file os.unlink(tmp.name) # Convert depth output to numpy array depth_array = np.array(output_d) return depth_array except Exception as e: print(f"Error processing frame: {e}") return None @spaces.GPU def process_video(video_path, fps=0, seed=0, max_workers=32): """Process video, batch frames, and use L40s GPU to generate depth maps.""" temp_dir = None try: start_time = time.time() # Preprocess the video video = preprocess_video(video_path) # Use original video FPS if not specified if fps == 0: fps = video.fps frames = list(video.iter_frames(fps=fps)) total_frames = len(frames) print(f"Processing {total_frames} frames at {fps} FPS...") # Create temporary directory for frame sequence temp_dir = tempfile.mkdtemp() frames_dir = os.path.join(temp_dir, "frames") os.makedirs(frames_dir, exist_ok=True) # Process frames in larger batches (based on GPU VRAM) batch_size = 50 # Increased batch size to fully utilize the GPU's capabilities processed_frames = [] with ThreadPoolExecutor(max_workers=max_workers) as executor: for i in range(0, total_frames, batch_size): futures = [executor.submit(process_frame, frames[j], seed) for j in range(i, min(i + batch_size, total_frames))] for j, future in enumerate(futures): try: result = future.result() if result is not None: # Save frame frame_path = os.path.join(frames_dir, f"frame_{i+j:06d}.png") Image.fromarray(result).save(frame_path) # Collect processed frame for preview processed_frames.append(result) # Update preview (only showing every 10th frame to avoid clutter) if (i + j + 1) % 10 == 0: elapsed_time = time.time() - start_time yield processed_frames[-1], None, None, f"Processed {i+j+1}/{total_frames} frames... Elapsed: {elapsed_time:.2f}s" except Exception as e: print(f"Error processing frame {i + j + 1}: {e}") print("Creating output files...") # Create output directory output_dir = os.path.join(os.path.dirname(video_path), "output") os.makedirs(output_dir, exist_ok=True) # Create ZIP of frame sequence zip_filename = f"depth_frames_{int(time.time())}.zip" zip_path = os.path.join(output_dir, zip_filename) shutil.make_archive(zip_path[:-4], 'zip', frames_dir) # Create MP4 video video_filename = f"depth_video_{int(time.time())}.mp4" video_path = os.path.join(output_dir, video_filename) try: # FFmpeg settings for high-quality MP4 stream = ffmpeg.input( os.path.join(frames_dir, 'frame_%06d.png'), pattern_type='sequence', framerate=fps ) stream = ffmpeg.output( stream, video_path, vcodec='libx264', pix_fmt='yuv420p', crf=17, # High quality threads=max_workers ) ffmpeg.run(stream, overwrite_output=True, capture_stdout=True, capture_stderr=True) print("MP4 video created successfully!") except ffmpeg.Error as e: print(f"Error creating video: {e.stderr.decode() if e.stderr else str(e)}") video_path = None print("Processing complete!") yield None, zip_path, video_path, f"Processing complete! Total time: {time.time() - start_time:.2f} seconds" except Exception as e: print(f"Error: {e}") yield None, None, None, f"Error processing video: {e}" finally: if temp_dir and os.path.exists(temp_dir): try: shutil.rmtree(temp_dir) except Exception as e: print(f"Error cleaning up temp directory: {e}") def process_wrapper(video, fps=0, seed=0, max_workers=32): if video is None: raise gr.Error("Please upload a video.") try: outputs = [] for output in process_video(video, fps, seed, max_workers): outputs.append(output) yield output return outputs[-1] except Exception as e: raise gr.Error(f"Error processing video: {str(e)}") # Custom CSS for styling custom_css = """ .title-container { text-align: center; padding: 10px 0; } #title { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif; font-size: 36px; font-weight: bold; color: #000000; padding: 10px; border-radius: 10px; display: inline-block; background: linear-gradient( 135deg, #e0f7fa, #e8f5e9, #fff9c4, #ffebee, #f3e5f5, #e1f5fe, #fff3e0, #e8eaf6 ); background-size: 400% 400%; animation: gradient-animation 15s ease infinite; } @keyframes gradient-animation { 0% { background-position: 0% 50%; } 50% { background-position: 100% 50%; } 100% { background-position: 0% 50%; } } """ # Gradio Interface with gr.Blocks(css=custom_css) as demo: gr.HTML('''
Video Depth Estimation
''') with gr.Row(): with gr.Column(): video_input = gr.Video(label="Upload Video", interactive=True, show_label=True) fps_slider = gr.Slider(minimum=0, maximum=60, step=1, value=0, label="Output FPS") seed_slider = gr.Slider(minimum=0, maximum=999999999, step=1, value=0, label="Seed") max_workers_slider = gr.Slider(minimum=1, maximum=32, step=1, value=32, label="Max Workers") btn = gr.Button("Process Video", elem_id="submit-button") with gr.Column(): preview_image = gr.Image(label="Live Preview", show_label=True) output_frames_zip = gr.File(label="Download Frame Sequence (ZIP)") output_video = gr.File(label="Download Video (MP4)") time_textbox = gr.Textbox(label="Status", interactive=False) btn.click(fn=process_wrapper , inputs=[video_input, fps_slider, seed_slider, max_workers_slider], outputs=[preview_image, output_frames_zip, output_video, time_textbox]) demo.queue() if __name__ == "__main__": demo.launch(debug=True)