Handle video codecs that fail in opencv

run.py

@@ -6,6 +6,7 @@ load_dotenv()
 os.environ["HF_XET_HIGH_PERFORMANCE"] = os.getenv("HF_XET_HIGH_PERFORMANCE", "1")
 os.environ["HF_HUB_DISABLE_XET"] = os.getenv("HF_HUB_DISABLE_XET", "0")
 os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1"
+os.environ["OPENCV_FFMPEG_LOGLEVEL"] = "-8"
 seed = None
 if "SEED" in os.environ:
     try:

toolkit/dataloader_mixins.py

@@ -560,55 +560,9 @@ class ImageProcessingDTOMixin:
             unique_frame_idxs = sorted(set(frames_to_extract))
             processed_frames = {}  # frame_idx -> processed frame (duplicates reuse it)
 
-            # Set frame position
+            def process_frame(rgb_frame):
-            pos = unique_frame_idxs[0]
-            if pos > 0:
-                cap.set(cv2.CAP_PROP_POS_FRAMES, pos)
-
-                # Silently verify position was set correctly (no warnings unless debug mode)
-                if hasattr(self.dataset_config, 'debug') and self.dataset_config.debug:
-                    actual_pos = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
-                    if actual_pos != pos:
-                        print_acc(f"Warning: Failed to set exact frame position. Requested: {pos}, Actual: {actual_pos}")
-
-            for frame_idx in unique_frame_idxs:
-                # skip past frames between targets without decoding them to images
-                while pos < frame_idx and cap.grab():
-                    pos += 1
-
-                ret, frame = cap.read()
-                if ret:
-                    pos += 1
-                else:
-                    # Try to provide more detailed error information
-                    actual_frame = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
-                    frame_pos_info = f"Requested frame: {frame_idx}, Actual frame position: {actual_frame}"
-                    
-                    # Try to read the next available frame as a fallback
-                    fallback_success = False
-                    for fallback_offset in [1, -1, 5, -5, 10, -10]:
-                        fallback_pos = max(0, min(frame_idx + fallback_offset, max_frame_index))
-                        cap.set(cv2.CAP_PROP_POS_FRAMES, fallback_pos)
-                        fallback_ret, fallback_frame = cap.read()
-                        if fallback_ret:
-                            # Only log in debug mode
-                            if hasattr(self.dataset_config, 'debug') and self.dataset_config.debug:
-                                print_acc(f"Falling back to nearby frame {fallback_pos} instead of {frame_idx}")
-                            frame = fallback_frame
-                            fallback_success = True
-                            # resync sequential position after the fallback seek
-                            pos = fallback_pos + 1
-                            break
-                    else:
-                        # No fallback worked, raise a more detailed exception
-                        video_info = f"Video: {self.path}, Total frames: {total_frames}, FPS: {video_fps}"
-                        raise Exception(f"Failed to read frame {frame_idx} from video. {frame_pos_info}. {video_info}")
-                
-                # Convert BGR to RGB
-                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-                
                 # Convert to PIL Image
-                img = Image.fromarray(frame)
+                img = Image.fromarray(rgb_frame)
 
                 # Apply the same processing as for single images
                 img = img.convert('RGB')
@@ -631,7 +585,84 @@ class ImageProcessingDTOMixin:
                 if transform:
                     img = transform(img)
 
-                processed_frames[frame_idx] = img
+                return img
+
+            decode_with_pyav = False
+
+            # Set frame position
+            pos = unique_frame_idxs[0]
+            if pos > 0:
+                cap.set(cv2.CAP_PROP_POS_FRAMES, pos)
+
+                # Silently verify position was set correctly (no warnings unless debug mode)
+                if hasattr(self.dataset_config, 'debug') and self.dataset_config.debug:
+                    actual_pos = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
+                    if actual_pos != pos:
+                        print_acc(f"Warning: Failed to set exact frame position. Requested: {pos}, Actual: {actual_pos}")
+
+            for frame_idx in unique_frame_idxs:
+                # skip past frames between targets without decoding them to images
+                while pos < frame_idx and cap.grab():
+                    pos += 1
+
+                ret, frame = cap.read()
+                if ret:
+                    pos += 1
+                else:
+                    # Try to read the next available frame as a fallback
+                    fallback_success = False
+                    for fallback_offset in [1, -1, 5, -5, 10, -10]:
+                        fallback_pos = max(0, min(frame_idx + fallback_offset, max_frame_index))
+                        cap.set(cv2.CAP_PROP_POS_FRAMES, fallback_pos)
+                        fallback_ret, fallback_frame = cap.read()
+                        if fallback_ret:
+                            # Only log in debug mode
+                            if hasattr(self.dataset_config, 'debug') and self.dataset_config.debug:
+                                print_acc(f"Falling back to nearby frame {fallback_pos} instead of {frame_idx}")
+                            frame = fallback_frame
+                            fallback_success = True
+                            # resync sequential position after the fallback seek
+                            pos = fallback_pos + 1
+                            break
+                    else:
+                        # No fallback worked. cv2's bundled ffmpeg cannot decode some codecs
+                        # at all (e.g. AV1 has no software decoder there), so retry the
+                        # remaining frames with PyAV below.
+                        decode_with_pyav = True
+                        break
+                
+                # Convert BGR to RGB
+                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+                processed_frames[frame_idx] = process_frame(frame)
+
+            if decode_with_pyav:
+                # cv2 could not decode this video (e.g. AV1: OpenCV's bundled ffmpeg has no
+                # software AV1 decoder). Decode the still-missing frames in one sequential
+                # PyAV pass; PyAV ships libdav1d so it handles codecs cv2 cannot.
+                import av
+
+                needed = {i for i in unique_frame_idxs if i not in processed_frames}
+                last_av_frame = None
+                with av.open(self.path) as container:
+                    decoded_idx = -1
+                    for av_frame in container.decode(video=0):
+                        decoded_idx += 1
+                        last_av_frame = av_frame
+                        if decoded_idx in needed:
+                            processed_frames[decoded_idx] = process_frame(av_frame.to_ndarray(format='rgb24'))
+                            needed.discard(decoded_idx)
+                            if not needed:
+                                break
+
+                if needed:
+                    if last_av_frame is None:
+                        video_info = f"Video: {self.path}, Total frames: {total_frames}, FPS: {video_fps}"
+                        raise Exception(f"Failed to read frames {sorted(needed)} from video with both cv2 and PyAV. {video_info}")
+                    # metadata frame count overshot the real stream; reuse the last decoded frame
+                    tail_frame = process_frame(last_av_frame.to_ndarray(format='rgb24'))
+                    for frame_idx in needed:
+                        processed_frames[frame_idx] = tail_frame
 
             # assemble in extraction order; stretched clips repeat decoded frames
             frames = [processed_frames[frame_idx] for frame_idx in frames_to_extract]