Reading and writing the formats from your own code#

def __init__(self, h5f: h5py.File, num_kps=None, **create_kwargs):
    """
    Constructs an unsegmented pose writer
    """
    self.h5f = h5f
    self.num_kps = num_kps
    self.timeline_grp = self.h5f.create_group("/timeline", track_order=True)
    self.pose_grps: Dict[int, List[Any]] = {}
    self.start_frame = 0
    self.create_kwargs = create_kwargs

def add_pose(self, frame_num: int, pose_id: int, pose: ndarray):
    """
    Add a pose
    """
    pose_grp, data, indices, indptr, last_frame_num = self._pose_grp(
        pose_id, frame_num
    )
    new_rows = frame_num - last_frame_num
    add_empty_rows_grp(indptr, data, new_rows)
    new_data = []
    new_indices = []
    for limb_idx, limb in get_pose_nz(pose):
        new_data.append(limb)
        new_indices.append(limb_idx)
    if new_data:
        grow_ds(data, len(new_data))
        grow_ds(indices, len(new_indices))
        data[-len(new_data) :] = new_data
        indices[-len(new_indices) :] = new_indices
    self.pose_grps[pose_id][-1] = frame_num

def end_shot(self):
    """
    Ends a shot. This should only be called once at the end of writing.
    """
    self.timeline_grp.attrs["start_frame"] = self.start_frame
    timeline_last_frame_num = 0
    for pose_grp, data, indices, indptr, last_frame_num in self.pose_grps.values():
        add_empty_rows_grp(indptr, data, 1)
        pose_grp.attrs["end_frame"] = last_frame_num + 1
        timeline_last_frame_num = max(timeline_last_frame_num, last_frame_num)
    self.timeline_grp.attrs["end_frame"] = timeline_last_frame_num + 1

def start_shot(self, start_frame: int = 0):
    """
    Start a shot. This should only be called once at the beginning of writing.
    """
    self.start_frame = start_frame

def __init__(self, h5f: h5py.File, num_kps=None, **create_kwargs):
    """
    Constructs a shot segmented writer
    """
    self.h5f = h5f
    self.num_kps = num_kps
    self.h5f.create_group("/timeline", track_order=True)

    self.pose_data: Dict[int, Dict[int, ndarray]] = {}
    self.pose_starts: Dict[int, int] = {}
    self.pose_ends: Dict[int, int] = {}
    self.shot_idx = 0
    self.shot_start = 0
    self.last_frame = 0
    self.create_kwargs = create_kwargs

def add_pose(self, frame_num: int, pose_id: int, pose: ndarray):
    """
    Add a pose
    """
    if pose_id not in self.pose_data:
        self.pose_starts[pose_id] = frame_num
        self.pose_data[pose_id] = {}
    self.pose_data[pose_id][frame_num] = pose
    self.pose_ends[pose_id] = frame_num
    self.last_frame = frame_num

def end_shot(self):
    """
    End the current shot
    """
    shot_grp = self.h5f.create_group(
        f"/timeline/shot{self.shot_idx}", track_order=True
    )
    shot_grp.attrs["start_frame"] = self.shot_start
    shot_grp.attrs["end_frame"] = self.last_frame + 1
    for pose_id, poses in self.pose_data.items():
        data: List[ndarray] = []
        indices: List[int] = []
        indptr: List[int] = []
        pose_first_frame = self.pose_starts[pose_id]
        pose_last_frame = self.pose_ends[pose_id] + 1
        last_frame_num = pose_first_frame - 1

        def add_empty_rows(num_rows):
            for _ in range(num_rows):
                indptr.append(len(data))

        for frame_num, pose in poses.items():
            add_empty_rows(frame_num - last_frame_num)
            for limb_idx, limb in get_pose_nz(pose):
                data.append(limb)
                indices.append(limb_idx)
            last_frame_num = frame_num
        # Extra empty row to insert final nnz entry
        add_empty_rows(1)

        pose_group = create_csr(
            self.h5f,
            f"/timeline/shot{self.shot_idx}/pose{pose_id}",
            self.num_kps,
            data=data,
            indices=indices,
            indptr=indptr,
            **self.create_kwargs,
        )
        pose_group.attrs["start_frame"] = pose_first_frame
        pose_group.attrs["end_frame"] = pose_last_frame
    self.pose_data = {}
    self.shot_idx += 1
    self.shot_start = self.last_frame + 1

def register_frame(self, frame_num: int):
    """
    Register frame_num as existing within the current shot
    """
    self.last_frame = frame_num

def start_shot(self, start_frame=None):
    """
    Start a new shot
    """
    if start_frame is not None:
        self.shot_start = start_frame

def __init__(self, h5f: h5py.File, bundle_cls=DumpReaderPoseBundle, infinite=False):
    """
    Constructs the reader from a HDF5 file and. If `infinite` is True, all
    shot iterators will terminate with a final empty shot which infinitely
    yields empty pose bundles.
    """
    self.h5f = h5f
    self.limbs = self.h5f.attrs["limbs"]
    assert self.h5f.attrs["fmt_type"] == "trackshots"
    self.mk_bundle = partial(bundle_cls, cls=POSE_CLASSES[self.h5f.attrs["mode"]])
    self.empty_bundle = self.mk_bundle({})
    self.infinite = infinite

def __iter__(self):
    """
    Returns a shot iterator of all shots
    """
    for shot_idx, shot_name, start_frame, end_frame, mk_shot in self._iter():
        yield mk_shot()

def iter_from_frame(self, start_frame: int):
    """
    Returns a shot iterator starting from frame 0-index `start_frame`.
    """
    started = False
    for (
        shot_idx,
        shot_name,
        shot_start_frame,
        shot_end_frame,
        mk_shot,
    ) in self._iter(start_frame):
        if started:
            yield mk_shot()
        elif shot_start_frame <= start_frame and (
            shot_end_frame is None or start_frame < shot_end_frame
        ):
            yield mk_shot()
            started = True

def iter_from_shot(self, start_shot: int):
    """
    Returns a shot iterator starting from shot 0-index `start_shot`.
    """
    for shot_idx, shot_name, start_frame, end_frame, mk_shot in self._iter():
        if shot_idx >= start_shot:
            yield mk_shot()

def __iter__(self):
    """
    Returns a pose bundle iterator of all frames.
    """
    return iter(self.shot_reader)

def iter_from(self, start_frame: int):
    """
    Returns a pose bundle iterator starting at frame `start_frame`.
    """
    return self.shot_reader.iter_from(start_frame)

def __iter__(self):
    """
    Returns a pose bundle iterator of all frames.
    """
    return self.iter_from(self.start_frame)

def iter_from(self, start_frame):
    """
    Returns a pose bundle iterator starting at frame `start_frame`.
    """
    for frame in range(start_frame, self.end_frame):
        bundle = {}
        for pose_num, start_frame, end_frame, sparse_pose in self.poses:
            if start_frame <= frame < end_frame:
                row_num = frame - start_frame
                bundle[pose_num] = sparse_pose.get_row(row_num)
        yield self.mk_bundle(bundle)