diff --git a/manim/scene/scene_file_writer.py b/manim/scene/scene_file_writer.py index eed0db1e53..49d8492a9a 100644 --- a/manim/scene/scene_file_writer.py +++ b/manim/scene/scene_file_writer.py @@ -7,6 +7,8 @@ import json import shutil from pathlib import Path +from queue import Queue +from threading import Thread from typing import TYPE_CHECKING, Any import av @@ -16,6 +18,7 @@ from pydub import AudioSegment from manim import __version__ +from manim.typing import PixelArray from .. import config, logger from .._config.logger_utils import set_file_logger @@ -359,6 +362,33 @@ def end_animation(self, allow_write: bool = False): if write_to_movie() and allow_write: self.close_partial_movie_stream() + def listen_and_write(self): + """ + For internal use only: blocks until new frame is available on the queue. + """ + while True: + num_frames, frame_data = self.queue.get() + if frame_data is None: + break + + self.encode_and_write_frame(frame_data, num_frames) + + def encode_and_write_frame(self, frame: PixelArray, num_frames: int) -> None: + """ + For internal use only: takes a given frame in ``np.ndarray`` format and + write it to the stream + """ + for _ in range(num_frames): + # Notes: precomputing reusing packets does not work! + # I.e., you cannot do `packets = encode(...)` + # and reuse it, as it seems that `mux(...)` + # consumes the packet. + # The same issue applies for `av_frame`, + # reusing it renders weird-looking frames. + av_frame = av.VideoFrame.from_ndarray(frame, format="rgba") + for packet in self.video_stream.encode(av_frame): + self.video_container.mux(packet) + def write_frame( self, frame_or_renderer: np.ndarray | OpenGLRenderer, num_frames: int = 1 ): @@ -379,16 +409,9 @@ def write_frame( if config.renderer == RendererType.OPENGL else frame_or_renderer ) - for _ in range(num_frames): - # Notes: precomputing reusing packets does not work! - # I.e., you cannot do `packets = encode(...)` - # and reuse it, as it seems that `mux(...)` - # consumes the packet. - # The same issue applies for `av_frame`, - # reusing it renders weird-looking frames. - av_frame = av.VideoFrame.from_ndarray(frame, format="rgba") - for packet in self.video_stream.encode(av_frame): - self.video_container.mux(packet) + + msg = (num_frames, frame) + self.queue.put(msg) if is_png_format() and not config["dry_run"]: image: Image = ( @@ -430,7 +453,7 @@ def save_final_image(self, image: np.ndarray): image.save(self.image_file_path) self.print_file_ready_message(self.image_file_path) - def finish(self): + def finish(self) -> None: """ Finishes writing to the FFMPEG buffer or writing images to output directory. @@ -440,8 +463,6 @@ def finish(self): frame in the default image directory. """ if write_to_movie(): - if hasattr(self, "writing_process"): - self.writing_process.terminate() self.combine_to_movie() if config.save_sections: self.combine_to_section_videos() @@ -455,7 +476,7 @@ def finish(self): if self.subcaptions: self.write_subcaption_file() - def open_partial_movie_stream(self, file_path=None): + def open_partial_movie_stream(self, file_path=None) -> None: """Open a container holding a video stream. This is used internally by Manim initialize the container holding @@ -499,13 +520,20 @@ def open_partial_movie_stream(self, file_path=None): self.video_container = video_container self.video_stream = stream - def close_partial_movie_stream(self): + self.queue: Queue[tuple[int, PixelArray | None]] = Queue() + self.writer_thread = Thread(target=self.listen_and_write, args=()) + self.writer_thread.start() + + def close_partial_movie_stream(self) -> None: """Close the currently opened video container. Used internally by Manim to first flush the remaining packages in the video stream holding a partial file, and then close the corresponding container. """ + self.queue.put((-1, None)) + self.writer_thread.join() + for packet in self.video_stream.encode(): self.video_container.mux(packet)