feat(lib): smarter files reversing (#439)

jeertmans · web-flow · commit e911ec30964b · 2025-01-28T23:05:59.000Z
* feat(lib): smarter files reversing Implement a smarter generation of reversed files by splitting the video into smaller segments. Closes #434 * chore(lib): change default length * chore: use suffix * chore(docs): update * chore(docs): fix docs and add changelog entry * chore(tests): coverage * chore(docs): typo
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,6 +10,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 (unreleased)=
 ## [Unreleased](https://github.com/jeertmans/manim-slides/compare/v5.4.2...HEAD)
 
+(unreleased-added)=
+### Added
+
+- Added `max_duration_before_split_reverse` and `num_processes` class variables.
+  [#439](https://github.com/jeertmans/manim-slides/pull/439)
+
+(unreleased-changed)=
+### Changed
+
+- Automatically split large video animations into smaller chunks
+  for lightweight (and potentially faster) reversed animations generation.
+  [#439](https://github.com/jeertmans/manim-slides/pull/439)
+
 (v5.4.2)=
 ## [v5.4.2](https://github.com/jeertmans/manim-slides/compare/v5.4.1...v5.4.2)
 
diff --git a/manim_slides/slide/base.py b/manim_slides/slide/base.py
@@ -36,6 +36,8 @@ class BaseSlide:
     disable_caching: bool = False
     flush_cache: bool = False
     skip_reversing: bool = False
+    max_duration_before_split_reverse: float | None = 4.0
+    num_processes: int | None = None
 
     def __init__(
         self, *args: Any, output_folder: Path = FOLDER_PATH, **kwargs: Any
@@ -530,10 +532,11 @@ def _save_slides(
 
         for pre_slide_config in tqdm(
             self._slides,
-            desc=f"Concatenating animation files to '{scene_files_folder}' and generating reversed animations",
+            desc=f"Concatenating animations to '{scene_files_folder}' and generating reversed animations",
             leave=self._leave_progress_bar,
             ascii=True if platform.system() == "Windows" else None,
             disable=not self._show_progress_bar,
+            unit=" slides",
         ):
             if pre_slide_config.skip_animations:
                 continue
@@ -557,7 +560,15 @@ def _save_slides(
                 if skip_reversing:
                     rev_file = dst_file
                 else:
-                    reverse_video_file(dst_file, rev_file)
+                    reverse_video_file(
+                        dst_file,
+                        rev_file,
+                        max_segment_duration=self.max_duration_before_split_reverse,
+                        num_processes=self.num_processes,
+                        leave=self._leave_progress_bar,
+                        ascii=True if platform.system() == "Windows" else None,
+                        disable=not self._show_progress_bar,
+                    )
 
             slides.append(
                 SlideConfig.from_pre_slide_config_and_files(
diff --git a/manim_slides/slide/manim.py b/manim_slides/slide/manim.py
@@ -11,7 +11,7 @@
 
 class Slide(BaseSlide, Scene):  # type: ignore[misc]
     """
-    Inherits from :class:`Scene<manim.scene.scene.Scene>` and provide necessary tools
+    Inherits from :class:`Scene<manim.scene.scene.Scene>` and provides necessary tools
     for slides rendering.
 
     :param args: Positional arguments passed to scene object.
@@ -20,15 +20,26 @@ class Slide(BaseSlide, Scene):  # type: ignore[misc]
     :cvar bool disable_caching: :data:`False`: Whether to disable the use of
         cached animation files.
     :cvar bool flush_cache: :data:`False`: Whether to flush the cache.
-
         Unlike with Manim, flushing is performed before rendering.
     :cvar bool skip_reversing: :data:`False`: Whether to generate reversed animations.
-
         If set to :data:`False`, and no cached reversed animation
         exists (or caching is disabled) for a given slide,
         then the reversed animation will be simply the same
         as the original one, i.e., ``rev_file = file``,
         for the current slide config.
+    :cvar typing.Optional[float] max_duration_before_split_reverse: :data:`4.0`: Maximum duration
+        before of a video animation before it is reversed by splitting the file into smaller chunks.
+        Generating reversed animations can require an important amount of
+        memory (because the whole video needs to be kept in memory),
+        and splitting the video into multiple chunks usually speeds
+        up the process (because it can be done in parallel) while taking
+        less memory.
+        Set this to :data:`None` to disable splitting the file into chunks.
+    :cvar typing.Optional[int] num_processes: :data:`None`: Number of processes
+        to use for parallelizable operations.
+        If :data:`None`, defaults to :func:`os.process_cpu_count`.
+        This is currently used when generating reversed animations, and can
+        increase memory consumption.
     """
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
diff --git a/manim_slides/utils.py b/manim_slides/utils.py
@@ -2,9 +2,12 @@
 import os
 import tempfile
 from collections.abc import Iterator
+from multiprocessing import Pool
 from pathlib import Path
+from typing import Any, Optional
 
 import av
+from tqdm import tqdm
 
 from .logger import logger
 
@@ -89,8 +92,9 @@ def link_nodes(*nodes: av.filter.context.FilterContext) -> None:
         c.link_to(n)
 
 
-def reverse_video_file(src: Path, dest: Path) -> None:
+def reverse_video_file_in_one_chunk(src_and_dest: tuple[Path, Path]) -> None:
     """Reverses a video file, writing the result to `dest`."""
+    src, dest = src_and_dest
     with (
         av.open(str(src)) as input_container,
         av.open(str(dest), mode="w") as output_container,
@@ -120,8 +124,70 @@ def reverse_video_file(src: Path, dest: Path) -> None:
 
         for _ in range(frames_count):
             frame = graph.pull()
-            frame.pict_type = 5  # Otherwise we get a warning saying it is changed
+            frame.pict_type = "NONE"  # Otherwise we get a warning saying it is changed
             output_container.mux(output_stream.encode(frame))
 
         for packet in output_stream.encode():
             output_container.mux(packet)
+
+
+def reverse_video_file(
+    src: Path,
+    dest: Path,
+    max_segment_duration: Optional[float] = 4.0,
+    num_processes: Optional[int] = None,
+    **tqdm_kwargs: Any,
+) -> None:
+    """Reverses a video file, writing the result to `dest`."""
+    with av.open(str(src)) as input_container:  # Fast path if file is short enough
+        input_stream = input_container.streams.video[0]
+        if max_segment_duration is None:
+            return reverse_video_file_in_one_chunk((src, dest))
+        elif input_stream.duration:
+            if (
+                float(input_stream.duration * input_stream.time_base)
+                <= max_segment_duration
+            ):
+                return reverse_video_file_in_one_chunk((src, dest))
+        else:  # pragma: no cover
+            logger.debug(
+                f"Could not determine duration of {src}, falling back to segmentation."
+            )
+
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            tmpdir = Path(tmpdirname)
+            with av.open(
+                str(tmpdir / f"%04d.{src.suffix}"),
+                "w",
+                format="segment",
+                options={"segment_time": str(max_segment_duration)},
+            ) as output_container:
+                output_stream = output_container.add_stream(
+                    template=input_stream,
+                )
+
+                for packet in input_container.demux(input_stream):
+                    if packet.dts is None:
+                        continue
+
+                    packet.stream = output_stream
+                    output_container.mux(packet)
+
+            src_files = list(tmpdir.iterdir())
+            rev_files = [
+                src_file.with_stem("rev_" + src_file.stem) for src_file in src_files
+            ]
+
+            with Pool(num_processes, maxtasksperchild=1) as pool:
+                for _ in tqdm(
+                    pool.imap_unordered(
+                        reverse_video_file_in_one_chunk, zip(src_files, rev_files)
+                    ),
+                    desc="Reversing large file by cutting it in segments",
+                    total=len(src_files),
+                    unit=" files",
+                    **tqdm_kwargs,
+                ):
+                    pass  # We just consume the iterator
+
+            concatenate_video_files(rev_files[::-1], dest)
diff --git a/tests/test_slide.py b/tests/test_slide.py
@@ -314,6 +314,26 @@ def construct(self) -> None:
                 self.play(dot.animate.move_to(LEFT))
                 self.play(dot.animate.move_to(DOWN))
 
+    def test_split_reverse(self) -> None:
+        @assert_renders
+        class _(CESlide):
+            max_duration_before_split_reverse = 3.0
+
+            def construct(self) -> None:
+                self.wait(2.0)
+                for _ in range(3):
+                    self.next_slide()
+                    self.wait(10.0)
+
+        @assert_renders
+        class __(CESlide):
+            max_duration_before_split_reverse = None
+
+            def construct(self) -> None:
+                self.wait(5.0)
+                self.next_slide()
+                self.wait(5.0)
+
     def test_file_too_long(self) -> None:
         @assert_renders
         class _(CESlide):
diff --git a/uv.lock b/uv.lock
