Update app.py

app.py

@@ -5,16 +5,18 @@ import cv2
 import gradio as gr
 import numpy as np
 import matplotlib.cm as cm
-import matplotlib
+import matplotlib  # New import for the updated colormap API
 import subprocess
 import sys
 import spaces
 
 from video_depth_anything.video_depth import VideoDepthAnything
-from utils.dc_utils import save_video
+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],
@@ -60,114 +62,8 @@ title = "# Video Depth Anything + 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."""
 
-def get_video_info(video_path, max_len=-1, target_fps=-1):
-    """Extract video information without loading all frames into memory."""
-    cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        raise ValueError(f"Could not open video file: {video_path}")
-    
-    # Get video properties
-    original_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    original_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    original_fps = cap.get(cv2.CAP_PROP_FPS)
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    
-    # Adjust based on max_len parameter
-    if max_len > 0 and max_len < total_frames:
-        frame_count = max_len
-    else:
-        frame_count = total_frames
-    
-    # Use target_fps if specified
-    if target_fps > 0:
-        fps = target_fps
-    else:
-        fps = original_fps
-    
-    cap.release()
-    
-    return {
-        'width': original_width,
-        'height': original_height,
-        'fps': fps,
-        'original_fps': original_fps,
-        'frame_count': frame_count,
-        'total_frames': total_frames
-    }
-
-def process_frame(frame, max_res):
-    """Process a single frame to the desired resolution."""
-    if max_res > 0:
-        h, w = frame.shape[:2]
-        scale = min(max_res / w, max_res / h)
-        if scale < 1:
-            new_w, new_h = int(w * scale), int(h * scale)
-            frame = cv2.resize(frame, (new_w, new_h))
-    return frame
-
-def read_video_frames_chunked(video_path, max_len=-1, target_fps=-1, max_res=-1, chunk_size=32):
-    """Read video frames in chunks to avoid loading the entire video into memory."""
-    cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        raise ValueError(f"Could not open video file: {video_path}")
-    
-    original_fps = cap.get(cv2.CAP_PROP_FPS)
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    
-    # Determine actual number of frames to process
-    if max_len > 0 and max_len < total_frames:
-        frame_count = max_len
-    else:
-        frame_count = total_frames
-    
-    # Use target_fps if specified
-    if target_fps > 0:
-        fps = target_fps
-        # Calculate frame skip if downsampling fps
-        if target_fps < original_fps:
-            skip = int(round(original_fps / target_fps)) - 1
-        else:
-            skip = 0
-    else:
-        fps = original_fps
-        skip = 0
-    
-    frame_idx = 0
-    processed_count = 0
-    
-    while processed_count < frame_count:
-        frames_chunk = []
-        # Read frames up to chunk size or remaining frames
-        chunk_limit = min(chunk_size, frame_count - processed_count)
-        
-        while len(frames_chunk) < chunk_limit:
-            ret, frame = cap.read()
-            if not ret:
-                break
-                
-            # Process frame if we're not skipping it
-            if frame_idx % (skip + 1) == 0:
-                # Convert from BGR to RGB
-                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-                # Resize if necessary
-                frame = process_frame(frame, max_res)
-                frames_chunk.append(frame)
-                processed_count += 1
-                
-                if processed_count >= frame_count:
-                    break
-                    
-            frame_idx += 1
-        
-        if frames_chunk:
-            yield frames_chunk, fps
-        
-        if processed_count >= frame_count or len(frames_chunk) < chunk_limit:
-            break
-    
-    cap.release()
-
 @spaces.GPU(enable_queue=True)
