diff --git a/fog_x/trajectory.py b/fog_x/trajectory.py index d84d3d4..c9266bc 100644 --- a/fog_x/trajectory.py +++ b/fog_x/trajectory.py @@ -44,7 +44,7 @@ def __init__( path: Text, mode="r", cache_dir: Optional[Text] = "/tmp/fog_x/cache/", - num_pre_initialized_h264_streams: int = 5, + lossy_compression: bool = True, feature_name_separator: Text = "/", ) -> None: """ @@ -76,6 +76,7 @@ def __init__( self.mode = mode self.stream_id_to_info = {} # stream_id: StreamInfo self.is_closed = False + self.lossy_compression = lossy_compression # check if the path exists # if not, create a new file and start data collection @@ -100,17 +101,6 @@ def _get_current_timestamp(self): def __len__(self): raise NotImplementedError - # def _pre_initialize_h264_streams(self, num_streams: int): - # """ - # Pre-initialize a configurable number of H.264 video streams. - # """ - # for i in range(num_streams): - # encoding = "libx264" - # stream = self.container_file.add_stream(encoding) - # stream.time_base = Fraction(1, 1000) - # stream.pix_fmt = "yuv420p" - # self.pre_initialized_image_streams.append(stream) - def __getitem__(self, key): """ get the value of the feature @@ -638,7 +628,7 @@ def is_packet_valid(packet): ] # Check if the stream is using rawvideo, meaning it's a pickled stream - if packet.stream.codec_context.codec.name == "libx264": + if packet.stream.codec_context.codec.name == "ffv1" or packet.stream.codec_context.codec.name == "libx264": data = pickle.loads(bytes(packet)) # Encode the image data as needed, example shown for raw images @@ -690,7 +680,7 @@ def _encode_frame(self, data: Any, stream: Any, timestamp: int) -> List[av.Packe encoding = stream.codec_context.codec.name feature_type = FeatureType.from_data(data) logger.debug(f"Encoding {stream.metadata.get('FEATURE_NAME')} with {encoding}") - if encoding == "libx264": + if encoding == "ffv1" or encoding == "libx264": if feature_type.dtype == "float32": frame = self._create_frame_depth(data, stream) else: @@ -792,13 +782,21 @@ def is_packet_valid(packet): def _add_stream_to_container(self, container, feature_name, encoding, feature_type): stream = container.add_stream(encoding) - if encoding == "libx264": + if encoding == "ffv1": stream.width = feature_type.shape[0] stream.height = feature_type.shape[1] stream.codec_context.options = { "preset": "fast", # Set preset to 'fast' for quicker encoding "tune": "zerolatency", # Reduce latency - "profile": "baseline", # Use baseline profile + } + + if encoding == "libx264": + stream.width = feature_type.shape[0] + stream.height = feature_type.shape[1] + stream.codec_context.options = { + "preset": "ultrafast", # Set preset to 'ultrafast' for quicker encoding + "tune": "zerolatency", # Reduce latency + "profile": "baseline", # no b frame } stream.metadata["FEATURE_NAME"] = feature_name @@ -840,7 +838,10 @@ def _get_encoding_of_feature( feature_type = FeatureType.from_data(feature_value) data_shape = feature_type.shape if len(data_shape) >= 2 and data_shape[0] >= 100 and data_shape[1] >= 100: - vid_coding = "libx264" + if self.lossy_compression: + vid_coding = "libx264" + else: + vid_coding = "ffv1" else: vid_coding = "rawvideo" return vid_coding