public-mirrors/3b1b-manim
1
0
Fork 0
mirror of https://github.com/3b1b/manim.git synced 2025-04-13 09:47:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Make manim into a module and refactor configuration

Browse source
This commit is contained in:
Devin Neal 2018-12-22 14:27:22 -08:00
parent 465951ab88
commit 27f677e45f
6 changed files with 290 additions and 257 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

142
config.py Normal file
View file

@ -0,0 +1,142 @@
#!/usr/bin/env python
import argparse
import colour
import constants
import os
import sys
def parse_cli():
try:
parser = argparse.ArgumentParser()
module_location = parser.add_mutually_exclusive_group()
module_location.add_argument(
"file",
nargs="?",
help="path to file holding the python code for the scene",
)
parser.add_argument(
"scene_name",
nargs="?",
help="Name of the Scene class you want to see",
)
module_location.add_argument("--livestream", action="store_true")
parser.add_argument("--to-twitch", action="store_true")
optional_args = [
("-p", "--preview"),
("-w", "--write_to_movie"),
("-s", "--show_last_frame"),
("-l", "--low_quality"),
("-m", "--medium_quality"),
("-g", "--save_pngs"),
("-f", "--show_file_in_finder"),
("-t", "--transparent"),
("-q", "--quiet"),
("-a", "--write_all")
]
for short_arg, long_arg in optional_args:
parser.add_argument(short_arg, long_arg, action="store_true")
parser.add_argument("-o", "--output_name")
parser.add_argument("-n", "--start_at_animation_number")
parser.add_argument("-r", "--resolution")
parser.add_argument("-c", "--color")
args = parser.parse_args()
if args.file is None and not args.livestream:
parser.print_help()
sys.exit(2)
else:
return args
except argparse.ArgumentError as err:
print(str(err))
sys.exit(2)
def get_configuration(args):
if args.output_name is not None:
output_name_root, output_name_ext = os.path.splitext(
args.output_name)
expected_ext = '.png' if args.show_last_frame else '.mp4'
if output_name_ext not in ['', expected_ext]:
print("WARNING: The output will be to (doubly-dotted) %s%s" %
output_name_root % expected_ext)
output_name = args.output_name
else:
# If anyone wants .mp4.mp4 and is surprised to only get .mp4, or such... Well, too bad.
output_name = output_name_root
else:
output_name = args.output_name
config = {
"file": args.file,
"scene_name": args.scene_name,
"open_video_upon_completion": args.preview,
"show_file_in_finder": args.show_file_in_finder,
# By default, write to file
"write_to_movie": args.write_to_movie or not args.show_last_frame,
"show_last_frame": args.show_last_frame,
"save_pngs": args.save_pngs,
# If -t is passed in (for transparent), this will be RGBA
"saved_image_mode": "RGBA" if args.transparent else "RGB",
"movie_file_extension": ".mov" if args.transparent else ".mp4",
"quiet": args.quiet or args.write_all,
"ignore_waits": args.preview,
"write_all": args.write_all,
"output_name": output_name,
"start_at_animation_number": args.start_at_animation_number,
"end_at_animation_number": None,
}
# Camera configuration
config["camera_config"] = {}
if args.low_quality:
config["camera_config"].update(constants.LOW_QUALITY_CAMERA_CONFIG)
config["frame_duration"] = constants.LOW_QUALITY_FRAME_DURATION
elif args.medium_quality:
config["camera_config"].update(constants.MEDIUM_QUALITY_CAMERA_CONFIG)
config["frame_duration"] = constants.MEDIUM_QUALITY_FRAME_DURATION
else:
config["camera_config"].update(constants.PRODUCTION_QUALITY_CAMERA_CONFIG)
config["frame_duration"] = constants.PRODUCTION_QUALITY_FRAME_DURATION
# If the resolution was passed in via -r
if args.resolution:
if "," in args.resolution:
height_str, width_str = args.resolution.split(",")
height = int(height_str)
width = int(width_str)
else:
height = int(args.resolution)
width = int(16 * height / 9)
config["camera_config"].update({
"pixel_height": height,
"pixel_width": width,
})
if args.color:
try:
config["camera_config"]["background_color"] = colour.Color(args.color)
except AttributeError as err:
print("Please use a valid color")
print(err)
sys.exit(2)
# If rendering a transparent image/move, make sure the
# scene has a background opacity of 0
if args.transparent:
config["camera_config"]["background_opacity"] = 0
# Arguments related to skipping
stan = config["start_at_animation_number"]
if stan is not None:
if "," in stan:
start, end = stan.split(",")
config["start_at_animation_number"] = int(start)
config["end_at_animation_number"] = int(end)
else:
config["start_at_animation_number"] = int(stan)
config["skip_animations"] = any([
config["show_last_frame"] and not config["write_to_movie"],
config["start_at_animation_number"],
])
return config

