Batch size conf. with env. variable and debug logging

Author: codeSamuraii

src/native/src/frame_extractor.mm

@@ -3,6 +3,14 @@
 #import <CoreMedia/CoreMedia.h>
 #import <VideoToolbox/VideoToolbox.h>
 #include "frame_extractor.h"
+#include <cstdlib>
+#include <iostream>
+
+#define DEBUG_LOG(msg) do { \
+    if (std::getenv("VITEO_DEBUG")) { \
+        std::cerr << "[VITEO DEBUG] " << msg << std::endl; \
+    } \
+} while(0)
 
 namespace viteo {
 
@@ -21,20 +29,30 @@
     int64_t currentFrame = 0;
 
     // Internal batch buffer for performance
-    static constexpr size_t BATCH_SIZE = 16;
+    size_t batch_size = 16;
     std::vector<uint8_t> batch_buffer;
     size_t batch_count = 0;
     size_t batch_index = 0;
 
     bool isOpen = false;
 
-    Impl() {}
+    Impl() {
+        const char* batch_env = std::getenv("VITEO_BATCH_SIZE");
+        if (batch_env) {
+            int val = std::atoi(batch_env);
+            if (val > 0 && val <= 256) {
+                batch_size = val;
+                DEBUG_LOG("Setting batch size to " << batch_size);
+            }
+        }
+    }
 
     ~Impl() {
         close();
         // ARC handles cleanup automatically
     }
 
+    /// Releases all resources and resets state
     void close() {
         @autoreleasepool {
             if (reader) {
@@ -47,46 +65,116 @@ void close() {
             isOpen = false;
             currentFrame = 0;
         }
+        DEBUG_LOG("Closed video resources");
     }
 
-    bool open(const std::string& path) {
-        close();
+    /// Loads asset from file path
+    AVAsset* loadAsset(const std::string& path) {
+        NSString* nsPath = [NSString stringWithUTF8String:path.c_str()];
+        NSURL* url = [NSURL fileURLWithPath:nsPath];
+        AVAsset* loadedAsset = [AVAsset assetWithURL:url];
 
-        @autoreleasepool {
-            NSString* nsPath = [NSString stringWithUTF8String:path.c_str()];
-            NSURL* url = [NSURL fileURLWithPath:nsPath];
+        if (loadedAsset) {
+            DEBUG_LOG("Loaded asset from: " << path);
+        } else {
+            DEBUG_LOG("Failed to load asset from: " << path);
+        }
 
-            asset = [AVAsset assetWithURL:url];
-            if (!asset) return false;
+        return loadedAsset;
+    }
+
+    /// Extracts video track from asset
+    AVAssetTrack* extractVideoTrack(AVAsset* videoAsset) {
+        #pragma clang diagnostic push
+        #pragma clang diagnostic ignored "-Wdeprecated-declarations"
+        NSArray* tracks = [videoAsset tracksWithMediaType:AVMediaTypeVideo];
+        #pragma clang diagnostic pop
+
+        if (tracks.count == 0) {
+            DEBUG_LOG("No video tracks found");
+            return nil;
+        }
+
+        DEBUG_LOG("Found " << tracks.count << " video track(s)");
+        return tracks[0];
+    }
 
-            #pragma clang diagnostic push
-            #pragma clang diagnostic ignored "-Wdeprecated-declarations"
-            NSArray* tracks = [asset tracksWithMediaType:AVMediaTypeVideo];
-            #pragma clang diagnostic pop
+    /// Caches video metadata from track
+    void cacheMetadata(AVAssetTrack* track, AVAsset* videoAsset) {
+        CGSize size = [track naturalSize];
+        cachedWidth = static_cast<int>(size.width);
+        cachedHeight = static_cast<int>(size.height);
+        cachedFPS = [track nominalFrameRate];
+
+        CMTime duration = [videoAsset duration];
+        cachedTotalFrames = static_cast<int64_t>(
+            CMTimeGetSeconds(duration) * cachedFPS
+        );
+
+        DEBUG_LOG("Video metadata: " << cachedWidth << "x" << cachedHeight
+                  << " @ " << cachedFPS << " fps, "
+                  << cachedTotalFrames << " total frames");
+    }
 
-            if (tracks.count == 0) return false;
+    /// Opens video file and initializes extraction
+    bool open(const std::string& path) {
+        close();
 
-            videoTrack = tracks[0];
+        @autoreleasepool {
+            asset = loadAsset(path);
+            if (!asset) return false;
 
-            CGSize size = [videoTrack naturalSize];
-            cachedWidth = static_cast<int>(size.width);
-            cachedHeight = static_cast<int>(size.height);
-            cachedFPS = [videoTrack nominalFrameRate];
+            videoTrack = extractVideoTrack(asset);
+            if (!videoTrack) return false;
 
-            CMTime duration = [asset duration];
-            cachedTotalFrames = static_cast<int64_t>(
-                CMTimeGetSeconds(duration) * cachedFPS
-            );
+            cacheMetadata(videoTrack, asset);
 
             // Allocate batch buffer
             size_t frame_size = cachedWidth * cachedHeight * 4;
-            batch_buffer.resize(BATCH_SIZE * frame_size);
+            batch_buffer.resize(batch_size * frame_size);
+            DEBUG_LOG("Allocated batch buffer for " << batch_size << " frames");
 
             isOpen = true;
             return setupReader(0);
         }
     }
 
+    /// Creates output settings dictionary for hardware accelerated decoding
+    NSDictionary* createOutputSettings() {
+        return @{
+            (id)kCVPixelBufferPixelFormatTypeKey: @(kCVPixelFormatType_32BGRA),
+            (id)kCVPixelBufferMetalCompatibilityKey: @YES,
+            (id)kCVPixelBufferIOSurfacePropertiesKey: @{},
+            AVVideoDecompressionPropertiesKey: @{
+                (id)kVTDecompressionPropertyKey_UsingHardwareAcceleratedVideoDecoder: @YES,
+                (id)kVTDecompressionPropertyKey_PropagatePerFrameHDRDisplayMetadata: @NO,
+            },
+        };
+    }
+
+    /// Configures track output for optimal performance
+    AVAssetReaderTrackOutput* createTrackOutput(AVAssetTrack* track, NSDictionary* settings) {
+        AVAssetReaderTrackOutput* trackOutput = [[AVAssetReaderTrackOutput alloc]
+            initWithTrack:track outputSettings:settings];
+
+        trackOutput.alwaysCopiesSampleData = NO;
+        trackOutput.supportsRandomAccess = YES;
+
+        DEBUG_LOG("Created track output with hardware acceleration");
+        return trackOutput;
+    }
+
+    /// Applies time range for seeking to specific frame
+    void applyTimeRange(AVAssetReader* videoReader, int64_t startFrame) {
+        if (startFrame > 0) {
+            CMTime startTime = CMTimeMake(startFrame, cachedFPS);
+            CMTime duration = CMTimeSubtract([asset duration], startTime);
+            videoReader.timeRange = CMTimeRangeMake(startTime, duration);
+            DEBUG_LOG("Seeking to frame " << startFrame);
+        }
+    }
+
+    /// Initializes reader for frame extraction
     bool setupReader(int64_t startFrame) {
         @autoreleasepool {
             if (reader) {
@@ -97,43 +185,26 @@ bool setupReader(int64_t startFrame) {
 
             NSError* error = nil;
             reader = [[AVAssetReader alloc] initWithAsset:asset error:&error];
-            if (error || !reader) return false;
-
-            // Configure for maximum performance with BGRA output
-            NSDictionary* outputSettings = @{
-                (id)kCVPixelBufferPixelFormatTypeKey: @(kCVPixelFormatType_32BGRA),
-                (id)kCVPixelBufferMetalCompatibilityKey: @YES,
-                (id)kCVPixelBufferIOSurfacePropertiesKey: @{},
-                // Add VideoToolbox hardware acceleration hints
-                AVVideoDecompressionPropertiesKey: @{
-                    (id)kVTDecompressionPropertyKey_UsingHardwareAcceleratedVideoDecoder: @YES,
-                    (id)kVTDecompressionPropertyKey_PropagatePerFrameHDRDisplayMetadata: @NO,
-                },
-            };
-
-            output = [[AVAssetReaderTrackOutput alloc]
-                initWithTrack:videoTrack outputSettings:outputSettings];
-
-            // Critical performance settings
-            output.alwaysCopiesSampleData = NO;  // Avoid unnecessary copies
-            output.supportsRandomAccess = YES;   // Enable seeking
+            if (error || !reader) {
+                DEBUG_LOG("Failed to create reader: " << (error ? [[error localizedDescription] UTF8String] : "unknown error"));
+                return false;
+            }
+
+            NSDictionary* outputSettings = createOutputSettings();
+            output = createTrackOutput(videoTrack, outputSettings);
 
             if (![reader canAddOutput:output]) {
+                DEBUG_LOG("Cannot add output to reader");
                 reader = nil;
                 output = nil;
                 return false;
             }
 
             [reader addOutput:output];
-
-            // Set time range if seeking
-            if (startFrame > 0) {
-                CMTime startTime = CMTimeMake(startFrame, cachedFPS);
-                CMTime duration = CMTimeSubtract([asset duration], startTime);
-                reader.timeRange = CMTimeRangeMake(startTime, duration);
-            }
+            applyTimeRange(reader, startFrame);
 
             if (![reader startReading]) {
+                DEBUG_LOG("Failed to start reading");
                 reader = nil;
                 output = nil;
                 return false;
@@ -142,10 +213,47 @@ bool setupReader(int64_t startFrame) {
             currentFrame = startFrame;
             batch_count = 0;
             batch_index = 0;
+            DEBUG_LOG("Reader initialized successfully");
             return true;
         }
     }
 
+    /// Copies frame from pixel buffer to destination
+    void copyFrameData(CVImageBufferRef imageBuffer, uint8_t* dst) {
+        uint8_t* src = (uint8_t*)CVPixelBufferGetBaseAddress(imageBuffer);
+        size_t bytesPerRow = CVPixelBufferGetBytesPerRow(imageBuffer);
+        size_t data_width = cachedWidth * 4;
+        size_t data_size = cachedHeight * data_width;
+
+        if (bytesPerRow == data_width) {
+            memcpy(dst, src, data_size);
+        } else {
+
+            for (int y = 0; y < cachedHeight; y++) {
+                memcpy(dst + y * data_width,
+                       src + y * bytesPerRow,
+                       data_width);
+            }
+        }
+    }
+
+    /// 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) {
@@ -157,51 +265,39 @@ void loadBatch() {
         batch_count = 0;
 
         @autoreleasepool {
-            while (batch_count < BATCH_SIZE) {
-                if (reader.status != AVAssetReaderStatusReading) break;
+            while (batch_count < batch_size) {
+                if (reader.status != AVAssetReaderStatusReading) {
+                    DEBUG_LOG("Reader stopped, loaded " << batch_count << " frames");
+                    break;
+                }
 
                 CMSampleBufferRef sampleBuffer = [output copyNextSampleBuffer];
-                if (!sampleBuffer) break;
-
-                CVImageBufferRef imageBuffer = CMSampleBufferGetImageBuffer(sampleBuffer);
-                if (imageBuffer) {
-                    CVPixelBufferLockBaseAddress(imageBuffer, kCVPixelBufferLock_ReadOnly);
-
-                    uint8_t* src = (uint8_t*)CVPixelBufferGetBaseAddress(imageBuffer);
-                    size_t bytesPerRow = CVPixelBufferGetBytesPerRow(imageBuffer);
-                    uint8_t* dst = batch_buffer.data() + (batch_count * frame_size);
-
-                    if (bytesPerRow == cachedWidth * 4) {
-                        memcpy(dst, src, frame_size);
-                    } else {
-                        size_t copy_width = cachedWidth * 4;
-                        for (int y = 0; y < cachedHeight; y++) {
-                            memcpy(dst + y * copy_width,
-                                   src + y * bytesPerRow,
-                                   copy_width);
-                        }
-                    }
-
-                    CVPixelBufferUnlockBaseAddress(imageBuffer, kCVPixelBufferLock_ReadOnly);
-                    batch_count++;
-                    currentFrame++;
+                if (!sampleBuffer) {
+                    DEBUG_LOG("No more sample buffers, loaded " << batch_count << " frames");
+                    break;
                 }
 
+                processSampleBuffer(sampleBuffer, frame_size);
                 CFRelease(sampleBuffer);
             }
         }
 
         batch_index = 0;
+        if (batch_count > 0) {
+            DEBUG_LOG("Loaded batch of " << batch_count << " frames");
+        }
     }
 
-    /// Get pointer to next frame from batch
+    /// Returns pointer to next frame from batch
     uint8_t* nextFrame() {
         if (!isOpen) return nullptr;
 
-        // Load new batch if needed
         if (batch_index >= batch_count) {
             loadBatch();
-            if (batch_count == 0) return nullptr;
+            if (batch_count == 0) {
+                DEBUG_LOG("No more frames available");
+                return nullptr;
+            }
         }
 
         size_t frame_size = cachedWidth * cachedHeight * 4;
@@ -210,8 +306,10 @@ void loadBatch() {
         return frame_ptr;
     }
 
+    /// Resets reader to specified frame index
     void reset(int64_t frameIndex) {
         if (!isOpen) return;
+        DEBUG_LOG("Resetting to frame " << frameIndex);
         setupReader(frameIndex);
     }
 };