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 _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 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