Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 15 additions & 3 deletions src/python/pose_format/utils/reader.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,8 @@ def unpack_f(self, s_format: str):
return self.unpack(getattr(ConstStructs, s_format))

def unpack_empty_tensor(self, s: struct.Struct, shape: Tuple):
arr = np.empty(shape, s.format)
# Zero-stride view: keeps the shape without allocating shape-sized memory
arr = np.broadcast_to(np.array(0, dtype=s.format), shape)
self.advance(s, int(np.prod(shape)))
return arr

Expand Down Expand Up @@ -220,6 +221,17 @@ def skip(self, s: struct.Struct, times=1):
self.read_skipped += s.size * times
super().skip(s, times)

def bytes_left(self):
# Unlike BufferReader, the buffer only holds what was already read from the file,
# so count the file's remaining bytes as well (v0.1 infers frame count from this)
current = self.reader.tell() # remember the position, since seeking moves it
file_size = self.reader.seek(0, 2) # seek 0 bytes from the end (SEEK_END): reads nothing, returns the file size
self.reader.seek(current) # restore the position so read_chunk continues from the right place
return file_size - self.read_offset

def buffered_bytes_left(self):
return len(self.buffer) - self.read_offset + self.read_skipped

def read_chunk(self, chunk_size: int):
if self.read_offset > self.reader.tell():
self.reader.seek(self.read_offset, 0) # 0 means absolute seek
Expand All @@ -230,8 +242,8 @@ def read_chunk(self, chunk_size: int):
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 self.buffered_bytes_left() < n:
self.read_chunk(n - self.buffered_bytes_left())


if __name__ == "__main__":
Expand Down
11 changes: 11 additions & 0 deletions src/python/pose_format/utils/reader_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,17 @@ def test_unpack_numpy_writeable(self):

arr -= 0.1

def test_unpack_empty_tensor(self):
""" Test that unpack_empty_tensor keeps shape and dtype, and advances the buffer"""
buffer = struct.pack("<ffff", 1., 2.5, 3.5, 4.5)
reader = BufferReader(buffer)

arr = reader.unpack_empty_tensor(ConstStructs.float, (2, 2))

self.assertEqual(arr.shape, (2, 2))
self.assertEqual(arr.dtype, np.float32)
self.assertEqual(reader.bytes_left(), 0, msg="Buffer should advance as if the data was read")

def test_unpack_torch(self):
""" Test that unpack_torch returns the correct value"""
buffer = struct.pack("<ffff", 1., 2.5, 3.5, 4.5)
Expand Down
17 changes: 17 additions & 0 deletions src/python/tests/pose_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@

from pose_format.numpy.pose_body import NumPyPoseBody
from pose_format.pose import Pose
from pose_format.pose_body import EmptyPoseBody
from pose_format.pose_header import (PoseHeader, PoseHeaderComponent,
PoseHeaderDimensions)

Expand Down Expand Up @@ -323,6 +324,22 @@ def test_pose_remove_components(self):
self.assertEqual(c_copy.name, c_orig.name) # with the same name
self.assertEqual(c_copy.points, c_orig.points) # and the same points

def test_read_empty_pose_body_shape_matches_numpy_pose_body(self):
data_dir = Path(__file__).parent / "data"
for pose_file in sorted(data_dir.glob("*.pose")):
with self.subTest(pose_file=pose_file.name):
with open(pose_file, 'rb') as f:
numpy_pose = Pose.read(f, pose_body=NumPyPoseBody)
if numpy_pose.header.version == 0: # EmptyPoseBody does not support the legacy v0.0 format
continue
with open(pose_file, 'rb') as f:
empty_pose = Pose.read(f, pose_body=EmptyPoseBody)

self.assertEqual(empty_pose.body.data.shape, numpy_pose.body.data.shape)
self.assertEqual(empty_pose.body.data.dtype, numpy_pose.body.data.dtype)
self.assertEqual(empty_pose.body.confidence.shape, numpy_pose.body.confidence.shape)
self.assertEqual(empty_pose.body.fps, numpy_pose.body.fps)

def test_pose_bbox(self):
data_dir = Path(__file__).parent / "data"
with open(data_dir / 'mediapipe.pose', 'rb') as f:
Expand Down
Loading