Update frame_slicer.py

frame_slicer.py

@@ -1,58 +1,126 @@
-import cv2
-import numpy as np
-import random
-
-def extract_video_frames(video_path, n_frames=30, frame_size=(96, 96)):
-    """
-    Simplified robust frame extractor for short videos (2-10 sec)
-    - Automatically handles varying video lengths
-    - Ensures consistent output shape
-    - Optimized for MP4/MPEG
-    """
-    # Open video
-    cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        print(f"Error: Could not open video {video_path}")
-        return None
-    
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    
-    # Basic video validation
-    if total_frames < 1 or fps < 1:
-        print(f"Error: Invalid video (frames:{total_frames}, fps:{fps})")
-        cap.release()
-        return None
-    
-    # Calculate how many frames to skip (adaptive based on video length)
-    video_length = total_frames / fps
-    frame_step = max(1, int(total_frames / n_frames))
-    
-    frames = []
-    last_good_frame = None
-    
-    for i in range(n_frames):
-        # Calculate position to read (spread evenly across video)
-        pos = min(int(i * (total_frames / n_frames)), total_frames - 1)
-        cap.set(cv2.CAP_PROP_POS_FRAMES, pos)
-        
-        ret, frame = cap.read()
-        
-        # Fallback strategies if read fails
-        if not ret or frame is None:
-            if last_good_frame is not None:
-                frame = last_good_frame.copy()
-            else:
-                # Generate placeholder frame (light gray)
-                frame = np.full((*frame_size[::-1], 3), 0.8, dtype=np.float32)
-        else:
-            # Process valid frame
-            frame = cv2.resize(frame, frame_size)
-            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            frame = frame.astype(np.float32) / 255.0
-            last_good_frame = frame
-        
-        frames.append(frame)
-    
-    cap.release()
-    return np.array(frames)
+--- START OF FILE frame_slicer.py ---
+
+import cv2
+import numpy as np
+import random
+import os
+
+def extract_video_frames(video_path, n_frames=30, frame_size=(96, 96)):
+    """
+    Extracts frames from a video, handling various lengths and potential errors.
+
+    Args:
+        video_path (str): Path to the video file.
+        n_frames (int): The target number of frames to extract.
+        frame_size (tuple): The target (width, height) for each frame.
+
+    Returns:
+        np.ndarray: An array of shape (n_frames, height, width, 3) with normalized
+                    pixel values (0-1), or None if extraction fails critically.
+                    Frames will be padded if the video is too short or has read errors.
+    """
+    if not os.path.exists(video_path):
+        print(f"Error: Video file not found at {video_path}")
+        return None
+
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        print(f"Error: Could not open video file {video_path}")
+        return None
+
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+
+    # Basic validation
+    if total_frames < 1:
+        print(f"Warning: Video has {total_frames} frames. Cannot extract.")
+        cap.release()
+        # Return array of zeros matching the expected shape
+        return np.zeros((n_frames, *frame_size[::-1], 3), dtype=np.float32)
+    if fps < 1:
+        print(f"Warning: Video has invalid FPS ({fps}). Proceeding, but timing might be off.")
+        # Use a default assumption if FPS is invalid but frames exist
+        fps = 30.0 # Or another sensible default
+
+    frames = []
+    extracted_count = 0
+    last_good_frame_processed = None # Store the last successfully processed frame
+
+    # Calculate indices of frames to attempt extraction (evenly spaced)
+    # Ensure indices are within the valid range [0, total_frames - 1]
+    indices = np.linspace(0, total_frames - 1, n_frames, dtype=int)
+
+    for i, frame_index in enumerate(indices):
+        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
+        ret, frame = cap.read()
+
+        processed_frame = None
+        if ret and frame is not None:
+            try:
+                # Process valid frame
+                frame_resized = cv2.resize(frame, frame_size)
+                frame_rgb = cv2.cvtColor(frame_resized, cv2.COLOR_BGR2RGB)
+                processed_frame = frame_rgb.astype(np.float32) / 255.0
+                last_good_frame_processed = processed_frame # Update last good frame
+                extracted_count += 1
+            except cv2.error as e:
+                print(f"Warning: OpenCV error processing frame {frame_index}: {e}")
+                # Fallback to last good frame if available
+                if last_good_frame_processed is not None:
+                    processed_frame = last_good_frame_processed.copy()
+                else: # If no good frame seen yet, create a placeholder
+                    processed_frame = np.zeros((*frame_size[::-1], 3), dtype=np.float32)
+            except Exception as e:
+                 print(f"Warning: Unexpected error processing frame {frame_index}: {e}")
+                 if last_good_frame_processed is not None:
+                    processed_frame = last_good_frame_processed.copy()
+                 else:
+                    processed_frame = np.zeros((*frame_size[::-1], 3), dtype=np.float32)
+
+        else:
+            # Handle read failure (e.g., end of video reached early, corrupted frame)
+            print(f"Warning: Failed to read frame at index {frame_index}. Using fallback.")
+            if last_good_frame_processed is not None:
+                processed_frame = last_good_frame_processed.copy()
+            else:
+                # If read fails and no previous frame exists, use a zero frame
+                processed_frame = np.zeros((*frame_size[::-1], 3), dtype=np.float32)
+
+        frames.append(processed_frame)
+
+    cap.release()
+
+    if extracted_count == 0 and total_frames > 0:
+         print("Warning: Failed to extract or process any valid frames, returning array of zeros.")
+         # This case should ideally be covered by fallbacks, but as a safeguard:
+         return np.zeros((n_frames, *frame_size[::-1], 3), dtype=np.float32)
+
+    # Ensure the final output always has n_frames by padding if necessary
+    # (This should technically be handled by the loop logic now, but double-check)
+    final_frames = np.array(frames)
+    if final_frames.shape[0] < n_frames:
+        print(f"Warning: Padding needed, final array shape {final_frames.shape} vs target {n_frames}")
+        if final_frames.shape[0] == 0: # If somehow array is empty
+             padding = np.zeros((n_frames, *frame_size[::-1], 3), dtype=np.float32)
+        else:
+            padding_needed = n_frames - final_frames.shape[0]
+            # Use the very last frame in the list (could be a fallback frame) for padding
+            last_frame_for_padding = final_frames[-1][np.newaxis, ...]
+            padding = np.repeat(last_frame_for_padding, padding_needed, axis=0)
+        final_frames = np.concatenate((final_frames, padding), axis=0)
+    elif final_frames.shape[0] > n_frames:
+        # Should not happen with linspace logic, but truncate if it does
+        print(f"Warning: More frames than expected ({final_frames.shape[0]}), truncating to {n_frames}")
+        final_frames = final_frames[:n_frames]
+
+
+    # Final check of output shape
+    if final_frames.shape != (n_frames, frame_size[1], frame_size[0], 3):
+         print(f"Error: Final frame array shape mismatch! Expected {(n_frames, frame_size[1], frame_size[0], 3)}, Got {final_frames.shape}")
+         # Attempt to reshape or return None/zeros? Returning zeros is safer.
+         return np.zeros((n_frames, *frame_size[::-1], 3), dtype=np.float32)
+
+
+    return final_frames
+
+--- END OF FILE frame_slicer.py ---