Improved stage_scenes and refactored get_movie_output_directory methods

constants.py

@@ -2,7 +2,7 @@ import os
 import numpy as np
 
 # Things anyone wishing to use this repository for their
-# own use will want to change this
+# own use will want to change
 MEDIA_DIR = os.path.join(
     os.path.expanduser('~'),
     "Dropbox (3Blue1Brown)/3Blue1Brown Team Folder"

extract_scene.py

@@ -146,7 +146,7 @@ def handle_scene(scene, **config):
         if (platform.system() == "Linux"):
             commands = ["xdg-open"]
         elif (platform.system() == "Windows"):
-            commands = ["start"]    
+            commands = ["start"]
 
         if config["show_file_in_finder"]:
             commands.append("-R")
@@ -238,10 +238,10 @@ def main():
     module = get_module(config["file"])
     scene_names_to_classes = dict(inspect.getmembers(module, is_scene))
 
-    config["output_directory"] = os.path.join(
-        ANIMATIONS_DIR,
-        config["file"].replace(".py", "")
-    )
+    # config["output_directory"] = os.path.join(
+    #     ANIMATIONS_DIR,
+    #     config["file"].replace(".py", "")
+    # )
 
     scene_kwargs = dict([
         (key, config[key])
@@ -250,7 +250,6 @@ def main():
             "frame_duration",
             "skip_animations",
             "write_to_movie",
-            "output_directory",
             "save_pngs",
             "movie_file_extension",
             "start_at_animation_number",

scene/scene.py

@@ -29,6 +29,32 @@ def add_extension_if_not_present(file_name, extension):
         return file_name
 
 
+def get_scene_output_directory(scene_class):
+    file_path = os.path.abspath(inspect.getfile(scene_class))
+
+    # TODO, is there a better way to do this?
+    parts = file_path.split(os.path.sep)
+    if "manim" in parts:
+        sub_parts = parts[parts.index("manim") + 1:]
+        file_path = os.path.join(*sub_parts)
+    file_path = file_path.replace(".pyc", "")
+    file_path = file_path.replace(".py", "")
+    return os.path.join(ANIMATIONS_DIR, file_path)
+
+
+def get_movie_output_directory(scene_class, camera_config, frame_duration):
+    directory = get_scene_output_directory(scene_class)
+    sub_dir = "%dp%d" % (
+        camera_config["pixel_shape"][0],
+        int(1.0 / frame_duration)
+    )
+    return os.path.join(directory, sub_dir)
+
+
+def get_image_output_directory(scene_class, sub_dir="images"):
+    return os.path.join(get_scene_output_directory(scene_class), sub_dir)
+
+
 class Scene(Container):
     CONFIG = {
         "camera_class": Camera,
@@ -41,14 +67,12 @@ class Scene(Container):
         "save_frames": False,
         "save_pngs": False,
         "pngs_mode": "RGBA",
-        "output_directory": ANIMATIONS_DIR,
         "movie_file_extension": ".mp4",
         "name": None,
         "always_continually_update": False,
         "random_seed": 0,
         "start_at_animation_number": None,
         "end_at_animation_number": None,
-        "include_render_quality_in_output_directory": True,
     }
 
     def __init__(self, **kwargs):
@@ -543,10 +567,10 @@ class Scene(Container):
         self.get_image().show()
 
     def get_image_file_path(self, name=None, dont_update=False):
-        folder = "images"
+        sub_dir = "images"
         if dont_update:
-            folder = str(self)
-        path = os.path.join(self.output_directory, folder)
+            sub_dir = str(self)
+        path = get_image_output_directory(self.__class__, sub_dir)
         file_name = add_extension_if_not_present(name or str(self), ".png")
         return os.path.join(path, file_name)
 
@@ -562,16 +586,9 @@ class Scene(Container):
         image.save(path)
 
     def get_movie_file_path(self, name=None, extension=None):
-        directory = self.output_directory
-        if self.include_render_quality_in_output_directory:
-            sub_dir = "%dp%d" % (
-                self.camera.pixel_shape[0],
-                int(1.0 / self.frame_duration)
-            )
-            directory = os.path.join(directory, sub_dir)
-        if not os.path.exists(directory):
-            os.makedirs(directory)
-
+        directory = get_movie_output_directory(
+            self.__class__, self.camera_config, self.frame_duration
+        )
         if extension is None:
             extension = self.movie_file_extension
         if name is None:

stage_scenes.py

@@ -1,59 +1,68 @@
 import inspect
-import itertools as it
 import os
-import shutil
 import sys
 
 from constants import ANIMATIONS_DIR
-from constants import STAGED_SCENES_DIR
+from constants import PRODUCTION_QUALITY_CAMERA_CONFIG
+from constants import PRODUCTION_QUALITY_FRAME_DURATION
 from extract_scene import get_module
 from extract_scene import is_scene
+from scene.scene import get_movie_output_directory
 
 
-def get_sorted_scene_names(module_name):
+def get_sorted_scene_classes(module_name):
     module = get_module(module_name)
     line_to_scene = {}
     for name, scene_class in inspect.getmembers(module, is_scene):
+        if inspect.getmodule(scene_class) != module:
+            continue
         lines, line_no = inspect.getsourcelines(scene_class)
-        line_to_scene[line_no] = name
+        line_to_scene[line_no] = scene_class
     return [
-        line_to_scene[line_no]
-        for line_no in sorted(line_to_scene.keys())
+        line_to_scene[index]
+        for index in sorted(line_to_scene.keys())
     ]
 
+
 def stage_animations(module_name):
-    scene_names = get_sorted_scene_names(module_name)
-    animation_dir = os.path.join(
-        ANIMATIONS_DIR, module_name.replace(".py", "")
+    scene_classes = get_sorted_scene_classes(module_name)
+    if len(scene_classes) == 0:
+        print("There are no rendered animations from this module")
+        return
+    animation_dir = get_movie_output_directory(
+        scene_classes[0],
+        PRODUCTION_QUALITY_CAMERA_CONFIG,
+        PRODUCTION_QUALITY_FRAME_DURATION,
     )
     files = os.listdir(animation_dir)
     sorted_files = []
-    for scene in scene_names:
-        for clip in filter(lambda f : f.startswith(scene), files):
-            sorted_files.append(
-                os.path.join(animation_dir, clip)
-            )
-    staged_scenes_dir = os.path.join(animation_dir, "staged_scenes")
+    for scene_class in scene_classes:
+        scene_name = scene_class.__name__
+        for clip in filter(lambda f: f.startswith(scene_name + "."), files):
+            sorted_files.append(os.path.join(animation_dir, clip))
+
+    animation_subdir = os.path.dirname(animation_dir)
     count = 0
     while True:
         staged_scenes_dir = os.path.join(
-            animation_dir, "staged_scenes_%d"%count
+            animation_subdir, "staged_scenes_%d" % count
         )
         if not os.path.exists(staged_scenes_dir):
             os.makedirs(staged_scenes_dir)
             break
-        #Otherwise, keep trying new names until 
-        #there is a free one
+        # Otherwise, keep trying new names until
+        # there is a free one
         count += 1
     for count, f in enumerate(sorted_files):
         symlink_name = os.path.join(
             staged_scenes_dir,
-            "Scene_%03d"%count + f.split(os.sep)[-1]
+            "Scene_%03d" % count + f.split(os.sep)[-1]
         )
         os.symlink(f, symlink_name)
 
+
 if __name__ == "__main__":
     if len(sys.argv) < 2:
         raise Exception("No module given.")
     module_name = sys.argv[1]
-    stage_animations(module_name)
+    stage_animations(module_name)