Fix frame_ids generation in generate_frames function and update tests…

… accordingly
mobiusml · Dec 6, 2024 · 696ae56 · 696ae56
1 parent c263f06
commit 696ae56
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/aana/integrations/external/decord.py b/aana/integrations/external/decord.py
@@ -182,7 +182,7 @@ def generate_frames(
         batch_timestamps = timestamps[i : i + batch_size]
         yield FramesDict(
             frames=batch_frames,
-            frame_ids=list(range(len(batch_frames))),
+            frame_ids=list(range(i, i + len(batch_frames))),
             timestamps=batch_timestamps,
             duration=duration,
         )

diff --git a/aana/tests/units/test_frame_extraction.py b/aana/tests/units/test_frame_extraction.py
@@ -41,6 +41,7 @@ def test_extract_frames_success(
         assert isinstance(result["frames"][0], Image)
     assert result["duration"] == expected_duration
     assert len(result["frames"]) == expected_num_frames
+    assert result["frame_ids"] == list(range(expected_num_frames))
     assert len(result["timestamps"]) == expected_num_frames
 
 
@@ -93,6 +94,7 @@ def test_generate_frames_success(
     params = VideoParams(extract_fps=extract_fps, fast_mode_enabled=fast_mode_enabled)
     gen_frame = generate_frames(video=video, params=params, batch_size=1)
     total_frames = 0
+    frame_ids = []
     for result in gen_frame:
         assert "frames" in result
         assert "timestamps" in result
@@ -107,7 +109,9 @@ def test_generate_frames_success(
             assert len(result["timestamps"]) == 1  # batch_size = 1
             total_frames += 1
         assert result["duration"] == expected_duration
+        frame_ids.extend(result["frame_ids"])
 
+    assert frame_ids == list(range(expected_num_frames))
     assert total_frames == expected_num_frames