def read_frames(self, num_frames: int = 1) -> TrackingData:
"""
Read the next num_frames frames from this frame store object and returns a TrackingData object.
:param num_frames: The number of frames to read from the frame store.
:return: A TrackingData object storing all frame info that was stored in this DLC Frame Store....
:raises: An EOFError if more frames were requested then were available.
"""
if not self.has_next(num_frames):
frames_left = self._header.number_of_frames - self._frames_processed
raise EOFError(
f"Only '{frames_left}' were available, and '{num_frames}' were requested."
)
self._frames_processed += num_frames
__, frame_h, frame_w, __, stride = self._header.to_list()[:5]
bp_lst = self._header.bodypart_names
track_data = TrackingData.empty_tracking_data(num_frames, len(bp_lst),
frame_w, frame_h, stride)
for frame_idx in range(track_data.get_frame_count()):
for bp_idx in range(track_data.get_bodypart_count()):
sparse_fmt_flag, has_offsets_flag = self._parse_flag_byte(
from_bytes(self._file.read(1), luint8))
data = zlib.decompress(
self._file.read(
int(from_bytes(self._file.read(8), luint64))))
if sparse_fmt_flag:
entry_len = int(from_bytes(data[:8], luint64))
if (entry_len == 0):
# If the length is 0 there is no data, continue with following frames.
continue
data = data[8:]
data, sparse_y = self._take_array(data,
dtype=luint32,
count=entry_len)
data, sparse_x = self._take_array(data,
dtype=luint32,
count=entry_len)
data, probs = self._take_array(data,
dtype=lfloat,
count=entry_len)
if (
has_offsets_flag
): # If offset flag is set to true, load in offset data....
self._init_offset_data(track_data)
data, off_y = self._take_array(data,
dtype=lfloat,
count=entry_len)
data, off_x = self._take_array(data,
dtype=lfloat,
count=entry_len)
track_data.get_offset_map()[frame_idx, sparse_y,
sparse_x, bp_idx,
1] = off_y
track_data.get_offset_map()[frame_idx, sparse_y,
sparse_x, bp_idx,
0] = off_x
track_data.get_prob_table(
frame_idx,
bp_idx)[sparse_y,
sparse_x] = probs # Set probability data...
else:
data, probs = self._take_array(data,
dtype=lfloat,
count=frame_w * frame_h)
probs = np.reshape(probs, (frame_h, frame_w))
if has_offsets_flag:
self._init_offset_data(track_data)
data, off_y = self._take_array(data,
dtype=lfloat,
count=frame_h * frame_w)
data, off_x = self._take_array(data,
dtype=lfloat,
count=frame_h * frame_w)
off_y = np.reshape(off_y, (frame_h, frame_w))
off_x = np.reshape(off_x, (frame_h, frame_w))
track_data.get_offset_map()[frame_idx, :, :, bp_idx,
1] = off_y
track_data.get_offset_map()[frame_idx, :, :, bp_idx,
0] = off_x
track_data.get_prob_table(frame_idx, bp_idx)[:] = probs
return track_data
def read_frames(self, num_frames: int = 1) -> TrackingData:
"""
Read the next num_frames frames from this frame store object and returns a TrackingData object.
:param num_frames: The number of frames to read from the frame store.
:return: A TrackingData object storing all frame info that was stored in this DLC Frame Store....
:raises: An EOFError if more frames were requested then were available.
"""
if not self.has_next(num_frames):
frames_left = self._header.number_of_frames - self._frames_processed
raise EOFError(
f"Only '{frames_left}' were available, and '{num_frames}' were requested."
)
temp_frame_idx = self._frames_processed
self._frames_processed += num_frames
__, frame_h, frame_w, __, stride = self._header.to_list()[:5]
bp_lst = self._header.bodypart_names
track_data = TrackingData.empty_tracking_data(num_frames, len(bp_lst),
frame_w, frame_h, stride)
for frame_idx in range(track_data.get_frame_count()):
frame_group = self._file[
f"{DLCH5FSConstants.FRAME_PREFIX}{temp_frame_idx}"]
for bp_idx in range(track_data.get_bodypart_count()):
bp_group = frame_group[bp_lst[bp_idx]]
is_sparse = bp_group.attrs[H5SubFrameKeys.IS_STORED_SPARSE]
has_offsets = bp_group.attrs[H5SubFrameKeys.INCLUDES_OFFSETS]
if (is_sparse):
sparse_x = bp_group[H5SubFrameKeys.X]
sparse_y = bp_group[H5SubFrameKeys.Y]
probs = bp_group[H5SubFrameKeys.PROBS]
if (has_offsets):
# If data has offsets, store them...
self._init_offset_data(track_data)
off_x = bp_group[H5SubFrameKeys.OFFSET_X]
off_y = bp_group[H5SubFrameKeys.OFFSET_Y]
track_data.get_offset_map()[frame_idx, sparse_y,
sparse_x, bp_idx,
1] = off_y
track_data.get_offset_map()[frame_idx, sparse_y,
sparse_x, bp_idx,
0] = off_x
track_data.get_source_map()[
frame_idx, sparse_y, sparse_x,
bp_idx] = probs # Set probability data...
else:
probs = bp_group[H5SubFrameKeys.PROBS]
if (has_offsets):
self._init_offset_data(track_data)
off_x = bp_group[H5SubFrameKeys.OFFSET_X]
off_y = bp_group[H5SubFrameKeys.OFFSET_Y]
track_data.get_offset_map()[frame_idx, :, :, bp_idx,
1] = off_y
track_data.get_offset_map()[frame_idx, :, :, bp_idx,
0] = off_x
track_data.get_source_map()[frame_idx, :, :,
bp_idx] = probs
return track_data