Skip to content

Commit c618a7f

Browse files
committed
Handle zero-detections in SkipFrames counter.
- For zero detections, the age is still incremented - `_detections` attrubute is set to None
1 parent 6a2e70e commit c618a7f

1 file changed

Lines changed: 9 additions & 1 deletion

File tree

dlclive/pose_estimation_pytorch/runner.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,11 @@ def update(self, pose: torch.Tensor, w: int, h: int) -> None:
8282
self._detections = dict(boxes=bboxes, scores=torch.ones(num_det))
8383
self._age += 1
8484

85+
def next_frame(self) -> None:
86+
"""Increment the frame counter and set detections to None (to handle no detections)"""
87+
self._detections = None
88+
self._age += 1
89+
8590

8691
@dataclass
8792
class TopDownConfig:
@@ -199,7 +204,10 @@ def get_pose(self, frame: np.ndarray) -> np.ndarray:
199204
if self.single_animal or self.n_individuals == 1:
200205
zero_pose = np.zeros((self.n_bodyparts, 3))
201206
else:
202-
zero_pose = np.zeros((self.n_individuals, self.n_bodyparts, 3))
207+
zero_pose = np.zeros((self.n_individuals, self.n_bodyparts, 3))
208+
# Update skip_frames even when returning early to maintain frame counter
209+
if self.top_down_config.skip_frames is not None:
210+
self.top_down_config.skip_frames.next_frame()
203211
return zero_pose
204212

205213
tensor = frame_batch # still CHW, batched

0 commit comments

Comments
 (0)