feat(pose): read pose as bytesio, without reading the entire file

Author: AmitMY

src/python/pose_format/pose.py

@@ -1,3 +1,4 @@
+from io import BytesIO
 from itertools import chain
 from typing import BinaryIO, Dict, List, Tuple, Type, Union
 
@@ -7,9 +8,10 @@
 from pose_format.pose_body import PoseBody
 from pose_format.pose_header import (PoseHeader, PoseHeaderComponent,
                                      PoseHeaderDimensions,
-                                     PoseNormalizationInfo)
+                                     PoseNormalizationInfo, PoseHeaderCache)
 from pose_format.utils.fast_math import distance_batch
-from pose_format.utils.reader import BufferReader
+from pose_format.utils.reader import BufferReader, BytesIOReader
+
 
 
 class Pose:
@@ -29,7 +31,7 @@ def __init__(self, header: PoseHeader, body: PoseBody):
         self.body = body
 
     @staticmethod
-    def read(buffer: bytes, pose_body: Type[PoseBody] = NumPyPoseBody, **kwargs):
+    def read(buffer: Union[bytes, BytesIO], pose_body: Type[PoseBody] = NumPyPoseBody, **kwargs):
         """
         Read Pose object from buffer.
 
@@ -45,7 +47,8 @@ def read(buffer: bytes, pose_body: Type[PoseBody] = NumPyPoseBody, **kwargs):
         Pose
             Pose object.
         """
-        reader = BufferReader(buffer)
+        reader = BufferReader(buffer) if isinstance(buffer, bytes) else BytesIOReader(buffer)
+        reader.expect_to_read(PoseHeaderCache.end_offset or 10 * 1024) # Expect to read the header at least (or 10kb)
         header = PoseHeader.read(reader)
         body = pose_body.read(header, reader, **kwargs)
 

src/python/pose_format/pose_body.py

