public-mirrors/3b1b-manim
1
0
Fork 0
mirror of https://github.com/3b1b/manim.git synced 2025-08-21 05:44:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Clean up SceneFileWriter

Browse source 
It had a number of vestigial functions no longer used, and some setup that could be made more organized.
This commit is contained in:
Grant Sanderson 2024-12-11 17:16:02 -06:00
parent 7f9a2cf487
commit 390064f6bd
1 changed files with 38 additions and 108 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

144
manimlib/scene/scene_file_writer.py
View file

@ -32,17 +32,15 @@ class SceneFileWriter(object):
self, self,
scene: Scene, scene: Scene,
write_to_movie: bool = False, write_to_movie: bool = False,
break_into_partial_movies: bool = False, subdivide_output: bool = False,
save_pngs: bool = False, # TODO, this currently does nothing
png_mode: str = "RGBA", png_mode: str = "RGBA",
save_last_frame: bool = False, save_last_frame: bool = False,
movie_file_extension: str = ".mp4", movie_file_extension: str = ".mp4",
# What python file is generating this scene # What python file is generating this scene
input_file_path: str = "", input_file_path: str = "",
# Where should this be written # Where should this be written
output_directory: str | None = None, output_directory: str = "",
file_name: str | None = None, file_name: str | None = None,
subdirectory_for_videos: bool = False,
open_file_upon_completion: bool = False, open_file_upon_completion: bool = False,
show_file_location_upon_completion: bool = False, show_file_location_upon_completion: bool = False,
quiet: bool = False, quiet: bool = False,
@ -57,8 +55,7 @@ class SceneFileWriter(object):
): ):
self.scene: Scene = scene self.scene: Scene = scene
self.write_to_movie = write_to_movie self.write_to_movie = write_to_movie
self.break_into_partial_movies = break_into_partial_movies self.subdivide_output = subdivide_output
self.save_pngs = save_pngs
self.png_mode = png_mode self.png_mode = png_mode
self.save_last_frame = save_last_frame self.save_last_frame = save_last_frame
self.movie_file_extension = movie_file_extension self.movie_file_extension = movie_file_extension
@ -66,7 +63,6 @@ class SceneFileWriter(object):
self.output_directory = output_directory self.output_directory = output_directory
self.file_name = file_name self.file_name = file_name
self.open_file_upon_completion = open_file_upon_completion self.open_file_upon_completion = open_file_upon_completion
self.subdirectory_for_videos = subdirectory_for_videos
self.show_file_location_upon_completion = show_file_location_upon_completion self.show_file_location_upon_completion = show_file_location_upon_completion
self.quiet = quiet self.quiet = quiet
self.total_frames = total_frames self.total_frames = total_frames
@ -81,40 +77,39 @@ class SceneFileWriter(object):
self.writing_process: sp.Popen | None = None self.writing_process: sp.Popen | None = None
self.progress_display: ProgressDisplay | None = None self.progress_display: ProgressDisplay | None = None
self.ended_with_interrupt: bool = False self.ended_with_interrupt: bool = False
self.init_output_directories() self.init_output_directories()
self.init_audio() self.init_audio()
# Output directories and files # Output directories and files
def init_output_directories(self) -> None: def init_output_directories(self) -> None:
out_dir = self.output_directory or ""
scene_name = self.file_name or self.get_default_scene_name()
if self.save_last_frame: if self.save_last_frame:
image_dir = guarantee_existence(os.path.join(out_dir, "images")) self.image_file_path = self.init_image_file_path()
image_file = add_extension_if_not_present(scene_name, ".png")
self.image_file_path = os.path.join(image_dir, image_file)
if self.write_to_movie: if self.write_to_movie:
if self.subdirectory_for_videos: self.movie_file_path = self.init_movie_file_path()
movie_dir = guarantee_existence(os.path.join(out_dir, "videos")) if self.subdivide_output:
else: self.partial_movie_directory = self.init_partial_movie_directory()
movie_dir = guarantee_existence(out_dir)
movie_file = add_extension_if_not_present(scene_name, self.movie_file_extension) def init_image_file_path(self) -> Path:
self.movie_file_path = os.path.join(movie_dir, movie_file) return self.get_output_file_rootname().with_suffix(".png")
if self.break_into_partial_movies:
self.partial_movie_directory = guarantee_existence(os.path.join( def init_movie_file_path(self) -> Path:
movie_dir, "partial_movie_files", scene_name, return self.get_output_file_rootname().with_suffix(self.movie_file_extension)
))
# A place to save mobjects def init_partial_movie_directory(self):
self.saved_mobject_directory = os.path.join( return guarantee_existence(self.get_output_file_rootname())
out_dir, "mobjects", str(self.scene)
def get_output_file_rootname(self) -> Path:
return Path(
guarantee_existence(self.output_directory),
self.get_output_file_name()
) )
def get_default_module_directory(self) -> str: def get_output_file_name(self) -> str:
path, _ = os.path.splitext(self.input_file_path) if self.file_name:
if path.startswith("_"): return self.file_name
path = path[1:] # Otherwise, use the name of the scene, potentially
return path # appending animation numbers
def get_default_scene_name(self) -> str:
name = str(self.scene) name = str(self.scene)
saan = self.scene.start_at_animation_number saan = self.scene.start_at_animation_number
eaan = self.scene.end_at_animation_number eaan = self.scene.end_at_animation_number
@ -124,26 +119,13 @@ class SceneFileWriter(object):
name += f"_{eaan}" name += f"_{eaan}"
return name return name
def get_resolution_directory(self) -> str:
pixel_height = self.scene.camera.pixel_height
fps = self.scene.camera.fps
return "{}p{}".format(
pixel_height, fps
)
# Directory getters # Directory getters
def get_image_file_path(self) -> str: def get_image_file_path(self) -> str:
return self.image_file_path return self.image_file_path
def get_next_partial_movie_path(self) -> str: def get_next_partial_movie_path(self) -> str:
result = os.path.join( result = Path(self.partial_movie_directory, f"{self.scene.num_plays:05}")
self.partial_movie_directory, return result.with_suffix(self.movie_file_extension)
"{:05}{}".format(
self.scene.num_plays,
self.movie_file_extension,
)
)
return result
def get_movie_file_path(self) -> str: def get_movie_file_path(self) -> str:
return self.movie_file_path return self.movie_file_path
@ -199,22 +181,19 @@ class SceneFileWriter(object):
# Writers # Writers
def begin(self) -> None: def begin(self) -> None:
if not self.break_into_partial_movies and self.write_to_movie: if not self.subdivide_output and self.write_to_movie:
self.open_movie_pipe(self.get_movie_file_path()) self.open_movie_pipe(self.get_movie_file_path())
def begin_animation(self) -> None: def begin_animation(self) -> None:
if self.break_into_partial_movies and self.write_to_movie: if self.subdivide_output and self.write_to_movie:
self.open_movie_pipe(self.get_next_partial_movie_path()) self.open_movie_pipe(self.get_next_partial_movie_path())
def end_animation(self) -> None: def end_animation(self) -> None:
if self.break_into_partial_movies and self.write_to_movie: if self.subdivide_output and self.write_to_movie:
self.close_movie_pipe() self.close_movie_pipe()
def finish(self) -> None: def finish(self) -> None:
if self.write_to_movie: if not self.subdivide_output and self.write_to_movie:
if self.break_into_partial_movies:
self.combine_movie_files()
else:
self.close_movie_pipe() self.close_movie_pipe()
if self.includes_sound: if self.includes_sound:
self.add_sound_to_video() self.add_sound_to_video()
@ -234,7 +213,6 @@ class SceneFileWriter(object):
width, height = self.scene.camera.get_pixel_shape() width, height = self.scene.camera.get_pixel_shape()
vf_arg = 'vflip' vf_arg = 'vflip'
# if self.pixel_format.startswith("yuv"):
vf_arg += f',eq=saturation={self.saturation}:gamma={self.gamma}' vf_arg += f',eq=saturation={self.saturation}:gamma={self.gamma}'
command = [ command = [
@ -246,7 +224,7 @@ class SceneFileWriter(object):
'-r', str(fps), # frames per second '-r', str(fps), # frames per second
'-i', '-', # The input comes from a pipe '-i', '-', # The input comes from a pipe
'-vf', vf_arg, '-vf', vf_arg,
'-an', # Tells FFMPEG not to expect any audio '-an', # Tells ffmpeg not to expect any audio
'-loglevel', 'error', '-loglevel', 'error',
] ]
if self.video_codec: if self.video_codec:
@ -273,8 +251,8 @@ class SceneFileWriter(object):
movie_path = Path(self.get_movie_file_path()) movie_path = Path(self.get_movie_file_path())
scene_name = movie_path.stem scene_name = movie_path.stem
insert_dir = Path(movie_path.parent, "inserts") insert_dir = Path(movie_path.parent, "inserts")
guarantee_existence(str(insert_dir)) guarantee_existence(insert_dir)
return Path(insert_dir, f"{scene_name}_{index}{movie_path.suffix}") return Path(insert_dir, f"{scene_name}_{index}").with_suffix(self.movie_file_extension)
def begin_insert(self): def begin_insert(self):
# Begin writing process # Begin writing process
@ -283,7 +261,7 @@ class SceneFileWriter(object):
index = 0 index = 0
while (insert_path := self.get_insert_file_path(index)).exists(): while (insert_path := self.get_insert_file_path(index)).exists():
index += 1 index += 1
self.inserted_file_path = str(insert_path) self.inserted_file_path = insert_path
self.open_movie_pipe(self.inserted_file_path) self.open_movie_pipe(self.inserted_file_path)
def end_insert(self): def end_insert(self):
@ -327,54 +305,6 @@ class SceneFileWriter(object):
else: else:
self.movie_file_path = self.temp_file_path self.movie_file_path = self.temp_file_path
def combine_movie_files(self) -> None:
kwargs = {
"remove_non_integer_files": True,
"extension": self.movie_file_extension,
}
if self.scene.start_at_animation_number is not None:
kwargs["min_index"] = self.scene.start_at_animation_number
if self.scene.end_at_animation_number is not None:
kwargs["max_index"] = self.scene.end_at_animation_number
else:
kwargs["remove_indices_greater_than"] = self.scene.num_plays - 1
partial_movie_files = get_sorted_integer_files(
self.partial_movie_directory,
**kwargs
)
if len(partial_movie_files) == 0:
log.warning("No animations in this scene")
return
# Write a file partial_file_list.txt containing all
# partial movie files
file_list = os.path.join(
self.partial_movie_directory,
"partial_movie_file_list.txt"
)
with open(file_list, 'w') as fp:
for pf_path in partial_movie_files:
if os.name == 'nt':
pf_path = pf_path.replace('\\', '/')
fp.write(f"file \'{pf_path}\'\n")
movie_file_path = self.get_movie_file_path()
commands = [
self.ffmpeg_bin,
'-y', # overwrite output file if it exists
'-f', 'concat',
'-safe', '0',
'-i', file_list,
'-loglevel', 'error',
'-c', 'copy',
movie_file_path
]
if not self.includes_sound:
commands.insert(-1, '-an')
combine_process = sp.Popen(commands)
combine_process.wait()
def add_sound_to_video(self) -> None: def add_sound_to_video(self) -> None:
movie_file_path = self.get_movie_file_path() movie_file_path = self.get_movie_file_path()
stem, ext = os.path.splitext(movie_file_path) stem, ext = os.path.splitext(movie_file_path)