Add direct encoding option to Trajectory class and optimize stream creation

- Introduced `force_direct_encoding` parameter to `add` and `add_by_dict` methods for direct codec encoding.
- Updated stream creation logic to conditionally use direct encoding or fallback to rawvideo.
- Enhanced batch data processing in `from_list_of_dicts` and `from_dict_of_lists` methods.
- Refactored `PyAVBackend` to support direct encoding and optimized stream handling.
- Removed deprecated test file for OpenX trajectory functionality.

Author: KeplerC

robodm/backend/codec_config.py

@@ -64,10 +64,17 @@ def is_valid_image_shape(shape: Tuple[int, ...],
             # AV1 also typically requires even dimensions for yuv420p
             if height % 2 != 0 or width % 2 != 0:
                 return False
+        elif codec_name == "ffv1":
+            # FFV1 can handle odd dimensions but requires minimal size
+            if height < 2 or width < 2:
+                return False
 
         # Test if the codec actually supports this resolution
-        return CodecConfig.is_codec_config_supported(width, height, "yuv420p",
-                                                     codec_name)
+        # For FFV1, test with rgb24 instead of yuv420p
+        if codec_name == "ffv1":
+            return CodecConfig.is_codec_config_supported(width, height, "rgb24", codec_name)
+        else:
+            return CodecConfig.is_codec_config_supported(width, height, "yuv420p", codec_name)
 
     @staticmethod
     def is_image_codec(codec_name: str) -> bool:
@@ -370,14 +377,16 @@ def get_pixel_format(self, codec: str, feature_type: FeatureType) -> Optional[st
         if codec in self.IMAGE_CODEC_CONFIGS:
             base_format = self.IMAGE_CODEC_CONFIGS[codec].get("pixel_format")
 
-            # For FFV1, adjust pixel format based on data type
-            if codec == "ffv1" and feature_type.dtype == "uint8":
+            # For FFV1, use RGB24 to avoid YUV conversion issues
+            if codec == "ffv1":
                 data_shape = feature_type.shape
                 if data_shape is not None and len(data_shape) == 3:
                     if data_shape[2] == 3:  # RGB
                         return "rgb24"
                     elif data_shape[2] == 4:  # RGBA
                         return "rgba"
+                # Fallback to rgb24 for any other FFV1 case
+                return "rgb24"
 
             return base_format
 

robodm/backend/pyav_backend.py

@@ -22,7 +22,7 @@
 import numpy as np
 
 from .base import ContainerBackend, StreamMetadata, PacketInfo, StreamConfig
-from robodm.feature import FeatureType
+from robodm import FeatureType
 from robodm.backend.codec_config import CodecConfig
 from .codec_manager import CodecManager
 
@@ -100,15 +100,28 @@ def encode_data_to_packets(
         data: Any, 
         stream_index: int, 
         timestamp: int,
-        codec_config: Any
+        codec_config: Any,
+        force_direct_encoding: bool = False
     ) -> List[PacketInfo]:
-        """Encode arbitrary data into packets with timestamp handling"""
+        """Encode arbitrary data into packets with timestamp handling
+        
+        Args:
+            data: Data to encode
+            stream_index: Target stream index
+            timestamp: Timestamp in milliseconds
+            codec_config: Codec configuration
+            force_direct_encoding: If True, encode directly to target format instead of rawvideo
+        """
         if stream_index not in self._idx_to_stream:
             raise ValueError(f"No stream with index {stream_index}")
 
         stream = self._idx_to_stream[stream_index]
         container_encoding = stream.codec_context.codec.name
 
+        # If force_direct_encoding is True, bypass rawvideo intermediate step
+        if force_direct_encoding and container_encoding != "rawvideo":
+            return self._encode_directly_to_target(data, stream_index, timestamp, codec_config)
+        
         # Create codec if it doesn't exist
         codec = self.codec_manager.get_codec_for_stream(stream_index)
         if codec is None:
@@ -124,6 +137,37 @@ def encode_data_to_packets(
                 return packets
 
         return []
+
+    def _encode_directly_to_target(self, data: Any, stream_index: int, timestamp: int, codec_config: Any) -> List[PacketInfo]:
+        """Encode data directly to the target codec format without intermediate rawvideo step"""
+        if stream_index not in self._idx_to_stream:
+            raise ValueError(f"No stream with index {stream_index}")
+        
+        stream = self._idx_to_stream[stream_index]
+        container_encoding = stream.codec_context.codec.name
+        
+        if container_encoding in {"ffv1", "libaom-av1", "libx264", "libx265"}:
+            # Direct video encoding
+            if isinstance(data, np.ndarray) and len(data.shape) >= 2:
+                frame = self._create_frame(data, stream)
+                frame.time_base = stream.time_base
+                frame.pts = timestamp
+                frame.dts = timestamp
+                
+                packets = []
+                for pkt in stream.encode(frame):  # type: ignore[attr-defined]
+                    packets.append(PacketInfo(
+                        data=bytes(pkt),
+                        pts=pkt.pts,
+                        dts=pkt.dts,
+                        stream_index=stream_index,
+                        time_base=(stream.time_base.numerator, stream.time_base.denominator),
+                        is_keyframe=bool(pkt.is_keyframe) if hasattr(pkt, 'is_keyframe') else False
+                    ))
+                return packets
+        
+        # Fallback to legacy encoding if direct encoding isn't supported
+        return self._legacy_encode_fallback(data, stream_index, timestamp, stream)
 
     def _get_feature_type_from_stream(self, stream: Any) -> Any:
         """Extract feature type information from stream metadata"""
@@ -612,12 +656,15 @@ def _create_frame(self, image_array, stream):
                 f"Got shape {image_array.shape}."
             )
 
-        # Create RGB frame and convert to YUV420p when required.
-        if encoding in {"libaom-av1", "ffv1", "libx264", "libx265"}:
-            frame = av.VideoFrame.from_ndarray(image_array, format="rgb24")
-            frame = frame.reformat(format="yuv420p")
-        else:
-            frame = av.VideoFrame.from_ndarray(image_array, format="rgb24")
+        # Create RGB frame
+        frame = av.VideoFrame.from_ndarray(image_array, format="rgb24")
+        
+        # Get the configured pixel format for this stream
+        configured_pix_fmt = stream.pix_fmt
+        
+        # Convert to the configured pixel format if different from RGB24
+        if configured_pix_fmt and configured_pix_fmt != "rgb24":
+            frame = frame.reformat(format=configured_pix_fmt)
 
         return frame 
 
@@ -769,4 +816,100 @@ def _transcode_raw_internal(self, packet: Any, output_stream: Any, output_contai
 
         except Exception as e:
             logger.error(f"Failed to transcode internal codec: {e}")
-            return False 
+            return False 
+
+    def create_streams_for_batch_data(
+        self,
+        sample_data: Dict[str, Any],
+        codec_config: Any,
+        feature_name_separator: str = "/"
+    ) -> Dict[str, int]:
+        """Create optimized streams for batch data processing.
+        
+        Analyzes sample data to determine optimal codec for each feature
+        and creates streams with target codec directly.
+        
+        Args:
+            sample_data: Sample data dict to analyze feature types
+            codec_config: Codec configuration
+            feature_name_separator: Separator for nested feature names
+            
+        Returns:
+            Dict mapping feature names to stream indices
+        """
+        if self.container is None:
+            raise RuntimeError("Container not opened")
+        
+        from robodm.utils.flatten import _flatten_dict
+        from robodm import FeatureType
+        
+        # Flatten the sample data
+        flattened_data = _flatten_dict(sample_data, sep=feature_name_separator)
+        
+        feature_to_stream_idx = {}
+        
+        for feature_name, sample_value in flattened_data.items():
+            # Determine feature type from sample
+            feature_type = FeatureType.from_data(sample_value)
+            
+            # Determine optimal codec for this feature
+            target_codec = codec_config.get_codec_for_feature(feature_type, feature_name)
+            container_codec = codec_config.get_container_codec(target_codec)
+            
+            # Create stream with target codec directly
+            stream = self.add_stream_for_feature(
+                feature_name=feature_name,
+                feature_type=feature_type,
+                codec_config=codec_config,
+                encoding=container_codec
+            )
+            
+            feature_to_stream_idx[feature_name] = stream.index
+            
+            logger.debug(f"Created stream for '{feature_name}' with codec '{container_codec}' (target: '{target_codec}')")
+        
+        return feature_to_stream_idx
+
+    def encode_batch_data_directly(
+        self,
+        data_batch: List[Dict[str, Any]],
+        feature_to_stream_idx: Dict[str, int],
+        codec_config: Any,
+        feature_name_separator: str = "/",
+        fps: int = 10
+    ) -> None:
+        """Encode a batch of data directly to target codecs without intermediate transcoding.
+        
+        Args:
+            data_batch: List of data dictionaries
+            feature_to_stream_idx: Mapping of feature names to stream indices
+            codec_config: Codec configuration
+            feature_name_separator: Separator for nested feature names
+            fps: Frames per second for timestamp calculation
+        """
+        from robodm.utils.flatten import _flatten_dict
+        
+        time_interval_ms = 1000 / fps
+        current_timestamp = 0
+        
+        for step_data in data_batch:
+            flattened_data = _flatten_dict(step_data, sep=feature_name_separator)
+            
+            for feature_name, value in flattened_data.items():
+                if feature_name in feature_to_stream_idx:
+                    stream_idx = feature_to_stream_idx[feature_name]
+                    
+                    # Encode directly to target format
+                    packet_infos = self.encode_data_to_packets(
+                        data=value,
+                        stream_index=stream_idx,
+                        timestamp=int(current_timestamp),
+                        codec_config=codec_config,
+                        force_direct_encoding=True
+                    )
+                    
+                    # Mux packets immediately
+                    for packet_info in packet_infos:
+                        self.mux_packet_info(packet_info)
+            
+            current_timestamp += time_interval_ms