feat: enable end-to-end bag to LeRobot conversion with streaming frame alignment

Cargo.toml

@@ -23,7 +23,7 @@ roboflow-distributed = { path = "crates/roboflow-distributed", version = "0.2.0"
 
 # External dependencies
 # Pinned to specific commit for reproducible builds
-robocodec = { git = "https://github.com/archebase/robocodec", rev = "b27acf0e5881c38f8cc7f09e25c15787a440ae13" }
+robocodec = { git = "https://github.com/archebase/robocodec", rev = "838fa3d464ae8f0c9eeb61e7fa617363a272cbc5" }
 chrono = { version = "0.4", features = ["serde"] }
 async-trait = "0.1"
 tokio = { version = "1.40", features = ["rt-multi-thread", "sync"] }

crates/roboflow-dataset/src/formats/alignment/buffer.rs

@@ -5,7 +5,6 @@
 //! Frame alignment with bounded memory footprint.
 
 use std::collections::{HashMap, HashSet};
-use std::sync::Arc;
 use std::time::Instant;
 
 use robocodec::CodecValue;
@@ -164,14 +163,12 @@ impl FrameAlignmentBuffer {
         timestamped_msg: &TimestampedMessage,
         feature_name: &str,
     ) -> Vec<AlignedFrame> {
-        use roboflow_media::ImageData;
-
         // Update current timestamp
         self.current_timestamp = timestamped_msg.log_time;
         let msg = &timestamped_msg.message;
 
         // Step 1: Extract and decode image data (if any)
-        let (image_info, decoded_data, final_is_encoded) =
+        let (image_info, decoded_data, _final_is_encoded) =
             self.process_image_data(msg, feature_name);
 
         // Step 2: Extract state/action values (before mutable borrow)
@@ -184,18 +181,13 @@ impl FrameAlignmentBuffer {
         // Step 4: Add feature to the partial frame
         entry.add_feature(feature_name);
 
-        // Step 5: Add image data to the frame (if we extracted any)
-        if let Some(data) = decoded_data {
-            entry.frame.images.insert(
+        // Step 5: Add image ref to the frame (if we extracted any)
+        if let Some(_data) = decoded_data {
+            let width = image_info.as_ref().map(|i| i.width).unwrap_or(0);
+            let height = image_info.as_ref().map(|i| i.height).unwrap_or(0);
+            entry.frame.image_refs.insert(
                 feature_name.to_string(),
-                Arc::new(ImageData {
-                    width: image_info.as_ref().map(|i| i.width).unwrap_or(0),
-                    height: image_info.as_ref().map(|i| i.height).unwrap_or(0),
-                    data,
-                    original_timestamp: timestamped_msg.log_time,
-                    is_encoded: final_is_encoded,
-                    is_depth: false,
-                }),
+                crate::formats::common::ImageRef { width, height },
             );
         }
 
@@ -531,15 +523,10 @@ impl FrameAlignmentBuffer {
         let mut total = 0usize;
 
         for partial in &self.active_frames {
-            // Estimate image memory usage
-            for image in partial.frame.images.values() {
-                if image.is_encoded {
-                    // Compressed image - use actual data size
-                    total += image.data.len();
-                } else {
-                    // RGB decoded image - width * height * 3
-                    total += (image.width as usize) * (image.height as usize) * 3;
-                }
+            // Estimate image memory usage (approximate, since we only have refs)
+            for image_ref in partial.frame.image_refs.values() {
+                // Estimate as uncompressed RGB
+                total += (image_ref.width as usize) * (image_ref.height as usize) * 3;
             }
 
             // Estimate state/action memory (small contribution)
@@ -604,15 +591,28 @@ impl FrameAlignmentBuffer {
         let mut to_remove = Vec::new();
 
         for (idx, partial) in self.active_frames.iter().enumerate() {
-            // Check if frame is complete by criteria
-            let is_data_complete = self
-                .completion_criteria
-                .is_complete(&partial.received_features);
+            // Check if frame is complete by explicit feature requirements
+            // (not just min_completeness - that causes immediate completion)
+            let has_required_features = !self.completion_criteria.features.is_empty()
+                && self
+                    .completion_criteria
+                    .features
+                    .iter()
+                    .all(|(feature, count)| {
+                        partial
+                            .received_features
+                            .iter()
+                            .filter(|f| *f == feature)
+                            .count()
+                            >= *count
+                    });
 
             // Check if frame is complete by time window (eligible time has passed)
             let is_time_complete = self.current_timestamp >= partial.eligible_timestamp;
 
-            if is_data_complete || is_time_complete {
+            // Only complete if required features are present OR time window expired
+            // Don't complete just for min_completeness - that causes premature completion
+            if has_required_features || is_time_complete {
                 to_remove.push(idx);
             }
         }

crates/roboflow-dataset/src/formats/alignment/config.rs

@@ -29,7 +29,7 @@ impl StreamingConfig {
     pub fn with_fps(fps: u32) -> Self {
         Self {
             fps,
-            // Default completion window: 3 frames worth of data
+            // Default completion window: 30 frames worth of data (~1 second)
             completion_window_ns: Self::default_completion_window(fps),
             feature_requirements: HashMap::new(),
             decoder_config: None,
@@ -38,9 +38,10 @@ impl StreamingConfig {
 
     /// Calculate default completion window based on FPS.
     pub fn default_completion_window(fps: u32) -> u64 {
-        // 3 frames at the given FPS
+        // 30 frames at the given FPS (1 second window)
+        // This ensures state messages have enough time to be matched with frames
         let frame_interval_ns = 1_000_000_000u64 / fps as u64;
-        frame_interval_ns * 3
+        frame_interval_ns * 30
     }
 
     /// Get the frame interval in nanoseconds.
@@ -87,8 +88,8 @@ mod tests {
         let config = StreamingConfig::with_fps(60);
         assert_eq!(config.fps, 60);
         // 60 FPS = 16.666... ms per frame
-        // 3 frames = 49,999,998ns (1_000_000_000 / 60 * 3 with integer division)
-        assert_eq!(config.completion_window_ns, 49_999_998);
+        // 30 frames = 499,999,980ns (1_000_000_000 / 60 * 30 with integer division)
+        assert_eq!(config.completion_window_ns, 499_999_980);
     }
 
     #[test]
@@ -101,8 +102,8 @@ mod tests {
     #[test]
     fn test_completion_window_ns() {
         let config = StreamingConfig::with_fps(30);
-        // 3 frames worth = 33,333,333 * 3 = 99,999,999
-        assert_eq!(config.completion_window_ns(), 99_999_999);
+        // 30 frames worth = 33,333,333 * 30 = 999,999,990
+        assert_eq!(config.completion_window_ns(), 999_999_990);
     }
 
     #[test]

crates/roboflow-dataset/src/formats/common/base.rs

@@ -24,7 +24,19 @@
 use roboflow_core::Result;
 use roboflow_media::{AudioData, CameraInfo, ImageData};
 use std::collections::HashMap;
-use std::sync::Arc;
+
+/// Lightweight image reference carrying only dimensions.
+///
+/// Used in [`AlignedFrame`] to record image metadata for parquet path
+/// generation without storing the full pixel data. The actual image
+/// data is routed directly to the video encoder via the fast-path.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct ImageRef {
+    /// Image width in pixels.
+    pub width: u32,
+    /// Image height in pixels.
+    pub height: u32,
+}
 
 /// Aligned frame data ready for writing to dataset formats.
 ///
@@ -56,9 +68,9 @@ pub struct AlignedFrame {
     /// Target timestamp for this frame (nanoseconds).
     pub timestamp: u64,
 
-    /// Image observations by feature name (e.g., "observation.camera_0").
-    /// Uses Arc for zero-copy sharing when the same image is referenced multiple times.
-    pub images: HashMap<String, Arc<ImageData>>,
+    /// Image references by feature name (e.g., "observation.camera_0").
+    /// Contains only dimensions — actual pixel data is routed directly to the encoder.
+    pub image_refs: HashMap<String, ImageRef>,
 
     /// State observations by feature name.
     pub states: HashMap<String, Vec<f32>>,
@@ -79,22 +91,17 @@ impl AlignedFrame {
         Self {
             frame_index,
             timestamp,
-            images: HashMap::new(),
+            image_refs: HashMap::new(),
             states: HashMap::new(),
             actions: HashMap::new(),
             timestamps: HashMap::new(),
             audio: HashMap::new(),
         }
     }
 
-    /// Add an image observation.
-    pub fn add_image(&mut self, feature: String, data: ImageData) {
-        self.images.insert(feature, Arc::new(data));
-    }
-
-    /// Add an image observation from Arc (zero-copy if already Arc-wrapped).
-    pub fn add_image_arc(&mut self, feature: String, data: Arc<ImageData>) {
-        self.images.insert(feature, data);
+    /// Add an image reference (dimensions only, pixel data routed to encoder separately).
+    pub fn add_image_ref(&mut self, feature: String, image_ref: ImageRef) {
+        self.image_refs.insert(feature, image_ref);
     }
 
     /// Add a state observation.
@@ -119,7 +126,7 @@ impl AlignedFrame {
 
     /// Check if the frame has any data.
     pub fn is_empty(&self) -> bool {
-        self.images.is_empty()
+        self.image_refs.is_empty()
             && self.states.is_empty()
             && self.actions.is_empty()
             && self.audio.is_empty()
@@ -473,7 +480,7 @@ mod tests {
         let frame = AlignedFrame::new(42, 1_000_000_000);
         assert_eq!(frame.frame_index, 42);
         assert_eq!(frame.timestamp, 1_000_000_000);
-        assert!(frame.images.is_empty());
+        assert!(frame.image_refs.is_empty());
         assert!(frame.states.is_empty());
         assert!(frame.actions.is_empty());
     }
@@ -512,22 +519,20 @@ mod tests {
     }
 
     #[test]
-    fn test_aligned_frame_add_image() {
+    fn test_aligned_frame_add_image_ref() {
         let mut frame = AlignedFrame::new(0, 0);
-        let img = ImageData::new(640, 480, vec![0u8; 640 * 480 * 3]);
-        frame.add_image("camera".to_string(), img);
-
-        assert!(frame.images.contains_key("camera"));
+        frame.add_image_ref(
+            "camera".to_string(),
+            ImageRef {
+                width: 640,
+                height: 480,
+            },
+        );
+
+        assert!(frame.image_refs.contains_key("camera"));
         assert!(!frame.is_empty());
-    }
-
-    #[test]
-    fn test_aligned_frame_add_image_arc() {
-        let mut frame = AlignedFrame::new(0, 0);
-        let img = Arc::new(ImageData::new(640, 480, vec![0u8; 640 * 480 * 3]));
-        frame.add_image_arc("camera".to_string(), img.clone());
-
-        assert!(frame.images.contains_key("camera"));
+        assert_eq!(frame.image_refs["camera"].width, 640);
+        assert_eq!(frame.image_refs["camera"].height, 480);
     }
 
     #[test]

crates/roboflow-dataset/src/formats/common/frame_buffer.rs

@@ -19,7 +19,7 @@ pub fn build_video_frame_buffer(images: &[ImageData]) -> Result<(VideoFrameBuffe
     let mut skipped = 0usize;
 
     for img in images {
-        if img.width == 0 || img.height == 0 {
+        if !img.is_encoded && (img.width == 0 || img.height == 0) {
             skipped += 1;
             continue;
         }

crates/roboflow-dataset/src/formats/common/mod.rs

@@ -26,7 +26,9 @@ pub mod progress;
 pub mod ring_buffer;
 
 // Re-export core types (shared across all formats)
-pub use base::{AlignedFrame, DatasetFrame, DatasetWriter, DatasetWriterError, WriterStats};
+pub use base::{
+    AlignedFrame, DatasetFrame, DatasetWriter, DatasetWriterError, ImageRef, WriterStats,
+};
 
 // Re-export frame types from roboflow_media (they were moved from roboflow_core)
 pub use roboflow_media::{AudioData, CameraInfo, ImageData, ImageDataError};

crates/roboflow-dataset/src/formats/lerobot/trait_impl.rs

@@ -7,7 +7,7 @@
 //! This module defines the [`LerobotWriterTrait`] which extends the common
 //! [`DatasetWriter`] trait with LeRobot-specific functionality.
 
-use crate::formats::common::{AlignedFrame, DatasetWriter, WriterStats};
+use crate::formats::common::{AlignedFrame, DatasetWriter, ImageRef, WriterStats};
 use crate::formats::lerobot::config::LerobotConfig;
 use roboflow_core::Result;
 
@@ -88,13 +88,13 @@ pub trait LerobotWriterTrait: DatasetWriter {
         self.write_frame(frame)
     }
 
-    /// Add image data for a camera frame.
+    /// Add image reference metadata for a camera frame.
     ///
     /// # Arguments
     ///
     /// * `camera` - Camera name (e.g., "cam_high")
-    /// * `data` - Image data
-    fn add_image(&mut self, camera: String, data: crate::formats::common::ImageData);
+    /// * `image_ref` - Image reference with dimensions
+    fn add_image_ref(&mut self, camera: String, image_ref: ImageRef);
 
     /// Finalize the dataset and write metadata files.
     ///
@@ -183,7 +183,7 @@ mod tests {
             self.metadata.register_task(task)
         }
 
-        fn add_image(&mut self, _camera: String, _data: crate::formats::common::ImageData) {}
+        fn add_image_ref(&mut self, _camera: String, _image_ref: ImageRef) {}
 
         fn metadata(&self) -> &crate::formats::lerobot::metadata::MetadataCollector {
             &self.metadata