extract_scene.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  1. import copy
  2. import inspect
  3. import sys
  4. from manimlib.config import get_custom_config
  5. from manimlib.logger import log
  6. from manimlib.scene.interactive_scene import InteractiveScene
  7. from manimlib.scene.scene import Scene
  8. class BlankScene(InteractiveScene):
  9. def construct(self):
  10. exec(get_custom_config()["universal_import_line"])
  11. self.embed()
  12. def is_child_scene(obj, module):
  13. if not inspect.isclass(obj):
  14. return False
  15. if not issubclass(obj, Scene):
  16. return False
  17. if obj == Scene:
  18. return False
  19. if not obj.__module__.startswith(module.__name__):
  20. return False
  21. return True
  22. def prompt_user_for_choice(scene_classes):
  23. name_to_class = {}
  24. max_digits = len(str(len(scene_classes)))
  25. for idx, scene_class in enumerate(scene_classes, start=1):
  26. name = scene_class.__name__
  27. print(f"{str(idx).zfill(max_digits)}: {name}")
  28. name_to_class[name] = scene_class
  29. try:
  30. user_input = input(
  31. "\nThat module has multiple scenes, " + \
  32. "which ones would you like to render?" + \
  33. "\nScene Name or Number: "
  34. )
  35. return [
  36. name_to_class[split_str] if not split_str.isnumeric() else scene_classes[int(split_str) - 1]
  37. for split_str in user_input.replace(" ", "").split(",")
  38. ]
  39. except IndexError:
  40. log.error("Invalid scene number")
  41. sys.exit(2)
  42. except KeyError:
  43. log.error("Invalid scene name")
  44. sys.exit(2)
  45. except EOFError:
  46. sys.exit(1)
  47. def get_scene_config(config):
  48. scene_parameters = inspect.signature(Scene).parameters.keys()
  49. return {
  50. key: config[key]
  51. for key in set(scene_parameters).intersection(config.keys())
  52. }
  53. def compute_total_frames(scene_class, scene_config):
  54. """
  55. When a scene is being written to file, a copy of the scene is run with
  56. skip_animations set to true so as to count how many frames it will require.
  57. This allows for a total progress bar on rendering, and also allows runtime
  58. errors to be exposed preemptively for long running scenes.
  59. """
  60. pre_config = copy.deepcopy(scene_config)
  61. pre_config["file_writer_config"]["write_to_movie"] = False
  62. pre_config["file_writer_config"]["save_last_frame"] = False
  63. pre_config["file_writer_config"]["quiet"] = True
  64. pre_config["skip_animations"] = True
  65. pre_scene = scene_class(**pre_config)
  66. pre_scene.run()
  67. total_time = pre_scene.time - pre_scene.skip_time
  68. return int(total_time * scene_config["camera_config"]["fps"])
  69. def scene_from_class(scene_class, scene_config, config):
  70. fw_config = scene_config["file_writer_config"]
  71. if fw_config["write_to_movie"] and config["prerun"]:
  72. fw_config["total_frames"] = compute_total_frames(scene_class, scene_config)
  73. return scene_class(**scene_config)
  74. def get_scenes_to_render(all_scene_classes, scene_config, config):
  75. if config["write_all"]:
  76. return [sc(**scene_config) for sc in all_scene_classes]
  77. names_to_classes = {sc.__name__ : sc for sc in all_scene_classes}
  78. scene_names = config["scene_names"]
  79. for name in set.difference(set(scene_names), names_to_classes):
  80. log.error(f"No scene named {name} found")
  81. scene_names.remove(name)
  82. if scene_names:
  83. classes_to_run = [names_to_classes[name] for name in scene_names]
  84. elif len(all_scene_classes) == 1:
  85. classes_to_run = [all_scene_classes[0]]
  86. else:
  87. classes_to_run = prompt_user_for_choice(all_scene_classes)
  88. return [
  89. scene_from_class(scene_class, scene_config, config)
  90. for scene_class in classes_to_run
  91. ]
  92. def get_scene_classes_from_module(module):
  93. if hasattr(module, "SCENES_IN_ORDER"):
  94. return module.SCENES_IN_ORDER
  95. else:
  96. return [
  97. member[1]
  98. for member in inspect.getmembers(
  99. module,
  100. lambda x: is_child_scene(x, module)
  101. )
  102. ]
  103. def main(config):
  104. module = config["module"]
  105. scene_config = get_scene_config(config)
  106. if module is None:
  107. # If no module was passed in, just play the blank scene
  108. return [BlankScene(**scene_config)]
  109. all_scene_classes = get_scene_classes_from_module(module)
  110. scenes = get_scenes_to_render(all_scene_classes, scene_config, config)
  111. return scenes