def on_frames(self, scmap: TrackingData) -> Union[None, Pose]:
for frame in range(scmap.get_frame_count()):
# Plot all probability maps
for bp, ax in zip(range(scmap.get_bodypart_count()),
self._axes.flat):
ax.clear()
ax.set_aspect("equal")
if not self.AXIS_ON:
ax.axis("off")
ax.set_title(
f"Bodypart: {self._parts[bp]}, Frame: {self._current_frame}"
)
if self.PROJECT_3D:
x, y = (
np.arange(scmap.get_frame_width()),
np.arange(scmap.get_frame_height()),
)
x, y = np.meshgrid(x, y)
z = (self._logify(scmap.get_prob_table(frame, bp)) if
(self.LOG_SCALE) else scmap.get_prob_table(frame, bp))
ax.plot_surface(x, y, z, cmap=self.COLOR_MAP)
z_range = ax.get_zlim()[1] - ax.get_zlim()[0]
ax.set_zlim(ax.get_zlim()[0],
ax.get_zlim()[0] + (z_range * self.Z_SHRINK_F))
else:
ax.pcolormesh(
self._logify(scmap.get_prob_table(frame, bp)) if
(self.LOG_SCALE) else scmap.get_prob_table(frame, bp),
cmap=self.COLOR_MAP,
)
# This reverses the y-axis data, so as probability maps match the video...
ax.set_ylim(ax.get_ylim()[::-1])
# Save chart to the buffer
self._figure.tight_layout()
self._figure.canvas.draw()
img = np.reshape(
np.frombuffer(self._figure.canvas.tostring_rgb(),
dtype="uint8"),
self._figure.canvas.get_width_height()[::-1] + (3, ),
)[:, :, ::-1]
if self._vid_writer is None:
height, width, colors = img.shape
self._vid_writer = cv2.VideoWriter(self.VIDEO_NAME,
self.OUTPUT_CODEC,
self.OUTPUT_FPS,
(width, height))
self._vid_writer.write(img)
self._current_frame += 1
# Return argmax values for each frame...
return scmap.get_poses_for(
scmap.get_max_scmap_points(num_max=self._num_outputs))
def on_frames(self, scmap: TrackingData) -> Union[None, Pose]:
# If we are just starting, write the header, body part names chunk, and magic for frame data chunk...
if self._current_frame == 0:
header = DLCFSHeader(self._num_frames, scmap.get_frame_height(),
scmap.get_frame_width(),
self._video_metadata["fps"],
scmap.get_down_scaling(),
*self._video_metadata["size"],
*self._crop_off, self._bodyparts)
if (self.IS_HDF5):
self._frame_writer = DLCH5FSWriter(
self._out_file,
header,
self.THRESHOLD if (self.SPARSIFY) else None,
)
else:
self._frame_writer = DLCFSWriter(
self._out_file, header, self.THRESHOLD if
(self.SPARSIFY) else None, self.COMPRESSION_LEVEL)
# Writing all of the frames in this batch...
self._frame_writer.write_data(scmap)
self._current_frame += scmap.get_frame_count()
return scmap.get_poses_for(
scmap.get_max_scmap_points(num_max=self._num_outputs))
def on_frames(self, scmap: TrackingData) -> Union[None, Pose]:
# If the video writer has not been created, create it now and compute all needed video dimensions...
if self._vid_writer is None:
self._compute_video_measurements(
scmap.get_frame_width(), scmap.get_frame_height()
)
for frame in range(scmap.get_frame_count()):
# Clear the canvas with the background color...
self._canvas[:] = self.BACKGROUND_COLOR
# Drawing the title...
self._draw_title(f"Frame {self._current_frame}")
for bp in range(len(self._parts)):
# Compute the current subplot we are on...
subplot_y = bp // self._grid_width
subplot_x = bp % self._grid_width
self._draw_subplot(
self._parts[bp],
subplot_x,
subplot_y,
scmap.get_prob_table(frame, bp),
)
self._vid_writer.write(self._canvas)
self._current_frame += 1
# Return just like argmax...
return scmap.get_poses_for(
scmap.get_max_scmap_points(num_max=self._num_outputs)
)
def _init_offset_data(cls, track_data: TrackingData):
if track_data.get_offset_map() is None:
# If tracking data is currently None, we need to create an empty array to store all data.
shape = (
track_data.get_frame_count(),
track_data.get_frame_height(),
track_data.get_frame_width(),
track_data.get_bodypart_count(),
2,
)
track_data.set_offset_map(np.zeros(shape, dtype=lfloat))
def write_data(self, data: TrackingData):
"""
Write the following frames to the file.
:param data: A TrackingData object, which contains frame data.
"""
# Some checks to make sure tracking data parameters match those set in the header:
self._current_frame += data.get_frame_count()
if self._current_frame > self._header.number_of_frames:
raise ValueError(
f"Data Overflow! '{self._header.number_of_frames}' frames expected, tried to write "
f"'{self._current_frame + 1}' frames.")
if data.get_bodypart_count() != len(self._header.bodypart_names):
raise ValueError(
f"'{data.get_bodypart_count()}' body parts does not match the "
f"'{len(self._header.bodypart_names)}' body parts specified in the header."
)
if (data.get_frame_width() != self._header.frame_width
or data.get_frame_height() != self._header.frame_height):
raise ValueError(
"Frame dimensions don't match ones specified in header!")
for frm_idx in range(data.get_frame_count()):
# Add this frame's offset to the offset list...
idx = self._current_frame - data.get_frame_count() + frm_idx
self._frame_offsets[idx] = self._out_file.tell(
) - self._fdat_offset
for bp in range(data.get_bodypart_count()):
frame = data.get_prob_table(frm_idx, bp)
offset_table = data.get_offset_map()
if offset_table is not None:
off_y = offset_table[frm_idx, :, :, bp, 1]
off_x = offset_table[frm_idx, :, :, bp, 0]
else:
off_y = None
off_x = None
if self._threshold is not None:
# Sparsify the data by removing everything below the threshold...
sparse_y, sparse_x = np.nonzero(frame > self._threshold)
probs = frame[(sparse_y, sparse_x)]
# Check if we managed to strip out at least 2/3rds of the data, and if so write the frame using the
# sparse format. Otherwise it is actually more memory efficient to just store the entire frame...
if len(frame.flat) >= (
len(sparse_y) *
DLCFSConstants.MIN_SPARSE_SAVING_FACTOR):
# Sparse indicator flag and the offsets included flag...
self._out_file.write(
to_bytes(True | ((offset_table is not None) << 1),
luint8))
# COMPRESSED DATA:
buffer = BytesIO()
buffer.write(to_bytes(
len(sparse_y),
luint64)) # The length of the sparse data entries.
buffer.write(sparse_y.astype(luint32).tobytes(
"C")) # Y coord data
buffer.write(sparse_x.astype(luint32).tobytes(
"C")) # X coord data
buffer.write(
probs.astype(lfloat).tobytes("C")) # Probabilities
if (
offset_table is not None
): # If offset table exists, write y offsets and then x offsets.
buffer.write(
off_y[(sparse_y,
sparse_x)].astype(lfloat).tobytes("C"))
buffer.write(
off_x[(sparse_y,
sparse_x)].astype(lfloat).tobytes("C"))
# Compress the sparse data and write it's length, followed by itself....
comp_data = zlib.compress(buffer.getvalue(),
self._compression_level)
self._out_file.write(to_bytes(len(comp_data), luint64))
self._out_file.write(comp_data)
continue
# If sparse optimization mode is off or the sparse format wasted more space, just write the entire
# frame...
self._out_file.write(
to_bytes(False | ((offset_table is not None) << 1),
luint8))
buffer = BytesIO()
buffer.write(frame.astype(lfloat).tobytes(
"C")) # The probability frame...
if offset_table is not None: # Y, then X offset data if it exists...
buffer.write(off_y.astype(lfloat).tobytes("C"))
buffer.write(off_x.astype(lfloat).tobytes("C"))
comp_data = zlib.compress(buffer.getvalue(),
self._compression_level)
self._out_file.write(to_bytes(len(comp_data), luint64))
self._out_file.write(comp_data)
if (self._current_frame >= self._header.number_of_frames):
# We have reached the end, dump the flup chunk
self._write_flup_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."
)
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 write_data(self, data: TrackingData):
"""
Write the following frames to the file.
:param data: A TrackingData object, which contains frame data.
"""
# Some checks to make sure tracking data parameters match those set in the header:
current_frame_tmp = self._current_frame
self._current_frame += data.get_frame_count()
if self._current_frame > self._header.number_of_frames:
raise ValueError(
f"Data Overflow! '{self._header.number_of_frames}' frames expected, tried to write "
f"'{self._current_frame + 1}' frames.")
if data.get_bodypart_count() != len(self._header.bodypart_names):
raise ValueError(
f"'{data.get_bodypart_count()}' body parts does not match the "
f"'{len(self._header.bodypart_names)}' body parts specified in the header."
)
if (data.get_frame_width() != self._header.frame_width
or data.get_frame_height() != self._header.frame_height):
raise ValueError(
"Frame dimensions don't match ones specified in header!")
for frm_idx in range(data.get_frame_count()):
frame_grp = self._file.create_group(
f"{DLCH5FSConstants.FRAME_PREFIX}{current_frame_tmp}")
for bp_idx in range(data.get_bodypart_count()):
bp_grp = frame_grp.create_group(
self._header.bodypart_names[bp_idx])
frame = data.get_prob_table(frm_idx, bp_idx)
offsets = data.get_offset_map()
if (offsets is not None):
off_y = offsets[frm_idx, :, :, bp_idx, 1]
off_x = offsets[frm_idx, :, :, bp_idx, 0]
else:
off_x, off_y = None, None
bp_grp.attrs[
H5SubFrameKeys.INCLUDES_OFFSETS] = offsets is not None
if (self._threshold is not None):
sparse_y, sparse_x = np.nonzero(frame > self._threshold)
probs = frame[(sparse_y, sparse_x)]
# Check if we managed to strip out at least 2/3rds of the data, and if so write the frame using the
# sparse format. Otherwise it is actually more memory efficient to just store the entire frame...
if (len(frame.flat) >=
(len(sparse_y) *
DLCH5FSConstants.MIN_SPARSE_SAVING_FACTOR)):
bp_grp.attrs[H5SubFrameKeys.IS_STORED_SPARSE] = True
out_x = bp_grp.create_dataset(H5SubFrameKeys.X,
sparse_x.shape,
np.uint32)
out_y = bp_grp.create_dataset(H5SubFrameKeys.Y,
sparse_y.shape,
np.uint32)
out_prob = bp_grp.create_dataset(
H5SubFrameKeys.PROBS, probs.shape, np.float32)
out_x[:] = sparse_x
out_y[:] = sparse_y
out_prob[:] = probs
if (offsets is not None):
off_x, off_y = off_x[(sparse_y,
sparse_x)], off_y[(sparse_y,
sparse_x)]
out_off_x = bp_grp.create_dataset(
H5SubFrameKeys.OFFSET_X, off_x.shape,
np.float32)
out_off_y = bp_grp.create_dataset(
H5SubFrameKeys.OFFSET_Y, off_y.shape,
np.float32)
out_off_x[:] = off_x
out_off_y[:] = off_y
continue
# User has disabled sparse optimizations or they wasted more space, stash entire frame...
bp_grp.attrs[H5SubFrameKeys.IS_STORED_SPARSE] = False
prob_data = bp_grp.create_dataset(H5SubFrameKeys.PROBS,
frame.shape, np.float32)
prob_data[:] = frame
if (offsets is not None):
out_x = bp_grp.create_dataset(H5SubFrameKeys.OFFSET_X,
off_x.shape, np.float32)
out_y = bp_grp.create_dataset(H5SubFrameKeys.OFFSET_Y,
off_y.shape, np.float32)
out_x[:] = off_x
out_y[:] = off_y
current_frame_tmp += 1
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