Update app.py

app.py

@@ -1,6 +1,5 @@
 import os
 import gc
-import torch
 import cv2
 import gradio as gr
 import numpy as np
@@ -10,62 +9,17 @@ import subprocess
 import sys
 import spaces
 
-from video_depth_anything.video_depth import VideoDepthAnything
 from utils.dc_utils import read_video_frames, save_video
-from huggingface_hub import hf_hub_download
 
-# Examples for the Gradio Demo.
-# Each example now contains 8 parameters:
-# [video_path, max_len, target_fps, max_res, stitch, grayscale, convert_from_color, blur]
-examples = [
-    ['assets/example_videos/octopus_01.mp4', -1, -1, 1280, True, True, True, 0.3],
-    ['assets/example_videos/chicken_01.mp4', -1, -1, 1280, True, True, True, 0.3],
-    ['assets/example_videos/gorilla_01.mp4', -1, -1, 1280, True, True, True, 0.3],
-    ['assets/example_videos/davis_rollercoaster.mp4', -1, -1, 1280, True, True, True, 0.3],
-    ['assets/example_videos/Tokyo-Walk_rgb.mp4', -1, -1, 1280, True, True, True, 0.3],
-    ['assets/example_videos/4158877-uhd_3840_2160_30fps_rgb.mp4', -1, -1, 1280, True, True, True, 0.3],
-    ['assets/example_videos/4511004-uhd_3840_2160_24fps_rgb.mp4', -1, -1, 1280, True, True, True, 0.3],
-    ['assets/example_videos/1753029-hd_1920_1080_30fps.mp4', -1, -1, 1280, True, True, True, 0.3],
-    ['assets/example_videos/davis_burnout.mp4', -1, -1, 1280, True, True, True, 0.3],
-    ['assets/example_videos/example_5473765-l.mp4', -1, -1, 1280, True, True, True, 0.3],
-    ['assets/example_videos/Istanbul-26920.mp4', -1, -1, 1280, True, True, True, 0.3],
-    ['assets/example_videos/obj_1.mp4', -1, -1, 1280, True, True, True, 0.3],
-    ['assets/example_videos/sheep_cut1.mp4', -1, -1, 1280, True, True, True, 0.3],
-]
-
-# Use GPU if available; otherwise, use CPU.
-DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-# Model configuration for different encoder variants.
-model_configs = {
-    'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
-    'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
-}
-encoder2name = {
-    'vits': 'Small',
-    'vitl': 'Large',
-}
-encoder = 'vitl'
-model_name = encoder2name[encoder]
-
-# Initialize the model.
-video_depth_anything = VideoDepthAnything(**model_configs[encoder])
-filepath = hf_hub_download(
-    repo_id=f"depth-anything/Video-Depth-Anything-{model_name}",
-    filename=f"video_depth_anything_{encoder}.pth",
-    repo_type="model"
-)
-video_depth_anything.load_state_dict(torch.load(filepath, map_location='cpu'))
-video_depth_anything = video_depth_anything.to(DEVICE).eval()
-
-title = "# Video Depth Anything + RGBD sbs output"
+title = "#RGBD sbs output"
 description = """**Video Depth Anything** + RGBD sbs output for viewing with Looking Glass Factory displays.
 Please refer to our [paper](https://arxiv.org/abs/2501.12375), [project page](https://videodepthanything.github.io/), and [github](https://github.com/DepthAnything/Video-Depth-Anything) for more details."""
 
 @spaces.GPU(enable_queue=True)
 