@@ -128,9 +128,9 @@ def read_v0_1_frames(cls,
         _frames = frames
         if start_frame is not None and start_frame > 0:
             if start_frame >= frames:
-                raise ValueError("Start frame is greater than the number of frames")
+                raise ValueError(f"Start frame {start_frame} is greater than the number of frames {frames}")
             # Advance to the start frame
-            reader.advance(s, int(np.prod((start_frame, *shape))))
+            reader.skip(s, int(np.prod((start_frame, *shape))))
             _frames -= start_frame
 
         remove_frames = None
@@ -142,7 +142,7 @@ def read_v0_1_frames(cls,
         tensor = tensor_reader(ConstStructs.float, shape=(_frames, *shape))
 
         if remove_frames is not None:
-            reader.advance(s, int(np.prod((remove_frames, *shape))))
+            reader.skip(s, int(np.prod((remove_frames, *shape))))
 
         return tensor
 

src/python/pose_format/utils/reader.py

@@ -1,6 +1,7 @@
 import struct
 from dataclasses import dataclass
-from typing import Tuple
+from io import BytesIO
+from typing import Tuple, Union
 
 import numpy as np
 
@@ -35,6 +36,8 @@ class ConstStructs:
 Struct format for unsigned integer"""
 
 
+
+
 class BufferReader:
     """
     Class is used to read binary data from buffer
@@ -45,11 +48,18 @@ class BufferReader:
             buffer from which to read data
         read_offset: int
             current read offset in buffer
+        read_skipped: int
+            how many bytes were skipped during reading
     """
 
-    def __init__(self, buffer: bytes):
-        self.buffer = buffer
+    def __init__(self, buffer: Union[bytearray, bytes]):
+        self.buffer: bytearray = buffer
+        self.total_bytes_read = len(buffer)
         self.read_offset = 0
+        self.read_skipped = 0
+
+    def expect_to_read(self, n: int):
+        pass
 
     def bytes_left(self):
         """
@@ -60,7 +70,7 @@ def bytes_left(self):
         int
             The number of bytes left to read.
         """
-        return len(self.buffer) - self.read_offset
+        return len(self.buffer) - self.read_offset +self.read_skipped
 
     def unpack_f(self, s_format: str):
         """
@@ -97,7 +107,9 @@ def unpack_numpy(self, s: struct.Struct, shape: Tuple):
         np.ndarray
             The unpacked NumPy array.
         """
-        arr = np.ndarray(shape, s.format, self.buffer, self.read_offset).copy()
+        self.expect_to_read(s.size * int(np.prod(shape)))
+
+        arr = np.ndarray(shape, s.format, self.buffer, self.read_offset - self.read_skipped).copy()
         self.advance(s, int(np.prod(shape)))
         return arr
 
@@ -155,7 +167,8 @@ def unpack(self, s: struct.Struct):
         -------
         Unpacked data as specified by the struct format.
         """
-        unpack: tuple = s.unpack_from(self.buffer, self.read_offset)
+        self.expect_to_read(s.size)
+        unpack: tuple = s.unpack_from(self.buffer, self.read_offset - self.read_skipped)
         self.advance(s)
         if len(unpack) == 1:
             return unpack[0]
@@ -174,6 +187,9 @@ def advance(self, s: struct.Struct, times=1):
         """
         self.read_offset += s.size * times
 
+    def skip(self, s: struct.Struct, times=1):
+        self.advance(s, times)
+
     def unpack_str(self) -> str:
         """
         Unpacks a string from the buffer.
@@ -184,10 +200,34 @@ def unpack_str(self) -> str:
             The unpacked string, encoded in UTF-8.
         """
         length: int = self.unpack(ConstStructs.ushort)
+        self.expect_to_read(length)
         bytes_: bytes = self.unpack_f("%ds" % length)
         return bytes_.decode("utf-8")
 
 
+class BytesIOReader(BufferReader):
+    def __init__(self, reader: BytesIO):
+        super().__init__(bytearray())
+        self.reader = reader
+
+    def skip(self, s: struct.Struct, times=1):
+        self.buffer = self.buffer[:self.read_offset] # remove the bytes that were not used
+        self.read_skipped += s.size * times
+        super().skip(s, times)
+
+    def read_chunk(self, chunk_size: int):
+        self.reader.seek(self.read_offset, 0) # 0 means absolute seek
+        self.buffer.extend(self.reader.read(chunk_size))
+        self.total_bytes_read += chunk_size
+
+        if not self.buffer:
+            raise EOFError("End of file reached")
+
+    def expect_to_read(self, n: int):
+        if self.bytes_left() < n:
+            self.read_chunk(n - self.bytes_left())
+
+
 if __name__ == "__main__":
     from tqdm import tqdm
 

src/python/pose_format/utils/reader_test.py

@@ -1,12 +1,15 @@
 import os
 import struct
+import tempfile
 from unittest import TestCase
 
 import numpy as np
 import torch
+from pose_format import Pose
+from pose_format.utils.generic import fake_pose
+from pose_format.utils.openpose import OpenPose_Components
 
 os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
-import tensorflow as tf
 
 from pose_format.utils.reader import BufferReader, ConstStructs
 
@@ -83,6 +86,8 @@ def test_unpack_torch(self):
 
     def test_unpack_tensorflow(self):
         """ Test that unpack_tensorflow returns the correct value"""
+        import tensorflow as tf
+
         buffer = struct.pack("<ffff", 1., 2.5, 3.5, 4.5)
         reader = BufferReader(buffer)
 
@@ -91,3 +96,21 @@ def test_unpack_tensorflow(self):
         res = tf.constant([[1., 2.5], [3.5, 4.5]])
         self.assertTrue(tf.reduce_all(tf.equal(arr, res)),
                         msg="Tensorflow unpacked array is not equal to expected array")
+
+    def test_file_reader_equal_buffer_reader(self):
+        pose = fake_pose(100, fps=25, components=OpenPose_Components)
+
+        file_path = tempfile.NamedTemporaryFile(delete=False)
+        with open(file_path.name, "wb") as f:
+            pose.write(f)
+
+        with open(file_path.name, "rb") as f:
+            pose_1 = Pose.read(f, start_frame=10, end_frame=50)
+
+        with open(file_path.name, "rb") as f:
+            pose_2 = Pose.read(f.read(), start_frame=10, end_frame=50)
+
+        self.assertEqual(pose_1.header, pose_2.header)
+        self.assertEqual(pose_1.body.fps, pose_2.body.fps)
+        self.assertTrue(np.all(pose_1.body.data == pose_2.body.data))
+        self.assertTrue(np.all(pose_1.body.confidence == pose_2.body.confidence))