+
 def infer_video_depth(
     input_video: str,
     max_len: int = -1,
@@ -180,122 +76,67 @@ def infer_video_depth(
     output_dir: str = './outputs',
     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)
     if not os.path.exists(output_dir):
         os.makedirs(output_dir)
-    
-    video_name = os.path.basename(input_video)
+
+    # 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')
-    
-    # Get video info first
-    video_info = get_video_info(input_video, max_len, target_fps)
-    fps = video_info['fps']
-    frame_count = video_info['frame_count']
-    
-    print(f"Processing video: {input_video}, {frame_count} frames at {fps} fps")
-    
-    # Process the video in chunks to manage memory
-    chunk_size = 32  # Adjust based on available memory
-    
-    # We'll collect depths as we go to calculate global min/max
-    all_depths = []
-    all_processed_frames = []
-    
-    # First pass to collect frames and depths
-    frame_idx = 0
-    for frames_chunk, fps in read_video_frames_chunked(input_video, max_len, target_fps, max_res, chunk_size):
-        print(f"Processing chunk: frames {frame_idx+1}-{frame_idx+len(frames_chunk)}/{frame_count}")
-        
-        # Process this chunk of frames
-        depths, _ = video_depth_anything.infer_video_depth(frames_chunk, fps, input_size=input_size, device=DEVICE)
-        
-        # Store results (we'll need both for the output videos)
-        all_processed_frames.extend(frames_chunk)
-        all_depths.extend(depths)
-        
-        frame_idx += len(frames_chunk)
-        
-        # Free memory
-        gc.collect()
-        torch.cuda.empty_cache()
-    
-    # Calculate global min/max for depth normalization
-    depths_array = np.array(all_depths)
-    d_min, d_max = depths_array.min(), depths_array.max()
-    
-    # Save the preprocessed video and depth visualization
-    save_video(all_processed_frames, processed_video_path, fps=fps)
-    save_video(all_depths, depth_vis_path, fps=fps, is_depths=True)
-    
-    # Free some memory before stitching
-    del all_processed_frames
-    gc.collect()
-    
-    # Process stitched video if requested
+    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:
-        # Use only the first 20 characters of the base name for the output filename
+        # 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)
+        # For each frame, create a visual depth image from the inferenced depths.
+        d_min, d_max = depths.min(), depths.max()
+        stitched_frames = []
+        for i in range(min(len(full_frames), len(depths))):
+            rgb_full = full_frames[i]  # Full-resolution RGB frame.
+            depth_frame = depths[i]
+            # Normalize the depth frame to the range [0, 255].
+            depth_norm = ((depth_frame - d_min) / (d_max - d_min) * 255).astype(np.uint8)
+            # Generate depth visualization:
+            if grayscale:
+                if convert_from_color:
+                    # First, generate a color depth image using the inferno colormap,
+                    # then convert that color image to grayscale.
+                    cmap = matplotlib.colormaps.get_cmap("inferno")
+                    depth_color = (cmap(depth_norm / 255.0)[..., :3] * 255).astype(np.uint8)
+                    depth_gray = cv2.cvtColor(depth_color, cv2.COLOR_RGB2GRAY)
+                    depth_vis = np.stack([depth_gray] * 3, axis=-1)
+                else:
+                    # Directly generate a grayscale image from the normalized depth values.
+                    depth_vis = np.stack([depth_norm] * 3, axis=-1)
+            else:
+                # Generate a color depth image using the inferno colormap.
+                cmap = matplotlib.colormaps.get_cmap("inferno")
+                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 an odd kernel size.
+                depth_vis = cv2.GaussianBlur(depth_vis, (kernel_size, kernel_size), 0)
+            # Resize the depth visualization to match the full-resolution RGB frame.
+            H_full, W_full = rgb_full.shape[:2]
+            depth_vis_resized = cv2.resize(depth_vis, (W_full, H_full))
+            # Concatenate the full-resolution RGB frame (left) and the resized depth visualization (right).
+            stitched = cv2.hconcat([rgb_full, depth_vis_resized])
+            stitched_frames.append(stitched)
+        stitched_frames = np.array(stitched_frames)
+        # Use only the first 20 characters of the base name for the output filename and append '_RGBD.mp4'
         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')
-        
-        # For stitching: read the original video in full resolution and stitch frames one by one
-        stitched_frames = []
-        
-        # Process in chunks for memory efficiency
-        frame_idx = 0
-        for frames_chunk, _ in read_video_frames_chunked(input_video, max_len, target_fps, -1, chunk_size):  # No max_res for original resolution
-            print(f"Stitching chunk: frames {frame_idx+1}-{frame_idx+len(frames_chunk)}/{frame_count}")
-            
-            # Process each frame in the chunk
-            for i, rgb_full in enumerate(frames_chunk):
-                depth_idx = frame_idx + i
-                if depth_idx >= len(all_depths):
-                    break
-                    
-                depth_frame = all_depths[depth_idx]
-                
-                # Normalize the depth frame
-                depth_norm = ((depth_frame - d_min) / (d_max - d_min) * 255).astype(np.uint8)
-                
-                # Generate depth visualization
-                if grayscale:
-                    if convert_from_color:
-                        # Convert from color to grayscale
-                        cmap = matplotlib.colormaps.get_cmap("inferno")
-                        depth_color = (cmap(depth_norm / 255.0)[..., :3] * 255).astype(np.uint8)
-                        depth_gray = cv2.cvtColor(depth_color, cv2.COLOR_RGB2GRAY)
-                        depth_vis = np.stack([depth_gray] * 3, axis=-1)
-                    else:
-                        # Directly use grayscale
-                        depth_vis = np.stack([depth_norm] * 3, axis=-1)
-                else:
-                    # Use color visualization
-                    cmap = matplotlib.colormaps.get_cmap("inferno")
-                    depth_vis = (cmap(depth_norm / 255.0)[..., :3] * 255).astype(np.uint8)
-                
-                # Apply 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)
-                
-                # Resize depth visualization to match original resolution
-                H_full, W_full = rgb_full.shape[:2]
-                depth_vis_resized = cv2.resize(depth_vis, (W_full, H_full))
-                
-                # Concatenate RGB and depth
-                stitched = cv2.hconcat([rgb_full, depth_vis_resized])
-                stitched_frames.append(stitched)
-            
-            frame_idx += len(frames_chunk)
-            
-            # Free memory after processing each chunk
-            gc.collect()
-        
-        # Save the stitched video
         save_video(stitched_frames, stitched_video_path, fps=fps)
         
-        # Merge audio from the input video
+        # Merge audio from the input video into the stitched video using ffmpeg.
         temp_audio_path = stitched_video_path.replace('_RGBD.mp4', '_RGBD_audio.mp4')
         cmd = [
             "ffmpeg",
@@ -311,15 +152,11 @@ def infer_video_depth(
         ]
         subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         os.replace(temp_audio_path, stitched_video_path)
-        
-        # Free memory
-        del stitched_frames
-    
-    # Clean up
-    del all_depths
+
     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]
 
 def construct_demo():
@@ -371,4 +208,6 @@ def construct_demo():
 
 if __name__ == "__main__":
     demo = construct_demo()
+    #demo.queue()  # Enable asynchronous processing.
+    #demo.launch(share=True)
     demo.queue(max_size=2).launch()