Skip to content

Commit

Permalink
feat(lib): smarter files reversing
Browse files Browse the repository at this point in the history
Implement a smarter generation of reversed files by splitting the video into smaller segments.

Closes #434
  • Loading branch information
jeertmans committed May 22, 2024
1 parent b080739 commit 4fc2454
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 4 deletions.
11 changes: 9 additions & 2 deletions manim_slides/slide/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -475,10 +475,11 @@ def _save_slides(self, use_cache: bool = True) -> None:

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",
):
slide_files = files[pre_slide_config.slides_slice]

Expand All @@ -492,7 +493,13 @@ def _save_slides(self, use_cache: bool = True) -> None:

# We only reverse video if it was not present
if not use_cache or not rev_file.exists():
reverse_video_file(dst_file, rev_file)
reverse_video_file(
dst_file,
rev_file,
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(
Expand Down
68 changes: 66 additions & 2 deletions manim_slides/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -120,8 +124,68 @@ 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: float = 1,
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 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:
logger.debug(

Check warning on line 151 in manim_slides/utils.py

View check run for this annotation

Codecov / codecov/patch

manim_slides/utils.py#L151

Added line #L151 was not covered by tests
f"Could not determine duration of {src}, falling back to segmentation."
)

with tempfile.TemporaryDirectory() as tmpdirname:
tmpdir = Path(tmpdirname)
with av.open(
str(tmpdir / "%04d.mp4"),
"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(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)

0 comments on commit 4fc2454

Please sign in to comment.