Update app.py

Fix for infer_frames_depth

app.py

@@ -105,37 +105,65 @@ def process_frame(frame, max_res):
             frame = cv2.resize(frame, (new_w, new_h))
     return frame
 
-def frame_generator(video_path, max_len=-1, target_fps=-1, max_res=-1, skip_frames=0):
-    """Generate frames from a video file one at a time."""
+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)
-    frame_count = 0
+    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
     
-    # Calculate frame skip if target_fps is specified
-    if target_fps > 0 and target_fps < original_fps:
-        skip = int(round(original_fps / target_fps)) - 1
+    # 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:
-        skip = skip_frames
+        fps = original_fps
+        skip = 0
     
     frame_idx = 0
-    while True:
-        ret, frame = cap.read()
-        if not ret or (max_len > 0 and frame_count >= max_len):
-            break
+    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)
         
-        # 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
-            processed_frame = process_frame(frame, max_res)
-            yield processed_frame
-            frame_count += 1
-            
-        frame_idx += 1
+        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()
 
@@ -164,159 +192,110 @@ def infer_video_depth(
     fps = video_info['fps']
     frame_count = video_info['frame_count']
     
-    # Set up VideoWriters
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    
-    # Setup for processing batches of frames
-    batch_size = 8  # Process frames in small batches to balance efficiency and memory usage
-    processed_frames = []
-    depth_frames = []
-    stitched_frames = []
-    
-    # Initialize min/max depth values for depth normalization
-    d_min, d_max = float('inf'), float('-inf')
-    depth_values = []
-    
-    # First pass: Process frames for depth inference and collect min/max depth values
     print(f"Processing video: {input_video}, {frame_count} frames at {fps} fps")
     
-    # Process frames in batches for depth inference
-    frame_gen = frame_generator(input_video, max_len, target_fps, max_res)
-    batch_count = 0
+    # Process the video in chunks to manage memory
+    chunk_size = 32  # Adjust based on available memory
     
-    for i, frame in enumerate(frame_gen):
-        if i % 10 == 0:
-            print(f"Processing frame {i+1}/{frame_count}")
-        
-        processed_frames.append(frame)
-        
-        # When we have a full batch or reached the end, process it
-        if len(processed_frames) == batch_size or i == frame_count - 1:
-            # Process the batch for depth
-            with torch.no_grad():
-                batch_depths = video_depth_anything.infer_frames_depth(
-                    processed_frames, 
-                    input_size=input_size, 
-                    device=DEVICE
-                )
-            
-            # Collect depth statistics and frames
-            for depth in batch_depths:
-                d_min = min(d_min, depth.min())
-                d_max = max(d_max, depth.max())
-                depth_values.append(depth)
-            
-            # Clear batch for next iteration
-            processed_frames = []
-            batch_count += 1
-            
-            # Free up memory
-            torch.cuda.empty_cache()
-            gc.collect()
+    # We'll collect depths as we go to calculate global min/max
+    all_depths = []
+    all_processed_frames = []
     
-    # Save the processed video
-    height, width = depth_values[0].shape[:2] if depth_values else (0, 0)
-    video_writer = cv2.VideoWriter(processed_video_path, fourcc, fps, (width, height))
-    
-    # Reprocess frames to save original and depth videos
-    frame_gen = frame_generator(input_video, max_len, target_fps, max_res)
-    depth_writer = cv2.VideoWriter(depth_vis_path, fourcc, fps, (width, height))
-    
-    for i, (frame, depth) in enumerate(zip(frame_gen, depth_values)):
-        # Save original frame (convert RGB to BGR for OpenCV)
-        video_writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+    # 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}")
         
-        # Normalize and visualize depth
-        depth_norm = ((depth - d_min) / (d_max - d_min) * 255).astype(np.uint8)
-        if grayscale:
-            if convert_from_color:
-                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:
-                depth_vis = np.stack([depth_norm] * 3, axis=-1)
-        else:
-            cmap = matplotlib.colormaps.get_cmap("inferno")
-            depth_vis = (cmap(depth_norm / 255.0)[..., :3] * 255).astype(np.uint8)
+        # 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)
         
-        # Apply 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)
+        frame_idx += len(frames_chunk)
         
-        # Save depth visualization (convert RGB to BGR for OpenCV)
-        depth_writer.write(cv2.cvtColor(depth_vis, cv2.COLOR_RGB2BGR))
+        # 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()
     
-    video_writer.release()
-    depth_writer.release()
+    # 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
     stitched_video_path = None
     if stitch:
-        # For stitching: read the original video in full resolution
-        video_info_full = get_video_info(input_video, max_len, target_fps)
-        original_frame_gen = frame_generator(input_video, max_len, target_fps, -1)  # No resizing
-        
-        # Create a new writer for the stitched video
+        # Use only the first 20 characters of the base name for the output filename
         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')
         
-        # Get dimensions of the first frame to setup the video writer
-        first_frame = next(frame_generator(input_video, 1, -1, -1))
-        H_full, W_full = first_frame.shape[:2]
+        # For stitching: read the original video in full resolution and stitch frames one by one
+        stitched_frames = []
         
-        # Set up the stitched video writer
-        stitched_writer = cv2.VideoWriter(
-            stitched_video_path, 
-            fourcc, 
-            fps, 
-            (W_full * 2, H_full)  # Width is doubled for side-by-side
-        )
-        
-        # Reset frame generator
-        original_frame_gen = frame_generator(input_video, max_len, target_fps, -1)
-        
-        # Process each frame
-        for i, (rgb_full, depth) in enumerate(zip(original_frame_gen, depth_values)):
-            if i % 10 == 0:
-                print(f"Stitching frame {i+1}/{frame_count}")
+        # 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}")
             
-            # Normalize and visualize depth
-            depth_norm = ((depth - d_min) / (d_max - d_min) * 255).astype(np.uint8)
-            
-            # Generate depth visualization
-            if grayscale:
-                if convert_from_color:
-                    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)
+            # 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:
-                    depth_vis = np.stack([depth_norm] * 3, axis=-1)
-            else:
-                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
-                depth_vis = cv2.GaussianBlur(depth_vis, (kernel_size, kernel_size), 0)
-            
-            # Resize depth to match original frame
-            H_full, W_full = rgb_full.shape[:2]
-            depth_vis_resized = cv2.resize(depth_vis, (W_full, H_full))
+                    # 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)
             
-            # Concatenate RGB and depth
-            stitched = cv2.hconcat([rgb_full, depth_vis_resized])
+            frame_idx += len(frames_chunk)
             
-            # Write to video (convert RGB to BGR for OpenCV)
-            stitched_writer.write(cv2.cvtColor(stitched, cv2.COLOR_RGB2BGR))
+            # Free memory after processing each chunk
+            gc.collect()
         
-        stitched_writer.release()
+        # Save the stitched video
+        save_video(stitched_frames, stitched_video_path, fps=fps)
         
-        # Merge audio from the input video into the stitched video
+        # Merge audio from the input video
         temp_audio_path = stitched_video_path.replace('_RGBD.mp4', '_RGBD_audio.mp4')
         cmd = [
             "ffmpeg",
@@ -332,8 +311,12 @@ 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()