Double buffering and dynamic version

Author: codeSamuraii

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "viteo"
-version = "0.1.1"
+dynamic = ["version"]
 description = "Hardware-accelerated video frame extraction for Apple Silicon"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -59,6 +59,11 @@ Issues = "https://github.com/codeSamuraii/viteo/issues"
 cmake.build-type = "Release"
 wheel.packages = ["src/viteo"]
 
+[tool.scikit-build.metadata.version]
+provider = "scikit_build_core.metadata.regex"
+input = "src/viteo/__init__.py"
+regex = '''^__version__ = ["'](?P<value>[^"']*)["']'''
+
 [tool.scikit-build.cmake.define]
 CMAKE_OSX_ARCHITECTURES = "arm64"
 

src/native/include/frame_extractor.h

@@ -10,7 +10,7 @@ namespace viteo {
 /// High-performance video frame extractor for Apple Silicon
 class FrameExtractor {
 public:
-    FrameExtractor(size_t batch_size = 8);
+    FrameExtractor();
     ~FrameExtractor();
 
     /// Open video file for extraction

src/native/src/bindings.cpp

@@ -30,7 +30,7 @@ NB_MODULE(_viteo, m) {
     m.doc() = "Hardware-accelerated video frame extraction for Apple Silicon";
 
     nb::class_<FrameExtractor>(m, "FrameExtractor")
-        .def(nb::init<size_t>(), nb::arg("batch_size") = 8, "Create new frame extractor")
+        .def(nb::init<>(), "Create new frame extractor")
         .def("open", &FrameExtractor::open, nb::arg("path"),
             "Open video file for extraction")
         .def("next_frame",

src/native/src/frame_extractor.mm

@@ -28,20 +28,18 @@
     int64_t cachedTotalFrames = 0;
     int64_t currentFrame = 0;
 
-    // Internal batch buffer for performance
-    size_t batch_size;
-    std::vector<uint8_t> batch_buffer;
-    size_t batch_count = 0;
-    size_t batch_index = 0;
+    std::vector<uint8_t> frame_buffer;
+    std::vector<uint8_t> prefetch_buffer;
+    bool has_prefetched_frame = false;
 
     bool isOpen = false;
     bool debugLogging = false;
 
-    Impl(size_t batch_size_param) : batch_size(batch_size_param) {
+    Impl() {
         if (std::getenv("VITEO_DEBUG")) {
             debugLogging = true;
         }
-        DEBUG_LOG("Setting batch size to " << batch_size);
+        DEBUG_LOG("Initialized frame extractor");
     }
 
     ~Impl() {
@@ -126,13 +124,15 @@ bool open(const std::string& path) {
 
             cacheMetadata(videoTrack, asset);
 
-            // Allocate batch buffer
             size_t frame_size = cachedWidth * cachedHeight * 4;
-            batch_buffer.resize(batch_size * frame_size);
-            DEBUG_LOG("Allocated batch buffer for " << batch_size << " frames");
+            frame_buffer.resize(frame_size);
+            prefetch_buffer.resize(frame_size);
+            DEBUG_LOG("Allocated frame buffers (" << (frame_size * 2 / 1024 / 1024) << " MB)");
 
             isOpen = true;
-            return setupReader(0);
+            if (!setupReader(0)) return false;
+
+            return prefetchFrame();
         }
     }
 
@@ -208,8 +208,6 @@ bool setupReader(int64_t startFrame) {
             }
 
             currentFrame = startFrame;
-            batch_count = 0;
-            batch_index = 0;
             DEBUG_LOG("Reader initialized successfully");
             return true;
         }
@@ -233,85 +231,64 @@ void copyFrameData(CVImageBufferRef imageBuffer, uint8_t* dst) {
         }
     }
 
-    /// Processes single sample buffer and adds to batch
-    bool processSampleBuffer(CMSampleBufferRef sampleBuffer, size_t frame_size) {
-        CVImageBufferRef imageBuffer = CMSampleBufferGetImageBuffer(sampleBuffer);
-        if (!imageBuffer) return false;
-
-        CVPixelBufferLockBaseAddress(imageBuffer, kCVPixelBufferLock_ReadOnly);
-
-        uint8_t* dst = batch_buffer.data() + (batch_count * frame_size);
-        copyFrameData(imageBuffer, dst);
-
-        CVPixelBufferUnlockBaseAddress(imageBuffer, kCVPixelBufferLock_ReadOnly);
-        batch_count++;
-        currentFrame++;
-
-        return true;
-    }
-
-    /// Load next batch of frames into internal buffer
-    void loadBatch() {
-        if (!reader || !output || !isOpen) {
-            batch_count = 0;
-            return;
+    bool prefetchFrame() {
+        if (!isOpen || !reader || !output) {
+            return false;
         }
 
-        size_t frame_size = cachedWidth * cachedHeight * 4;
-        batch_count = 0;
-
         @autoreleasepool {
-            while (batch_count < batch_size) {
-                if (reader.status != AVAssetReaderStatusReading) {
-                    DEBUG_LOG("Reader stopped, loaded " << batch_count << " frames");
-                    break;
-                }
-
-                CMSampleBufferRef sampleBuffer = [output copyNextSampleBuffer];
-                if (!sampleBuffer) {
-                    DEBUG_LOG("No more sample buffers, loaded " << batch_count << " frames");
-                    break;
-                }
-
-                processSampleBuffer(sampleBuffer, frame_size);
+            if (reader.status != AVAssetReaderStatusReading) {
+                return false;
+            }
+
+            CMSampleBufferRef sampleBuffer = [output copyNextSampleBuffer];
+            if (!sampleBuffer) {
+                return false;
+            }
+
+            CVImageBufferRef imageBuffer = CMSampleBufferGetImageBuffer(sampleBuffer);
+            if (!imageBuffer) {
                 CFRelease(sampleBuffer);
+                return false;
             }
-        }
 
-        batch_index = 0;
-        if (batch_count > 0) {
-            DEBUG_LOG("Loaded batch of " << batch_count << " frames");
+            CVPixelBufferLockBaseAddress(imageBuffer, kCVPixelBufferLock_ReadOnly);
+            copyFrameData(imageBuffer, prefetch_buffer.data());
+            CVPixelBufferUnlockBaseAddress(imageBuffer, kCVPixelBufferLock_ReadOnly);
+
+            CFRelease(sampleBuffer);
+            has_prefetched_frame = true;
+
+            return true;
         }
     }
 
-    /// Returns pointer to next frame from batch
     uint8_t* nextFrame() {
-        if (!isOpen) return nullptr;
-
-        if (batch_index >= batch_count) {
-            loadBatch();
-            if (batch_count == 0) {
-                DEBUG_LOG("No more frames available");
-                return nullptr;
-            }
+        if (!isOpen || !has_prefetched_frame) {
+            DEBUG_LOG("Not ready to extract frames");
+            return nullptr;
         }
 
-        size_t frame_size = cachedWidth * cachedHeight * 4;
-        uint8_t* frame_ptr = batch_buffer.data() + (batch_index * frame_size);
-        batch_index++;
-        return frame_ptr;
+        std::swap(frame_buffer, prefetch_buffer);
+        currentFrame++;
+
+        prefetchFrame();
+
+        return frame_buffer.data();
     }
 
-    /// Resets reader to specified frame index
     void reset(int64_t frameIndex) {
         if (!isOpen) return;
         DEBUG_LOG("Resetting to frame " << frameIndex);
-        setupReader(frameIndex);
+        has_prefetched_frame = false;
+        if (setupReader(frameIndex)) {
+            prefetchFrame();
+        }
     }
 };
 
 // Public interface implementation
-FrameExtractor::FrameExtractor(size_t batch_size_param) : impl(new Impl(batch_size_param)) {}
+FrameExtractor::FrameExtractor() : impl(new Impl()) {}
 FrameExtractor::~FrameExtractor() { delete impl; }
 
 bool FrameExtractor::open(const std::string& path) {

src/viteo/__init__.py

@@ -19,23 +19,21 @@
 from _viteo import FrameExtractor as _FrameExtractor
 from typing import Optional
 
-__version__ = "0.1.1"
+__version__ = "0.1.2"
 __all__ = ["FrameExtractor", "open"]
 
 
 class FrameExtractor(_FrameExtractor):
     """Hardware-accelerated video frame extractor for Apple Silicon."""
 
-    def __init__(self, path: Optional[str | pathlib.Path] = None, batch_size: int = 8):
+    def __init__(self, path: Optional[str | pathlib.Path] = None):
         """
         Initialize extractor and optionally open a video file.
 
         Args:
             path: Optional path to video file
-            batch_size: Number of frames to buffer internally (default: 8)
         """
-        super().__init__(batch_size)
-        self.batch_size = batch_size
+        super().__init__()
         if path:
             if not super().open(str(path)):
                 raise RuntimeError(f"Failed to open video: {path}")
@@ -47,20 +45,19 @@ def __exit__(self, *args):
         pass
 
 
-def open(path: str | pathlib.Path, batch_size: int = 8) -> FrameExtractor:
+def open(path: str | pathlib.Path) -> FrameExtractor:
     """
     Open a video file for frame extraction.
 
     Args:
         path: Path to video file
-        batch_size: Number of frames to buffer internally (default: 8)
 
     Returns:
         FrameExtractor configured for iteration
 
     Example:
-        with viteo.open("video.mp4", batch_size=16) as frames:
+        with viteo.open("video.mp4") as frames:
             for frame in frames:
                 process_frame(frame)
     """
-    return FrameExtractor(path, batch_size=batch_size)
+    return FrameExtractor(path)

tests/test_viteo.py

@@ -82,7 +82,7 @@ def test_constructor_with_path(sample_video):
     if not path.exists():
         pytest.skip(f"Test video not found: {path}")
 
-    extractor = viteo.FrameExtractor(str(path))
+    extractor = viteo.open(path)
 
     # Check that properties are set correctly
     assert extractor.width > 0
@@ -162,7 +162,7 @@ def test_reset(sample_video):
     if not path.exists():
         pytest.skip(f"Test video not found: {path}")
 
-    extractor = viteo.FrameExtractor(str(path))
+    extractor = viteo.open(path)
 
     # Get first frame
     first_frame = next(extractor)
@@ -220,7 +220,7 @@ def test_reset_out_of_bounds(sample_video):
     if not path.exists():
         pytest.skip(f"Test video not found: {path}")
 
-    extractor = viteo.FrameExtractor(str(path))
+    extractor = viteo.open(path)
 
     # Reset to a frame index way beyond the end of the video
     extractor.reset(1000000)
@@ -306,46 +306,19 @@ def test_performance(video_files):
     import os
     import sys
 
-    # Path to test videos
-    test_data = Path(__file__).parent / "test-data"
-
-    # Use videos provided on command line or all test videos
     if len(sys.argv) > 1:
         videos = [Path(p) for p in sys.argv[1:]]
     else:
-        videos = [
-            test_data / "video_4k.mp4",
-            test_data / "video_1080p.mp4",
-            test_data / "video_720p.mp4",
-            test_data / "video_480p.mp4",
-        ]
+        samples_dir = Path(__file__).parent / "samples"
+        videos = list(samples_dir.rglob("1080p_*.mp4", case_sensitive=False))
 
     # Run benchmark for each video
     for video_path in videos:
-        if not video_path.exists():
-            print(f"File not found: {video_path}")
+        if not video_path.is_file():
+            print(f"x Not found: {video_path}")
             continue
 
-        print(f"{'-'*20} {video_path.name} {'-'*20}")
-
-        try:
-            extractor = viteo.FrameExtractor(str(video_path))
-
-            # Print video properties
-            print(f"Video:")
-            print(f"* Resolution: {extractor.width}x{extractor.height}")
-            print(f"* FPS: {extractor.fps:.2f}")
-            print(f"* Total frames: {extractor.total_frames}")
-
-            # Extract frames and measure performance
-            fps, ms_per_frame = measure_performance(video_path)
-            print(f"Benchmark:")
-            print(f"* 256 frames extracted in {(256 * ms_per_frame / 1000):.3f}s")
-            print(f"* {fps:.1f} fps / {ms_per_frame:.3f}ms per frame")
-
-        except Exception as e:
-            print(f"\nxxx Error testing {video_path}: {e}\n")
-            import traceback
-            traceback.print_exc()
-
-        print('\n')
+        extractor = viteo.open(video_path)
+        num_frames = min(256, extractor.total_frames)
+        fps, ms_per_frame = measure_performance(video_path, num_frames)
+        print(f"* {video_path.name}: {fps:.2f} fps - {ms_per_frame:.2f}ms")