Merge pull request #2 from 3b1b/master

Updating Manim

manimlib/constants.py

@@ -31,14 +31,12 @@
 THIS_DIR = os.path.dirname(os.path.realpath(__file__))
 FILE_DIR = os.path.join(THIS_DIR, "files")
 TEX_DIR = os.path.join(FILE_DIR, "Tex")
-TEX_IMAGE_DIR = TEX_DIR  # TODO, What is this doing?
 # These two may be depricated now.
 MOBJECT_DIR = os.path.join(FILE_DIR, "mobjects")
 IMAGE_MOBJECT_DIR = os.path.join(MOBJECT_DIR, "image")
 
 for folder in [FILE_DIR, RASTER_IMAGE_DIR, SVG_IMAGE_DIR, VIDEO_DIR, TEX_DIR,
-               TEX_IMAGE_DIR, MOBJECT_DIR, IMAGE_MOBJECT_DIR,
-               STAGED_SCENES_DIR]:
+               MOBJECT_DIR, IMAGE_MOBJECT_DIR, STAGED_SCENES_DIR]:
     if not os.path.exists(folder):
         os.makedirs(folder)
 

manimlib/for_3b1b_videos/pi_creature_scene.py

@@ -213,7 +213,6 @@ def wait(self, time=1, blink=True):
             time_to_blink = self.total_wait_time % self.seconds_to_blink == 0
             if blink and self.any_pi_creatures_on_screen() and time_to_blink:
                 self.blink()
-                self.num_plays -= 1  # This shouldn't count as an animation
             else:
                 self.non_blink_wait()
             time -= 1

manimlib/scene/scene.py

@@ -5,7 +5,7 @@
 import os
 import random
 import shutil
-import subprocess as sp
+import subprocess
 import warnings
 
 from tqdm import tqdm as ProgressDisplay
@@ -24,6 +24,8 @@
 from manimlib.utils.output_directory_getters import add_extension_if_not_present
 from manimlib.utils.output_directory_getters import get_image_output_directory
 from manimlib.utils.output_directory_getters import get_movie_output_directory
+from manimlib.utils.output_directory_getters import get_partial_movie_output_directory
+from manimlib.utils.output_directory_getters import get_sorted_integer_files
 
 
 class Scene(Container):
@@ -35,11 +37,9 @@ class Scene(Container):
         "skip_animations": False,
         "ignore_waits": False,
         "write_to_movie": False,
-        "save_frames": False,
         "save_pngs": False,
         "pngs_mode": "RGBA",
         "movie_file_extension": ".mp4",
-        "name": None,
         "always_continually_update": False,
         "random_seed": 0,
         "start_at_animation_number": None,
@@ -57,38 +57,38 @@ def __init__(self, **kwargs):
         self.continual_animations = []
         self.foreground_mobjects = []
         self.num_plays = 0
-        self.saved_frames = []
-        self.shared_locals = {}
         self.frame_num = 0
-        self.current_scene_time = 0
+        self.time = 0
         self.original_skipping_status = self.skip_animations
         self.stream_lock = False
-        if self.name is None:
-            self.name = self.__class__.__name__
         if self.random_seed is not None:
             random.seed(self.random_seed)
             np.random.seed(self.random_seed)
 
         self.setup()
-        if self.write_to_movie:
-            self.open_movie_pipe()
         if self.livestreaming:
             return None
         try:
             self.construct(*self.construct_args)
         except EndSceneEarlyException:
             pass
-
-        # Always tack on one last frame, so that scenes
-        # with no play calls still display something
-        self.skip_animations = False
-        self.wait(self.frame_duration)
-
         self.tear_down()
 
         if self.write_to_movie:
-            self.close_movie_pipe()
-        print("Played a total of %d animations" % self.num_plays)
+            self.combine_movie_files()
+
+    def handle_play_like_call(func):
+        def wrapper(self, *args, **kwargs):
+            self.handle_animation_skipping()
+            should_write = self.write_to_movie and not self.skip_animations
+            if should_write:
+                self.open_movie_pipe()
+                func(self, *args, **kwargs)
+                self.close_movie_pipe()
+            else:
+                func(self, *args, **kwargs)
+            self.num_plays += 1
+        return wrapper
 
     def setup(self):
         """
@@ -109,11 +109,7 @@ def construct(self):
         pass  # To be implemented in subclasses
 
     def __str__(self):
-        return self.name
-
-    def set_name(self, name):
-        self.name = name
-        return self
+        return self.__class__.__name__
 
     def set_variables_as_attrs(self, *objects, **newly_named_objects):
         """