-def infer_video_depth(
-    input_video: str,
+def stitch_rgbd_videos(
+    processed_video: str,
+    depth_vis_video: str,
     max_len: int = -1,
     target_fps: int = -1,
     max_res: int = 1280,
@@ -77,24 +31,18 @@ def infer_video_depth(
     input_size: int = 518,
 ):
     # 1. Read input video frames for inference (downscaled to max_res).
-    frames, target_fps = read_video_frames(input_video, max_len, target_fps, max_res)
-    # 2. Perform depth inference using the model.
-    depths, fps = video_depth_anything.infer_video_depth(frames, target_fps, input_size=input_size, device=DEVICE)
-
-    video_name = os.path.basename(input_video)
+    frames, target_fps = read_video_frames(processed_video, max_len, target_fps, max_res)
+    
+    video_name = os.path.basename(processed_video)
     if not os.path.exists(output_dir):
         os.makedirs(output_dir)
-
-    # Save the preprocessed (RGB) video and the generated depth visualization.
-    processed_video_path = os.path.join(output_dir, os.path.splitext(video_name)[0] + '_src.mp4')
-    depth_vis_path = os.path.join(output_dir, os.path.splitext(video_name)[0] + '_vis.mp4')
-    save_video(frames, processed_video_path, fps=fps)
-    save_video(depths, depth_vis_path, fps=fps, is_depths=True)
-
+    
     stitched_video_path = None
     if stitch:
         # For stitching: read the original video in full resolution (without downscaling).
-        full_frames, _ = read_video_frames(input_video, max_len, target_fps, max_res=-1)
+        full_frames, _ = read_video_frames(processed_video, max_len, target_fps, max_res=-1)
+        depths, _ = read_video_frames(depth_vis_video, max_len, target_fps, max_res=-1)
+        
         # For each frame, create a visual depth image from the inferenced depths.
         d_min, d_max = depths.min(), depths.max()
         stitched_frames = []
@@ -134,7 +82,7 @@ def infer_video_depth(
         base_name = os.path.splitext(video_name)[0]
         short_name = base_name[:20]
         stitched_video_path = os.path.join(output_dir, short_name + '_RGBD.mp4')
-        save_video(stitched_frames, stitched_video_path, fps=fps)
+        save_video(stitched_frames, stitched_video_path, fps=target_fps)
         
         # Merge audio from the input video into the stitched video using ffmpeg.
         temp_audio_path = stitched_video_path.replace('_RGBD.mp4', '_RGBD_audio.mp4')
@@ -142,7 +90,7 @@ def infer_video_depth(
             "ffmpeg",
             "-y",
             "-i", stitched_video_path,
-            "-i", input_video,
+            "-i", processed_video,
             "-c:v", "copy",
             "-c:a", "aac",
             "-map", "0:v:0",
@@ -154,10 +102,9 @@ def infer_video_depth(
         os.replace(temp_audio_path, stitched_video_path)
 
     gc.collect()
-    torch.cuda.empty_cache()
 
-    # Return the preprocessed RGB video, depth visualization, and (if created) the stitched video.
-    return [processed_video_path, depth_vis_path, stitched_video_path]
+    # Return stitched video.
+    return [stitched_video_path]
 
 def construct_demo():
     with gr.Blocks(analytics_enabled=False) as demo:
@@ -168,11 +115,10 @@ def construct_demo():
         with gr.Row(equal_height=True):
             with gr.Column(scale=1):
                 # Video input component for file upload.
-                input_video = gr.Video(label="Input Video")
+                processed_video = gr.Video(label="Input Video")
+                depth_vis_video = gr.Video(label="Generated Depth Video")                    
             with gr.Column(scale=2):
                 with gr.Row(equal_height=True):
-                    processed_video = gr.Video(label="Preprocessed Video", interactive=False, autoplay=True, loop=True, show_share_button=True, scale=5)
-                    depth_vis_video = gr.Video(label="Generated Depth Video", interactive=False, autoplay=True, loop=True, show_share_button=True, scale=5)
                     stitched_video = gr.Video(label="Stitched RGBD Video", interactive=False, autoplay=True, loop=True, show_share_button=True, scale=5)
                     
         with gr.Row(equal_height=True):
@@ -189,19 +135,10 @@ def construct_demo():
             with gr.Column(scale=2):
                 pass
         
-        gr.Examples(
-            examples=examples,
-            inputs=[input_video, max_len, target_fps, max_res, stitch_option, grayscale_option, convert_from_color_option, blur_slider],
-            outputs=[processed_video, depth_vis_video, stitched_video],
-            fn=infer_video_depth,
-            cache_examples=False,
-            cache_mode="lazy",
-        )
-        
         generate_btn.click(
-            fn=infer_video_depth,
-            inputs=[input_video, max_len, target_fps, max_res, stitch_option, grayscale_option, convert_from_color_option, blur_slider],
-            outputs=[processed_video, depth_vis_video, stitched_video],
+            fn=stitch_rgbd_videos,
+            inputs=[processed_video, depth_vis_video, max_len, target_fps, max_res, stitch_option, grayscale_option, convert_from_color_option, blur_slider],
+            outputs=[stitched_video],
         )
     
     return demo