Merge pull request #303 from clamsproject/302-vdh-destroy-frame-list

Fix the issue of extract_frames_as_images destory input frame
clamsproject · Aug 21, 2024 · af83e13 · af83e13
2 parents 2e17432 + c94106c
commit af83e13
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 4 deletions.
diff --git a/mmif/utils/video_document_helper.py b/mmif/utils/video_document_helper.py
@@ -84,18 +84,19 @@ def extract_frames_as_images(video_document: Document, framenums: List[int], as_
     video = capture(video_document)
     cur_f = 0
     tot_fcount = video_document.get_property(FRAMECOUNT_DOCPROP_KEY)
+    framenums_copy = framenums.copy()
     while True:
-        if not framenums or cur_f > tot_fcount:
+        if not framenums_copy or cur_f > tot_fcount:
             break
         ret, frame = video.read()
-        if cur_f == framenums[0]:
+        if cur_f == framenums_copy[0]:
             if not ret:
                 sec = convert(cur_f, 'f', 's', video_document.get_property(FPS_DOCPROP_KEY))
                 warnings.warn(f'Frame #{cur_f} ({sec}s) could not be read from the video {video_document.id}.')
                 cur_f += 1
                 continue
             frames.append(Image.fromarray(frame[:, :, ::-1]) if as_PIL else frame)
-            framenums.pop(0)
+            framenums_copy.pop(0)
         cur_f += 1
     return frames
 

diff --git a/tests/black-2997fps.mp4 b/tests/black-2997fps.mp4
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -1,3 +1,4 @@
+import pathlib
 import unittest
 
 import pytest
@@ -31,7 +32,7 @@ def setUp(self):
             "properties": {
                 "mime": "video",
                 "id": "d1",
-                "location": "file:///home/snewman/Documents/test_vid.mp4"
+                "location": f"file://{pathlib.Path(__file__).parent}/black-2997fps.mp4"
             }
         })
         self.video_doc.add_property('fps', self.fps)
@@ -105,6 +106,23 @@ def test_convert_timeframe(self):
         for times in zip((3.337, 6.674), vdh.convert_timeframe(self.mmif_obj, timeframe_ann, 's')):
             self.assertAlmostEqual(*times, places=0)
 
+    def test_extract_frames_as_images(self):
+        frame_list = [5, 10, 15]
+        target_images = vdh.extract_frames_as_images(self.video_doc, frame_list, as_PIL=False)
+        self.assertEqual(3, len(target_images))
+        # check if the extract_frames_as_images destroy the input frame list
+        self.assertEqual(3, len(frame_list))
+        # check return empty list if the frame list is empty
+        empty_flist = []
+        empty_target_images = vdh.extract_frames_as_images(self.video_doc, empty_flist, as_PIL=False)
+        self.assertEqual([], empty_target_images)
+        # check there is an error if there is a frame in the list that does not exist
+        tot_fcount = self.video_doc.get_property('frameCount')
+        frame_list.append(tot_fcount + 1)
+        new_target_images = vdh.extract_frames_as_images(self.video_doc, frame_list, as_PIL=False)
+        self.assertEqual(4, len(frame_list))
+        self.assertEqual(3, len(new_target_images))
+
 
 class TestSequenceHelper(unittest.TestCase):