Update app.py

app.py

@@ -1,29 +1,16 @@
-# Copyright (2025) Bytedance Ltd. and/or its affiliates 
-
-# Licensed under the Apache License, Version 2.0 (the "License"); 
-# you may not use this file except in compliance with the License. 
-# You may obtain a copy of the License at 
-
-#     http://www.apache.org/licenses/LICENSE-2.0 
-
-# Unless required by applicable law or agreed to in writing, software 
-# distributed under the License is distributed on an "AS IS" BASIS, 
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
-# See the License for the specific language governing permissions and 
-# limitations under the License. 
-import spaces
-import gradio as gr
-import gc
-
-import numpy as np
 import os
+import gc
 import torch
+import cv2  # Wird für die Bildverarbeitung (z.B. hconcat, GaussianBlur) benötigt
+import gradio as gr
+import numpy as np
+import matplotlib.cm as cm
 
 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
 examples = [
     ['assets/example_videos/davis_rollercoaster.mp4', -1, -1, 1280],
     ['assets/example_videos/Tokyo-Walk_rgb.mp4', -1, -1, 1280],
@@ -39,6 +26,7 @@ examples = [
 
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
+# Model configuration
 model_configs = {
     'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
     'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
@@ -49,21 +37,23 @@ encoder2name = {
     'vitl': 'Large',
 }
 
-encoder='vitl'
+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")
+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"
 description = """Official demo for **Video Depth Anything**.
 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(duration=240)
+@gradio.processing_utils.threaded  # alternativ kann spaces.GPU genutzt werden, falls verfügbar
 def infer_video_depth(
     input_video: str,
     max_len: int = -1,
@@ -71,106 +61,100 @@ def infer_video_depth(
     max_res: int = 1280,
     output_dir: str = './outputs',
     input_size: int = 518,
+    stitch: bool = False,
+    grayscale: bool = False,
+    blur: float = 0.0,
 ):
+    # Read input video frames
     frames, target_fps = read_video_frames(input_video, max_len, target_fps, max_res)
+    # Infer depths 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)
     if not os.path.exists(output_dir):
         os.makedirs(output_dir)
 
-    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 the processed (RGB) video and the 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 = ""
+    if stitch:
+        # Create a stitched video (side-by-side): left: processed RGB, right: depth
+        d_min, d_max = depths.min(), depths.max()
+        stitched_frames = []
+        for i in range(min(len(frames), len(depths))):
+            rgb_frame = frames[i]
+            depth_frame = depths[i]
+            # Normalize the depth frame to [0, 255]
+            depth_norm = ((depth_frame - d_min) / (d_max - d_min) * 255).astype(np.uint8)
+            # Choose grayscale or colored mapping
+            if grayscale:
+                depth_vis = np.stack([depth_norm] * 3, axis=-1)
+            else:
+                cmap = cm.get_cmap("inferno")
+                # cmap liefert RGBA, hier verwenden wir nur die ersten drei Kanäle
+                depth_vis = (cmap(depth_norm / 255.0)[..., :3] * 255).astype(np.uint8)
+            # Apply Gaussian blur if requested
+            if blur > 0:
+                kernel_size = int(blur * 20) * 2 + 1  # ensures odd kernel size
+                depth_vis = cv2.GaussianBlur(depth_vis, (kernel_size, kernel_size), 0)
+            # Concatenate side-by-side
+            stitched = cv2.hconcat([rgb_frame, depth_vis])
+            stitched_frames.append(stitched)
+        stitched_frames = np.array(stitched_frames)
+        stitched_video_path = os.path.join(output_dir, os.path.splitext(video_name)[0] + '_stitched.mp4')
+        save_video(stitched_frames, stitched_video_path, fps=fps)
+
     gc.collect()
     torch.cuda.empty_cache()
 
-    return [processed_video_path, depth_vis_path]
+    # Return three outputs: processed RGB video, depth visualization, and (optionally) stitched video.
+    # Falls stitch nicht aktiviert ist, wird ein leerer String zurückgegeben.
+    return [processed_video_path, depth_vis_path, stitched_video_path]
 
 
 def construct_demo():
     with gr.Blocks(analytics_enabled=False) as demo:
         gr.Markdown(title)
         gr.Markdown(description)
-        gr.Markdown("### If you find this work useful, please help ⭐ the [\[Github Repo\]](https://github.com/DepthAnything/Video-Depth-Anything). Thanks for your attention!")
+        gr.Markdown("### If you find this work useful, please help ⭐ the [Github Repo](https://github.com/DepthAnything/Video-Depth-Anything). Thanks for your attention!")
 
         with gr.Row(equal_height=True):
             with gr.Column(scale=1):
-                input_video = gr.Video(label="Input Video")
-
-            # with gr.Tab(label="Output"):
+                input_video = gr.Video(label="Input Video", source="upload", type="filepath")
             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,
-                    )
-
+                    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):
             with gr.Column(scale=1):
-                with gr.Row(equal_height=False):
-                    with gr.Accordion("Advanced Settings", open=False):
-                        max_len = gr.Slider(
-                            label="max process length",
-                            minimum=-1,
-                            maximum=1000,
-                            value=500,
-                            step=1,
-                        )
-                        target_fps = gr.Slider(
-                            label="target FPS",
-                            minimum=-1,
-                            maximum=30,
-                            value=15,
-                            step=1,
-                        )
-                        max_res = gr.Slider(
-                            label="max side resolution",
-                            minimum=480,
-                            maximum=1920,
-                            value=1280,
-                            step=1,
-                        )
-                    generate_btn = gr.Button("Generate")
+                with gr.Accordion("Advanced Settings", open=False):
+                    max_len = gr.Slider(label="Max process length", minimum=-1, maximum=1000, value=500, step=1)
+                    target_fps = gr.Slider(label="Target FPS", minimum=-1, maximum=30, value=15, step=1)
+                    max_res = gr.Slider(label="Max side resolution", minimum=480, maximum=1920, value=1280, step=1)
+                    stitch_option = gr.Checkbox(label="Stitch RGB & Depth Videos", value=False)
+                    grayscale_option = gr.Checkbox(label="Output Depth as Grayscale", value=False)
+                    blur_slider = gr.Slider(minimum=0, maximum=1, step=0.01, label="Depth Blur Factor", value=0)
+                generate_btn = gr.Button("Generate")
             with gr.Column(scale=2):
                 pass
 
         gr.Examples(
             examples=examples,
-            inputs=[
-                input_video,
-                max_len,
-                target_fps,
-                max_res
-            ],
-            outputs=[processed_video, depth_vis_video],
+            inputs=[input_video, max_len, target_fps, max_res, stitch_option, grayscale_option, blur_slider],
+            outputs=[processed_video, depth_vis_video, stitched_video],
             fn=infer_video_depth,
             cache_examples="lazy",
         )
 
         generate_btn.click(
             fn=infer_video_depth,
-            inputs=[
-                input_video,
-                max_len,
-                target_fps,
-                max_res
-            ],
-            outputs=[processed_video, depth_vis_video],
+            inputs=[input_video, max_len, target_fps, max_res, stitch_option, grayscale_option, blur_slider],
+            outputs=[processed_video, depth_vis_video, stitched_video],
         )
 
     return demo
@@ -178,4 +162,4 @@ def construct_demo():
 if __name__ == "__main__":
     demo = construct_demo()
     demo.queue()
-    demo.launch(share=True)
+    demo.launch(share=True)