tiena2cva commited on
Commit
b69e2d9
·
1 Parent(s): 7e61c38

perf(pose): read dense video frames sequentially

Browse files
src/pozify/steps/pose_landmarker.py CHANGED
@@ -46,11 +46,16 @@ def _iter_video_frames(
46
  if not capture.isOpened():
47
  return
48
  if sample_count is None:
49
- frame_indices = range(manifest.total_frames)
50
- else:
51
- frame_indices = sample_frame_indices(
52
- manifest.total_frames, min(sample_count, manifest.total_frames)
53
- )
 
 
 
 
 
54
  for frame_index in frame_indices:
55
  capture.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
56
  ok, frame = capture.read()
 
46
  if not capture.isOpened():
47
  return
48
  if sample_count is None:
49
+ for frame_index in range(manifest.total_frames):
50
+ ok, frame = capture.read()
51
+ if not ok or frame is None:
52
+ break
53
+ yield frame_index, frame
54
+ return
55
+
56
+ frame_indices = sample_frame_indices(
57
+ manifest.total_frames, min(sample_count, manifest.total_frames)
58
+ )
59
  for frame_index in frame_indices:
60
  capture.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
61
  ok, frame = capture.read()
tests/test_pose_steps.py CHANGED
@@ -8,6 +8,7 @@ from types import SimpleNamespace
8
  import sys
9
  import tempfile
10
  import unittest
 
11
 
12
  import cv2
13
  import numpy as np
@@ -129,6 +130,50 @@ class PoseStepTests(unittest.TestCase):
129
 
130
  self.assertEqual(len(sequence.frames), 130)
131
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
132
  def test_pose_cleaning_interpolates_smooths_and_adds_normalized_fields(
133
  self,
134
  ) -> None:
 
8
  import sys
9
  import tempfile
10
  import unittest
11
+ from unittest.mock import patch
12
 
13
  import cv2
14
  import numpy as np
 
130
 
131
  self.assertEqual(len(sequence.frames), 130)
132
 
133
+ def test_dense_video_iteration_reads_sequentially_without_reseeking(self) -> None:
134
+ class FakeCapture:
135
+ def __init__(self) -> None:
136
+ self.set_calls: list[tuple[int, int]] = []
137
+ self.read_calls = 0
138
+
139
+ def isOpened(self) -> bool:
140
+ return True
141
+
142
+ def set(self, prop: int, value: int) -> None:
143
+ self.set_calls.append((prop, value))
144
+
145
+ def read(self) -> tuple[bool, object | None]:
146
+ if self.read_calls >= 3:
147
+ return False, None
148
+ self.read_calls += 1
149
+ return True, object()
150
+
151
+ def release(self) -> None:
152
+ return None
153
+
154
+ capture = FakeCapture()
155
+ manifest = VideoManifest(
156
+ video_path="fake.mp4",
157
+ fps=30.0,
158
+ duration_sec=0.1,
159
+ total_frames=3,
160
+ sampled_frames=3,
161
+ width=640,
162
+ height=480,
163
+ codec="mp4v",
164
+ container="mp4",
165
+ brightness_mean=120.0,
166
+ blur_laplacian_var=100.0,
167
+ quality_warnings=[],
168
+ analysis_allowed=True,
169
+ )
170
+
171
+ with patch("pozify.steps.pose_landmarker.cv2.VideoCapture", return_value=capture):
172
+ frames = list(pose_landmarker._iter_video_frames(manifest, sample_count=None))
173
+
174
+ self.assertEqual([frame_index for frame_index, _ in frames], [0, 1, 2])
175
+ self.assertEqual(capture.set_calls, [])
176
+
177
  def test_pose_cleaning_interpolates_smooths_and_adds_normalized_fields(
178
  self,
179
  ) -> None: