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