108
constants.py
View file

@ -1,6 +1,7 @@
import os
import numpy as np
# Initialize directories
env_MEDIA_DIR = os.getenv("MEDIA_DIR")
if env_MEDIA_DIR:
MEDIA_DIR = env_MEDIA_DIR
@ -12,17 +13,78 @@ else:
os.path.expanduser('~'),
"Dropbox (3Blue1Brown)/3Blue1Brown Team Folder"
)
if not os.path.isdir(MEDIA_DIR):
MEDIA_DIR = "media"
print(
f"Media will be stored in {MEDIA_DIR + os.sep}. You can change " + \
f"Media will be stored in {MEDIA_DIR + os.sep}. You can change "
"this behavior by writing a different directory to media_dir.txt."
)
with open("media_dir.txt", 'w') as media_file:
media_file.write(MEDIA_DIR)
VIDEO_DIR = os.path.join(MEDIA_DIR, "videos")
RASTER_IMAGE_DIR = os.path.join(MEDIA_DIR, "designs", "raster_images")
SVG_IMAGE_DIR = os.path.join(MEDIA_DIR, "designs", "svg_images")
# TODO, staged scenes should really go into a subdirectory of a given scenes directory
STAGED_SCENES_DIR = os.path.join(VIDEO_DIR, "staged_scenes")
###
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]:
if not os.path.exists(folder):
os.makedirs(folder)
TEX_USE_CTEX = False
TEX_TEXT_TO_REPLACE = "YourTextHere"
TEMPLATE_TEX_FILE = os.path.join(
THIS_DIR, "tex_template.tex" if not TEX_USE_CTEX
else "ctex_template.tex"
)
with open(TEMPLATE_TEX_FILE, "r") as infile:
TEMPLATE_TEXT_FILE_BODY = infile.read()
TEMPLATE_TEX_FILE_BODY = TEMPLATE_TEXT_FILE_BODY.replace(
TEX_TEXT_TO_REPLACE,
"\\begin{align*}\n" + TEX_TEXT_TO_REPLACE + "\n\\end{align*}",
)
HELP_MESSAGE = """
Usage:
python extract_scene.py <module> [<scene name>]
-p preview in low quality
-s show and save picture of last frame
-w write result to file [this is default if nothing else is stated]
-o <file_name> write to a different file_name
-l use low quality
-m use medium quality
-a run and save every scene in the script, or all args for the given scene
-q don't print progress
-f when writing to a movie file, export the frames in png sequence
-t use transperency when exporting images
-n specify the number of the animation to start from
-r specify a resolution
-c specify a background color
"""
SCENE_NOT_FOUND_MESSAGE = """
That scene is not in the script
"""
CHOOSE_NUMBER_MESSAGE = """
Choose number corresponding to desired scene/arguments.
(Use comma separated list for multiple entries)
Choice(s): """
INVALID_NUMBER_MESSAGE = "Fine then, if you don't want to give a valid number I'll just quit"
NO_SCENE_MESSAGE = """
There are no scenes inside that module
"""
LOW_QUALITY_FRAME_DURATION = 1. / 15
MEDIUM_QUALITY_FRAME_DURATION = 1. / 30
PRODUCTION_QUALITY_FRAME_DURATION = 1. / 60
@ -102,44 +164,9 @@ PI = np.pi
TAU = 2 * PI
DEGREES = TAU / 360
VIDEO_DIR = os.path.join(MEDIA_DIR, "videos")
RASTER_IMAGE_DIR = os.path.join(MEDIA_DIR, "designs", "raster_images")
SVG_IMAGE_DIR = os.path.join(MEDIA_DIR, "designs", "svg_images")
# TODO, staged scenes should really go into a subdirectory of a given scenes directory
STAGED_SCENES_DIR = os.path.join(VIDEO_DIR, "staged_scenes")
###
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]:
if not os.path.exists(folder):
os.makedirs(folder)
TEX_USE_CTEX = False
TEX_TEXT_TO_REPLACE = "YourTextHere"
TEMPLATE_TEX_FILE = os.path.join(
THIS_DIR, "tex_template.tex" if not TEX_USE_CTEX
else "ctex_template.tex"
)
with open(TEMPLATE_TEX_FILE, "r") as infile:
TEMPLATE_TEXT_FILE_BODY = infile.read()
TEMPLATE_TEX_FILE_BODY = TEMPLATE_TEXT_FILE_BODY.replace(
TEX_TEXT_TO_REPLACE,
"\\begin{align*}\n" + TEX_TEXT_TO_REPLACE + "\n\\end{align*}",
)
FFMPEG_BIN = "ffmpeg"
# Colors
COLOR_MAP = {
"DARK_BLUE": "#236B8E",
"DARK_BROWN": "#8B4513",
@ -202,12 +229,11 @@ locals().update(COLOR_MAP)
for name in [s for s in list(COLOR_MAP.keys()) if s.endswith("_C")]:
locals()[name.replace("_C", "")] = locals()[name]
# Streaming related configurations
IS_LIVE_STREAMING = False
# Streaming related configuration
LIVE_STREAM_NAME = "LiveStream"
IS_STREAMING_TO_TWITCH = False
TWITCH_STREAM_KEY = "YOUR_STREAM_KEY"
STREAMING_PROTOCOL = "tcp"
STREAMING_IP = "127.0.0.1"
STREAMING_PORT = "2000"
STREAMING_CLIENT = "ffplay"
STREAMING_URL = f"{STREAMING_PROTOCOL}://{STREAMING_IP}:{STREAMING_PORT}?listen"

186
extract_scene.py
View file

@ -1,173 +1,18 @@
# !/usr/bin/env python2
# !/usr/bin/env python
import sys
import argparse
# import imp
import constants
import importlib
import inspect
import itertools as it
import os
import subprocess as sp
import sys
import traceback
from constants import *
from scene.scene import Scene
from utils.sounds import play_error_sound
from utils.sounds import play_finish_sound
from colour import Color
HELP_MESSAGE = """
Usage:
python extract_scene.py <module> [<scene name>]
-p preview in low quality
-s show and save picture of last frame
-w write result to file [this is default if nothing else is stated]
-o <file_name> write to a different file_name
-l use low quality
-m use medium quality
-a run and save every scene in the script, or all args for the given scene
-q don't print progress
-f when writing to a movie file, export the frames in png sequence
-t use transperency when exporting images
-n specify the number of the animation to start from
-r specify a resolution
-c specify a background color
"""
SCENE_NOT_FOUND_MESSAGE = """
That scene is not in the script
"""
CHOOSE_NUMBER_MESSAGE = """
Choose number corresponding to desired scene/arguments.
(Use comma separated list for multiple entries)
Choice(s): """
INVALID_NUMBER_MESSAGE = "Fine then, if you don't want to give a valid number I'll just quit"
NO_SCENE_MESSAGE = """
There are no scenes inside that module
"""
def get_configuration():
try:
parser = argparse.ArgumentParser()
parser.add_argument(
"file", help="path to file holding the python code for the scene"
)
parser.add_argument(
"scene_name", help="Name of the Scene class you want to see"
)
optional_args = [
("-p", "--preview"),
("-w", "--write_to_movie"),
("-s", "--show_last_frame"),
("-l", "--low_quality"),
("-m", "--medium_quality"),
("-g", "--save_pngs"),
("-f", "--show_file_in_finder"),
("-t", "--transparent"),
("-q", "--quiet"),
("-a", "--write_all")
]
for short_arg, long_arg in optional_args:
parser.add_argument(short_arg, long_arg, action="store_true")
parser.add_argument("-o", "--output_name")
parser.add_argument("-n", "--start_at_animation_number")
parser.add_argument("-r", "--resolution")
parser.add_argument("-c", "--color")
args = parser.parse_args()
if args.output_name is not None:
output_name_root, output_name_ext = os.path.splitext(
args.output_name)
expected_ext = '.png' if args.show_last_frame else '.mp4'
if output_name_ext not in ['', expected_ext]:
print("WARNING: The output will be to (doubly-dotted) %s%s" %
output_name_root % expected_ext)
output_name = args.output_name
else:
# If anyone wants .mp4.mp4 and is surprised to only get .mp4, or such... Well, too bad.
output_name = output_name_root
else:
output_name = args.output_name
except argparse.ArgumentError as err:
print(str(err))
sys.exit(2)
config = {
"file": args.file,
"scene_name": args.scene_name,
"open_video_upon_completion": args.preview,
"show_file_in_finder": args.show_file_in_finder,
# By default, write to file
"write_to_movie": args.write_to_movie or not args.show_last_frame,
"show_last_frame": args.show_last_frame,
"save_pngs": args.save_pngs,
# If -t is passed in (for transparent), this will be RGBA
"saved_image_mode": "RGBA" if args.transparent else "RGB",
"movie_file_extension": ".mov" if args.transparent else ".mp4",
"quiet": args.quiet or args.write_all,
"ignore_waits": args.preview,
"write_all": args.write_all,
"output_name": output_name,
"start_at_animation_number": args.start_at_animation_number,
"end_at_animation_number": None,
}
# Camera configuration
config["camera_config"] = {}
if args.low_quality:
config["camera_config"].update(LOW_QUALITY_CAMERA_CONFIG)
config["frame_duration"] = LOW_QUALITY_FRAME_DURATION
elif args.medium_quality:
config["camera_config"].update(MEDIUM_QUALITY_CAMERA_CONFIG)
config["frame_duration"] = MEDIUM_QUALITY_FRAME_DURATION
else:
config["camera_config"].update(PRODUCTION_QUALITY_CAMERA_CONFIG)
config["frame_duration"] = PRODUCTION_QUALITY_FRAME_DURATION
# If the resolution was passed in via -r
if args.resolution:
if "," in args.resolution:
height_str, width_str = args.resolution.split(",")
height = int(height_str)
width = int(width_str)
else:
height = int(args.resolution)
width = int(16 * height / 9)
config["camera_config"].update({
"pixel_height": height,
"pixel_width": width,
})
if args.color:
try:
config["camera_config"]["background_color"] = Color(args.color)
except AttributeError as err:
print("Please use a valid color")
print(err)
sys.exit(2)
# If rendering a transparent image/move, make sure the
# scene has a background opacity of 0
if args.transparent:
config["camera_config"]["background_opacity"] = 0
# Arguments related to skipping
stan = config["start_at_animation_number"]
if stan is not None:
if "," in stan:
start, end = stan.split(",")
config["start_at_animation_number"] = int(start)
config["end_at_animation_number"] = int(end)
else:
config["start_at_animation_number"] = int(stan)
config["skip_animations"] = any([
config["show_last_frame"] and not config["write_to_movie"],
config["start_at_animation_number"],
])
return config
def handle_scene(scene, **config):
import platform
@ -223,26 +68,31 @@ def prompt_user_for_choice(name_to_obj):
print("%d: %s" % (count, name))
num_to_name[count] = name
try:
user_input = input(CHOOSE_NUMBER_MESSAGE)
user_input = input(constants.CHOOSE_NUMBER_MESSAGE)
return [
name_to_obj[num_to_name[int(num_str)]]
for num_str in user_input.split(",")
]
except:
print(INVALID_NUMBER_MESSAGE)
except KeyError:
print(constants.INVALID_NUMBER_MESSAGE)
sys.exit()
user_input = input(constants.CHOOSE_NUMBER_MESSAGE)
return [
name_to_obj[num_to_name[int(num_str)]]
for num_str in user_input.split(",")
]
def get_scene_classes(scene_names_to_classes, config):
if len(scene_names_to_classes) == 0:
print(NO_SCENE_MESSAGE)
print(constants.NO_SCENE_MESSAGE)
return []
if len(scene_names_to_classes) == 1:
return list(scene_names_to_classes.values())
if config["scene_name"] in scene_names_to_classes:
return [scene_names_to_classes[config["scene_name"]]]
if config["scene_name"] != "":
print(SCENE_NOT_FOUND_MESSAGE)
print(constants.SCENE_NOT_FOUND_MESSAGE)
return []
if config["write_all"]:
return list(scene_names_to_classes.values())
@ -254,16 +104,10 @@ def get_module(file_name):
return importlib.import_module(module_name)
def main():
config = get_configuration()
def main(config):
module = get_module(config["file"])
scene_names_to_classes = dict(inspect.getmembers(module, is_scene))
# config["output_directory"] = os.path.join(
# VIDEO_DIR,
# config["file"].replace(".py", "")
# )
scene_kwargs = dict([
(key, config[key])
for key in [
@ -288,7 +132,7 @@ def main():
try:
handle_scene(SceneClass(**scene_kwargs), **config)
play_finish_sound()
except:
except Exception:
print("\n\n")
traceback.print_exc()
print("\n\n")

40
manim.py
View file

@ -1,31 +1,11 @@
#!/usr/bin/python3
#!/usr/bin/env python
import config
import extract_scene
import stream_starter
from constants import *
from scene.scene import Scene
class Manim():
def __new__(cls):
kwargs = {
"scene_name": LIVE_STREAM_NAME,
"open_video_upon_completion": False,
"show_file_in_finder": False,
# By default, write to file
"write_to_movie": True,
"show_last_frame": False,
"save_pngs": False,
# If -t is passed in (for transparent), this will be RGBA
"saved_image_mode": "RGB",
"movie_file_extension": ".mp4",
"quiet": True,
"ignore_waits": False,
"write_all": False,
"name": LIVE_STREAM_NAME,
"start_at_animation_number": 0,
"end_at_animation_number": None,
"skip_animations": False,
"camera_config": HIGH_QUALITY_CAMERA_CONFIG,
"frame_duration": MEDIUM_QUALITY_FRAME_DURATION,
}
return Scene(**kwargs)
args = config.parse_cli()
if not args.livestream:
config = config.get_configuration(args)
extract_scene.main(config)
else:
stream_starter.start_livestream(args.to_twitch)

14
scene/scene.py
View file

@ -51,6 +51,8 @@ class Scene(Container):
"random_seed": 0,
"start_at_animation_number": None,
"end_at_animation_number": None,
"livestreaming": False,
"to_twitch": False,
}
def __init__(self, **kwargs):
@ -76,7 +78,7 @@ class Scene(Container):
self.setup()
if self.write_to_movie:
self.open_movie_pipe()
if IS_LIVE_STREAMING:
if self.livestreaming:
return None
try:
self.construct(*self.construct_args)
@ -464,7 +466,7 @@ class Scene(Container):
raise EndSceneEarlyException()
def play(self, *args, **kwargs):
if IS_LIVE_STREAMING:
if self.livestreaming:
self.stream_lock = False
if len(args) == 0:
warnings.warn("Called Scene.play with no animations")
@ -503,7 +505,7 @@ class Scene(Container):
self.continual_update(0)
self.num_plays += 1
if IS_LIVE_STREAMING:
if self.livestreaming:
self.stream_lock = True
thread.start_new_thread(self.idle_stream, ())
@ -653,8 +655,8 @@ class Scene(Container):
'-vcodec', 'libx264',
'-pix_fmt', 'yuv420p',
]
if IS_LIVE_STREAMING:
if IS_STREAMING_TO_TWITCH:
if self.livestreaming:
if self.to_twitch:
command += ['-f', 'flv']
command += ['rtmp://live.twitch.tv/app/' + TWITCH_STREAM_KEY]
else:
@ -668,7 +670,7 @@ class Scene(Container):
def close_movie_pipe(self):
self.writing_process.stdin.close()
self.writing_process.wait()
if IS_LIVE_STREAMING:
if self.livestreaming:
return True
if os.name == 'nt':
shutil.move(*self.args_to_rename_file)

57
stream_starter.py
View file

@ -1,13 +1,52 @@
from big_ol_pile_of_manim_imports import *
import subprocess
from scene.scene import Scene
from time import sleep
from manim import Manim
import code
import constants
import os
import subprocess
if not IS_STREAMING_TO_TWITCH:
FNULL = open(os.devnull, 'w')
subprocess.Popen([STREAMING_CLIENT, STREAMING_PROTOCOL + '://' + STREAMING_IP + ':' + STREAMING_PORT + '?listen'], stdout=FNULL, stderr=FNULL)
sleep(3)
manim = Manim()
def start_livestream(to_twitch=False):
class Manim():
print("YOUR STREAM IS READY!")
def __new__(cls):
kwargs = {
"scene_name": constants.LIVE_STREAM_NAME,
"open_video_upon_completion": False,
"show_file_in_finder": False,
# By default, write to file
"write_to_movie": True,
"show_last_frame": False,
"save_pngs": False,
# If -t is passed in (for transparent), this will be RGBA
"saved_image_mode": "RGB",
"movie_file_extension": ".mp4",
"quiet": True,
"ignore_waits": False,
"write_all": False,
"name": constants.LIVE_STREAM_NAME,
"start_at_animation_number": 0,
"end_at_animation_number": None,
"skip_animations": False,
"camera_config": constants.HIGH_QUALITY_CAMERA_CONFIG,
"frame_duration": constants.MEDIUM_QUALITY_FRAME_DURATION,
"livestreaming": True,
"to_twitch": to_twitch,
}
return Scene(**kwargs)
if not to_twitch:
FNULL = open(os.devnull, 'w')
subprocess.Popen(
[constants.STREAMING_CLIENT, constants.STREAMING_URL],
stdout=FNULL,
stderr=FNULL)
sleep(3)
manim = Manim()
print("YOUR STREAM IS READY!")
variables = globals().copy()
variables.update(locals())
shell = code.InteractiveConsole(variables)
shell.push("from big_ol_pile_of_manim_imports import *")
shell.interact()