@@ -214,6 +210,14 @@ def should_continually_update(self):
 
     ###
 
+    def get_time(self):
+        return self.time
+
+    def increment_time(self, d_time):
+        self.time += d_time
+
+    ###
+
     def get_top_level_mobjects(self):
         # Return only those which are not in the family
         # of another mobject from the scene
@@ -465,13 +469,14 @@ def handle_animation_skipping(self):
                 self.skip_animations = True
                 raise EndSceneEarlyException()
 
+    @handle_play_like_call
     def play(self, *args, **kwargs):
         if self.livestreaming:
             self.stream_lock = False
         if len(args) == 0:
             warnings.warn("Called Scene.play with no animations")
             return
-        self.handle_animation_skipping()
+
         animations = self.compile_play_args_to_animation_list(*args)
         for animation in animations:
             # This is where kwargs to play like run_time and rate_func
@@ -503,12 +508,10 @@ def play(self, *args, **kwargs):
             self.continual_update(total_run_time)
         else:
             self.continual_update(0)
-        self.num_plays += 1
 
         if self.livestreaming:
             self.stream_lock = True
             thread.start_new_thread(self.idle_stream, ())
-
         return self
 
     def idle_stream(self):
@@ -533,6 +536,7 @@ def get_mobjects_from_last_animation(self):
             return self.mobjects_from_last_animation
         return []
 
+    @handle_play_like_call
     def wait(self, duration=DEFAULT_WAIT_TIME):
         if self.should_continually_update():
             total_time = 0
@@ -554,7 +558,7 @@ def wait(self, duration=DEFAULT_WAIT_TIME):
     def wait_to(self, time, assert_positive=True):
         if self.ignore_waits:
             return
-        time -= self.current_scene_time
+        time -= self.get_time()
         if assert_positive:
             assert(time >= 0)
         elif time < 0:
@@ -575,16 +579,15 @@ def revert_to_original_skipping_status(self):
     def add_frames(self, *frames):
         if self.skip_animations:
             return
-        self.current_scene_time += len(frames) * self.frame_duration
+        self.increment_time(len(frames) * self.frame_duration)
         if self.write_to_movie:
             for frame in frames:
                 if self.save_pngs:
                     self.save_image(
-                        "frame" + str(self.frame_num), self.pngs_mode, True)
+                        "frame" + str(self.frame_num), self.pngs_mode, True
+                    )
                     self.frame_num = self.frame_num + 1
                 self.writing_process.stdin.write(frame.tostring())
-        if self.save_frames:
-            self.saved_frames += list(frames)
 
     # Display methods
 
@@ -615,17 +618,26 @@ def get_movie_file_path(self, name=None, extension=None):
         if extension is None:
             extension = self.movie_file_extension
         if name is None:
-            name = self.name
+            name = str(self)
         file_path = os.path.join(directory, name)
         if not file_path.endswith(extension):
             file_path += extension
         return file_path
 
+    def get_partial_movie_directory(self):
+        return get_partial_movie_output_directory(
+            self.__class__, self.camera_config, self.frame_duration
+        )
+
     def open_movie_pipe(self):
-        name = str(self)
-        file_path = self.get_movie_file_path(name)
-        temp_file_path = file_path.replace(name, name + "Temp")
-        print("Writing to %s" % temp_file_path)
+        directory = self.get_partial_movie_directory()
+        file_path = os.path.join(
+            directory, "{}{}".format(
+                self.num_plays,
+                self.movie_file_extension,
+            )
+        )
+        temp_file_path = file_path.replace(".", "_temp.")
         self.args_to_rename_file = (temp_file_path, file_path)
 
         fps = int(1 / self.frame_duration)
