From 3f0d594d9bc39638c0cca48209ee7ce03fca11a1 Mon Sep 17 00:00:00 2001 From: Aleksandr Movchan Date: Fri, 6 Dec 2024 12:14:24 +0000 Subject: [PATCH] Fix media_id generation in generate_frames function and update corresponding tests --- aana/integrations/external/decord.py | 2 +- aana/tests/units/test_frame_extraction.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/aana/integrations/external/decord.py b/aana/integrations/external/decord.py index 8de48107..424b1dc6 100644 --- a/aana/integrations/external/decord.py +++ b/aana/integrations/external/decord.py @@ -176,7 +176,7 @@ def generate_frames( batch_frames_array = video_reader.get_batch(batch).asnumpy() batch_frames = [] for frame_id, frame in enumerate(batch_frames_array): - img = Image(numpy=frame, media_id=f"{video.media_id}_frame_{frame_id}") + img = Image(numpy=frame, media_id=f"{video.media_id}_frame_{i+frame_id}") batch_frames.append(img) batch_timestamps = timestamps[i : i + batch_size] diff --git a/aana/tests/units/test_frame_extraction.py b/aana/tests/units/test_frame_extraction.py index 91cb7bdd..bc78355b 100644 --- a/aana/tests/units/test_frame_extraction.py +++ b/aana/tests/units/test_frame_extraction.py @@ -95,6 +95,7 @@ def test_generate_frames_success( gen_frame = generate_frames(video=video, params=params, batch_size=1) total_frames = 0 frame_ids = [] + frames_media_ids = [] for result in gen_frame: assert "frames" in result assert "timestamps" in result @@ -110,8 +111,12 @@ def test_generate_frames_success( total_frames += 1 assert result["duration"] == expected_duration frame_ids.extend(result["frame_ids"]) + frames_media_ids.extend([frame.media_id for frame in result["frames"]]) assert frame_ids == list(range(expected_num_frames)) + assert frames_media_ids == [ + f"{video.media_id}_frame_{frame_id}" for frame_id in range(expected_num_frames) + ] assert total_frames == expected_num_frames