Merge pull request #40 from IRFM/save-part-of-movie-fix

to_h264 method doesn't work properly when saving partial movies

src/python/librir/video_io/IRMovie.py

@@ -566,6 +566,8 @@ def to_h264(
             count = self.images
         if start_img + count > self.images:
             count = self.images - start_img
+        if count == 0:
+            raise RuntimeError("No images in selected range to save")
 
         logger.info(
             "Start saving in {} from {} to {}".format(
@@ -576,6 +578,10 @@ def to_h264(
         # retrieve attributes
         if attrs is None:
             attrs = self.attributes
+        if (frame_attributes is not None) and (len(frame_attributes) != count):
+            raise RuntimeError(
+                "Given frame attributes are not equal to the number of saved images"
+            )
 
         # adding custom attribute --> must be a dict[str]=str
         # attrs.update(self.additional_attributes)
@@ -609,7 +615,7 @@ def to_h264(
                 _frame_attributes = (
                     self.frame_attributes
                     if frame_attributes is None
-                    else frame_attributes[i]
+                    else frame_attributes[saved]
                 )
 
                 s.add_image(img, times[i] * 1e9, attributes=_frame_attributes)

tests/python/conftest.py

@@ -125,27 +125,35 @@ def forge_mock_metadata_in_last_line_for_2d_arrays(array):
 #     return [IRMovie.from_numpy_array(arr) for arr in VALID_UNIFORM_NUMPY_ARRAYS]
 
 
-@pytest.fixture(scope="session", params=VALID_SHAPES)
+@pytest.fixture(
+    scope="session", params=VALID_SHAPES, ids=[str(s) for s in VALID_SHAPES]
+)
 def array(request):
     shape = request.param
     arr = generate_mock_movie_data_uniform(*shape)
 
     yield arr
 
 
-@pytest.fixture(scope="session", params=VALID_3D_SHAPES)
+@pytest.fixture(
+    scope="session", params=VALID_3D_SHAPES, ids=[str(s) for s in VALID_3D_SHAPES]
+)
 def valid_3d_array(request):
     shape = request.param
     yield generate_mock_movie_data_uniform(*shape)
 
 
-@pytest.fixture(scope="session", params=INVALID_SHAPES)
+@pytest.fixture(
+    scope="session", params=INVALID_SHAPES, ids=[str(s) for s in INVALID_SHAPES]
+)
 def bad_array(request):
     shape = request.param
     yield generate_mock_movie_data_uniform(*shape)
 
 
-@pytest.fixture(scope="session", params=VALID_2D_SHAPES)
+@pytest.fixture(
+    scope="session", params=VALID_2D_SHAPES, ids=[str(s) for s in VALID_2D_SHAPES]
+)
 def valid_2D_array(request):
     shape = request.param
     arr = generate_mock_movie_data_uniform(*shape)

tests/python/test_IRMovie.py

@@ -1,4 +1,6 @@
+from pathlib import Path
 import random
+from typing import Dict
 from librir.video_io.utils import is_ir_file_corrupted, split_rush
 
 import numpy as np
@@ -186,3 +188,63 @@ def test_flip_camera_calibration_when_no_calibration(movie: IRMovie):
         movie.flip_calibration(False, True)
     with pytest.raises(RuntimeError):
         movie.flip_calibration(True, True)
+
+
+@pytest.mark.parametrize("attrs", ({}, {"additional": 123}), ids=("none", "additional"))
+def test_save_subset_of_movie_to_h264(movie: IRMovie, attrs: Dict[str, int]):
+    dest_filename = Path(f"{movie.filename}_subset.h264")
+    count = 1 if movie.images == 1 else movie.images // 2
+    frame_attributes = [{"additional_frame_attribute": i} for i in range(movie.images)]
+
+    with pytest.raises(RuntimeError):
+        movie.to_h264(
+            dest_filename,
+            start_img=0,
+            count=movie.images // 2,
+            attrs=attrs,
+            frame_attributes=frame_attributes,
+        )
+    with pytest.raises(RuntimeError):
+        movie.to_h264(
+            dest_filename,
+            start_img=movie.images // 2,
+            count=movie.images // 2,
+            attrs=attrs,
+            frame_attributes=frame_attributes,
+        )
+    with pytest.raises(RuntimeError):
+        movie.to_h264(
+            dest_filename,
+            start_img=movie.images // 2,
+            count=0,
+            attrs=attrs,
+            frame_attributes=frame_attributes,
+        )
+    with pytest.raises(RuntimeError):
+        movie.to_h264(
+            dest_filename,
+            start_img=movie.images // 2,
+            count=0,
+            attrs=attrs,
+            frame_attributes=None,
+        )
+    # assert not dest_filename.exists()
+
+    movie.to_h264(
+        dest_filename,
+        start_img=0,
+        count=count,
+        attrs=attrs,
+        frame_attributes=[{"additional_frame_attribute": i} for i in range(count)],
+    )
+    # assert dest_filename.exists()
+    expected_attrs = {k: str(v).encode() for k, v in attrs.items()}
+    expected_attrs["GOP"] = movie.attributes["GOP"]
+    # expected_attrs.update({k: str(v).encode() for k, v in attrs.items()})
+
+    with IRMovie.from_filename(dest_filename) as subset_movie:
+        # subset_movie.load_pos(0)
+        assert subset_movie.attributes == expected_attrs
+        assert len(subset_movie.frames_attributes) == subset_movie.images
+
+    # dest_filename.unlink()