Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(lib): smarter files reversing #439

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions manim_slides/slide/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -516,10 +516,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",
):
slide_files = files[pre_slide_config.slides_slice]

Expand All @@ -536,7 +537,13 @@ 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,
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 @@
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 @@

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 = 4,
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 / 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(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)
Loading