@@ -649,6 +661,7 @@ def open_movie_pipe(self):
             # should be transparent.
             command += [
                 '-vcodec', 'qtrle',
+                # '-vcodec', 'png',
             ]
         else:
             command += [
@@ -664,18 +677,56 @@ def open_movie_pipe(self):
                 command += [STREAMING_PROTOCOL + '://' + STREAMING_IP + ':' + STREAMING_PORT]
         else:
             command += [temp_file_path]
-        # self.writing_process = sp.Popen(command, stdin=sp.PIPE, shell=True)
-        self.writing_process = sp.Popen(command, stdin=sp.PIPE)
+        self.writing_process = subprocess.Popen(command, stdin=subprocess.PIPE)
 
     def close_movie_pipe(self):
         self.writing_process.stdin.close()
         self.writing_process.wait()
         if self.livestreaming:
             return True
-        if os.name == 'nt':
-            shutil.move(*self.args_to_rename_file)
+        shutil.move(*self.args_to_rename_file)
+
+    def combine_movie_files(self):
+        partial_movie_file_directory = self.get_partial_movie_directory()
+        kwargs = {
+            "remove_non_integer_files": True,
+            "extension": self.movie_file_extension,
+        }
+        if self.start_at_animation_number is not None:
+            kwargs["min_index"] = self.start_at_animation_number
+        if self.end_at_animation_number is not None:
+            kwargs["max_index"] = self.end_at_animation_number
         else:
-            os.rename(*self.args_to_rename_file)
+            kwargs["remove_indices_greater_than"] = self.num_plays - 1
+        partial_movie_files = get_sorted_integer_files(
+            partial_movie_file_directory,
+            **kwargs
+        )
+        # Write a file partial_file_list.txt containing all
+        # partial movie files
+        file_list = os.path.join(
+            partial_movie_file_directory,
+            "partial_movie_file_list.txt"
+        )
+        with open(file_list, 'w') as fp:
+            for pf_path in partial_movie_files:
+                fp.write("file {}\n".format(pf_path))
+
+        movie_file_path = self.get_movie_file_path()
+        commands = [
+            FFMPEG_BIN,
+            '-y',  # overwrite output file if it exists
+            '-f', 'concat',
+            '-safe', '0',
+            '-i', file_list,
+            '-c', 'copy',
+            '-an',  # Tells FFMPEG not to expect any audio
+            '-loglevel', 'error',
+            movie_file_path
+        ]
+        subprocess.call(commands)
+        os.remove(file_list)
+        print("File ready at {}".format(movie_file_path))
 
     def tex(self, latex):
         eq = TextMobject(latex)

manimlib/utils/output_directory_getters.py

@@ -1,7 +1,6 @@
-import inspect
 import os
+import numpy as np
 
-from manimlib.constants import THIS_DIR
 from manimlib.constants import VIDEO_DIR
 
 
@@ -35,6 +34,44 @@ def get_movie_output_directory(scene_class, camera_config, frame_duration):
     return guarantee_existance(os.path.join(directory, sub_dir))
 
 
+def get_partial_movie_output_directory(scene_class, camera_config, frame_duration):
+    directory = get_movie_output_directory(scene_class, camera_config, frame_duration)
+    return guarantee_existance(
+        os.path.join(directory, scene_class.__name__)
+    )
+
+
 def get_image_output_directory(scene_class, sub_dir="images"):
     directory = get_scene_output_directory(scene_class)
     return guarantee_existance(os.path.join(directory, sub_dir))
+
+
+def get_sorted_integer_files(directory,
+                             min_index=0,
+                             max_index=np.inf,
+                             remove_non_integer_files=False,
+                             remove_indices_greater_than=None,
+                             extension=None,
+                             ):
+    indexed_files = []
+    for file in os.listdir(directory):
+        if '.' in file:
+            index_str = file[:file.index('.')]
+        else:
+            index_str = file
+
+        full_path = os.path.join(directory, file)
+        if index_str.isdigit():
+            index = int(index_str)
+            if remove_indices_greater_than is not None:
+                if index > remove_indices_greater_than:
+                    os.remove(full_path)
+                    continue
+            if extension is not None and not file.endswith(extension):
+                continue
+            if index >= min_index and index < max_index:
+                indexed_files.append((index, file))
+        elif remove_non_integer_files:
+            os.remove(full_path)
+    indexed_files.sort(key=lambda p: p[0])
+    return map(lambda p: p[1], indexed_files)