From 7564beb393c8fda0798d6d4d5756984dcbe99b24 Mon Sep 17 00:00:00 2001 From: Grant Sanderson Date: Sun, 11 Oct 2015 15:22:47 -0700 Subject: [PATCH] Finally have faster video writing capabilities --- constants.py | 1 - displayer.py | 116 ++++++++---------------------- mobject/function_graphs.py | 1 + sample_script.py | 8 +-- scene/scene.py | 1 + script_wrapper.py | 18 +++-- scripts/matrix_as_transform_2d.py | 65 +++++++++++------ 7 files changed, 90 insertions(+), 120 deletions(-) diff --git a/constants.py b/constants.py index fed42eb4..7aeccdcf 100644 --- a/constants.py +++ b/constants.py @@ -59,7 +59,6 @@ TEX_DIR = os.path.join(FILE_DIR, "Tex") TEX_IMAGE_DIR = os.path.join(IMAGE_DIR, "Tex") MOBJECT_DIR = os.path.join(FILE_DIR, "mobjects") IMAGE_MOBJECT_DIR = os.path.join(MOBJECT_DIR, "image") -TMP_IMAGE_DIR = "/tmp/animation_images/" for folder in [IMAGE_DIR, GIF_DIR, MOVIE_DIR, TEX_DIR, TMP_IMAGE_DIR, TEX_IMAGE_DIR, MOBJECT_DIR, IMAGE_MOBJECT_DIR]: diff --git a/displayer.py b/displayer.py index 8633777e..cad94176 100644 --- a/displayer.py +++ b/displayer.py @@ -1,8 +1,9 @@ import numpy as np import itertools as it +import subprocess as sp import os +import sys from PIL import Image -import subprocess import cv2 from colour import Color import progressbar @@ -10,6 +11,8 @@ import progressbar from mobject import * from constants import * +FFMPEG_BIN = "ffmpeg" + def get_pixels(image_array): #TODO, FIX WIDTH/HEIGHT PROBLEM HERE if image_array is None: return np.zeros( @@ -126,97 +129,34 @@ def write_to_gif(scene, name): os.system("rm " + temppath) def write_to_movie(scene, name): - frames = scene.frames - progress_bar = progressbar.ProgressBar(maxval=len(frames)) - progress_bar.start() - print "writing " + name + "..." + filepath = os.path.join(MOVIE_DIR, name) + ".mp4" + print "Writing to %s"%filepath - filepath = os.path.join(MOVIE_DIR, name) - filedir = "/".join(filepath.split("/")[:-1]) - if not os.path.exists(filedir): - os.makedirs(filedir) - rate = int(1/scene.frame_duration) + fps = int(1/scene.display_config["frame_duration"]) + dim = (scene.display_config["width"], scene.display_config["height"]) - tmp_stem = os.path.join(TMP_IMAGE_DIR, name.replace("/", "_")) - suffix = "-%04d.png" - image_files = [] - for frame, count in zip(frames, it.count()): - progress_bar.update(int(0.9 * count)) - Image.fromarray(frame).save(tmp_stem + suffix%count) - image_files.append(tmp_stem + suffix%count) - commands = [ - "ffmpeg", - "-y", - "-loglevel", - "error", - "-i", - tmp_stem + suffix, - "-c:v", - "libx264", - "-vf", - "fps=%d,format=yuv420p"%rate, - filepath + ".mp4" + command = [ + FFMPEG_BIN, + '-y', # overwrite output file if it exists + '-f', 'rawvideo', + '-vcodec','rawvideo', + '-s', '%dx%d'%dim, # size of one frame + '-pix_fmt', 'rgb24', + '-r', str(fps), # frames per second + '-i', '-', # The imput comes from a pipe + '-an', # Tells FFMPEG not to expect any audio + '-vcodec', 'mpeg', + '-c:v', 'libx264', + '-pix_fmt', 'yuv420p', + '-loglevel', 'error', + filepath, ] - os.system(" ".join(commands)) - for image_file in image_files: - os.remove(image_file) - progress_bar.finish() + process = sp.Popen(command, stdin=sp.PIPE) + for frame in scene.frames: + process.stdin.write(frame.tostring()) - - # vs = VideoSink(scene.shape, filepath, rate) - # for frame in frames: - # vs.run(frame) - # vs.close() - # progress_bar.finish() - - - # filepath = os.path.join(MOVIE_DIR, name + ".mov") - # fourcc = cv2.cv.FOURCC(*"8bps") - # out = cv2.VideoWriter( - # filepath, fourcc, 1.0/scene.frame_duration, (DEFAULT_WIDTH, DEFAULT_HEIGHT), True - # ) - # progress = 0 - # for frame in frames: - # if progress == 0: - # print "Writing movie" - # progress_bar.update(progress) - # r, g, b = cv2.split(np.array(frame)) - # bgr_frame = cv2.merge([b, g, r]) - # out.write(bgr_frame) - # progress += 1 - # out.release() - # progress_bar.finish() - - -# class VideoSink(object): -# def __init__(self, size, filename="output", rate=10, byteorder="bgra") : -# self.size = size -# cmdstring = [ -# 'mencoder', -# '/dev/stdin', -# '-demuxer', 'rawvideo', -# '-rawvideo', 'w=%i:h=%i'%size[::-1]+":fps=%i:format=%s"%(rate,byteorder), -# '-o', filename+'.mp4', -# '-ovc', 'lavc', -# ] -# self.p = subprocess.Popen(cmdstring, stdin=subprocess.PIPE, shell=False) - -# def run(self, image): -# """ -# Image comes in as HEIGHTxWIDTHx3 numpy array, order rgb -# """ -# assert image.shape == self.size + (3,) -# r, g, b = [image[:,:,i].astype('uint32') for i in range(3)] -# a = np.ones(image.shape[:2], dtype = 'uint32') -# #hacky -# image = sum([ -# arr << 8**i -# for arr, i in zip(range(4), [a, r, g, b]) -# ]) -# self.p.stdin.write(image.tostring()) - -# def close(self): -# self.p.stdin.close() + process.stdin.close() + process.wait() diff --git a/mobject/function_graphs.py b/mobject/function_graphs.py index 08e8ebbc..cac7a577 100644 --- a/mobject/function_graphs.py +++ b/mobject/function_graphs.py @@ -196,6 +196,7 @@ class NumberPlane(Mobject1D): def generate_points(self): color = self.color faded = Color(rgb = self.fade_factor*np.array(color.get_rgb())) + #Vertical Lines freq_color_pairs = [ (self.x_faded_line_frequency, faded), (self.x_line_frequency, color) diff --git a/sample_script.py b/sample_script.py index f0b79852..8abb4098 100644 --- a/sample_script.py +++ b/sample_script.py @@ -16,10 +16,10 @@ from script_wrapper import command_line_create_scene class SampleScene(Scene): def construct(self): - plane = NumberPlane() - arrow1 = Arrow(ORIGIN, UP, color = "green", point_thickness = 5) - arrow2 = Arrow(ORIGIN, LEFT, color = "Red") - self.add(plane, arrow1, arrow2) + circle = Circle().repeat(6) + self.play(Transform(circle, Square(), run_time = 3)) + self.dither() + if __name__ == "__main__": command_line_create_scene() \ No newline at end of file diff --git a/scene/scene.py b/scene/scene.py index 8d33852a..31de65a3 100644 --- a/scene/scene.py +++ b/scene/scene.py @@ -51,6 +51,7 @@ class Scene(object): return self.name return self.__class__.__name__ + \ self.args_to_string(*self.construct_args) + def set_name(self, name): self.name = name return self diff --git a/script_wrapper.py b/script_wrapper.py index 62f64fb6..a9d1f017 100644 --- a/script_wrapper.py +++ b/script_wrapper.py @@ -3,6 +3,7 @@ import getopt import imp import itertools as it import inspect +import traceback from helpers import cammel_case_initials from scene import Scene @@ -58,8 +59,8 @@ def get_configuration(sys_argv): config["write_all"] = True config["quiet"] = True #By default, write to file - actions = set(["write", "preview", "save_image"]) - if len(actions.intersection(config)) == 0: + actions = ["write", "preview", "save_image"] + if not any([config[key] for key in actions]): config["write"] = True if len(args) > 0: @@ -78,13 +79,15 @@ def handle_scene(scene, **config): if config["write"]: scene.write_to_movie(os.path.join(config["movie_prefix"], name)) if config["save_image"]: - scene.show_frame() + if not config["write_all"]: + scene.show_frame() path = os.path.join(MOVIE_DIR, config["movie_prefix"], "images") if not os.path.exists(path): os.mkdir(path) scene.save_image(path, name) if config["quiet"]: + sys.stdout.close() sys.stdout = curr_stdout def prompt_user_for_args(SceneClass): @@ -143,7 +146,14 @@ def command_line_create_scene(movie_prefix = ""): if type(args) is not tuple: args = (args,) scene_kwargs["construct_args"] = args - handle_scene(SceneClass(**scene_kwargs), **config) + try: + handle_scene(SceneClass(**scene_kwargs), **config) + except: + print "\n\n" + traceback.print_exc() + print "\n\n" + + diff --git a/scripts/matrix_as_transform_2d.py b/scripts/matrix_as_transform_2d.py index 2477016f..30607d1b 100644 --- a/scripts/matrix_as_transform_2d.py +++ b/scripts/matrix_as_transform_2d.py @@ -15,6 +15,9 @@ from script_wrapper import command_line_create_scene MOVIE_PREFIX = "matrix_as_transform_2d/" +ARROW_CONFIG = {"point_thickness" : 2*DEFAULT_POINT_THICKNESS} +LIGHT_RED = "#F40" + def matrix_to_string(matrix): return "--".join(["-".join(map(str, row)) for row in matrix]) @@ -44,9 +47,21 @@ class ShowMultiplication(NumberLineScene): def args_to_string(num, show_original_line): end_string = "WithCopiedOriginalLine" if show_original_line else "" return str(num) + end_string + @staticmethod + def string_to_args(string): + parts = string.split() + if len(parts) == 2: + num, original_line = parts + show_original_line = original_line == "WithCopiedOriginalLine" + return float(num), False + else: + return float(parts[0]), False def construct(self, num, show_original_line): - config = {"density" : abs(num)*DEFAULT_POINT_DENSITY_1D} + config = { + "density" : abs(num)*DEFAULT_POINT_DENSITY_1D, + "point_thickness" : 2*DEFAULT_POINT_THICKNESS + } if abs(num) < 1: config["numerical_radius"] = SPACE_WIDTH/num @@ -54,7 +69,7 @@ class ShowMultiplication(NumberLineScene): if show_original_line: self.copy_original_line() self.dither() - self.show_multiplication(num, run_time = 2.0) + self.show_multiplication(num, run_time = 1.5) self.dither() def copy_original_line(self): @@ -62,7 +77,7 @@ class ShowMultiplication(NumberLineScene): copied_num_mobs = deepcopy(self.number_mobs) self.play( ApplyFunction( - lambda m : m.shift(DOWN).highlight("green"), + lambda m : m.shift(DOWN).highlight("lightgreen"), copied_line ), *[ ApplyMethod(mob.shift, DOWN) @@ -90,7 +105,7 @@ class ExamplesOfNonlinearOneDimensionalTransforms(NumberLineScene): def shift_zero((x, y, z)): return (2*x+4, y, z) self.nonlinear = text_mobject("Not a Linear Transform") - self.nonlinear.highlight("red").to_edge(UP) + self.nonlinear.highlight(LIGHT_RED).to_edge(UP, buff = 1.5) pairs = [ (sinx_plux_x, "numbers don't remain evenly spaced"), (shift_zero, "zero does not remain fixed") @@ -104,7 +119,7 @@ class ExamplesOfNonlinearOneDimensionalTransforms(NumberLineScene): self.clear() self.add(self.nonlinear) NumberLineScene.construct(self) - words = text_mobject(explanation).highlight("red") + words = text_mobject(explanation).highlight(LIGHT_RED) words.next_to(self.nonlinear, DOWN, buff = 0.5) self.add(words) @@ -136,7 +151,6 @@ class ShowTwoThenThree(ShowMultiplication): self.dither() - ######################################################## class TransformScene2D(Scene): @@ -144,7 +158,8 @@ class TransformScene2D(Scene): config = { "x_radius" : 2*SPACE_WIDTH, "y_radius" : 2*SPACE_HEIGHT, - "density" : DEFAULT_POINT_DENSITY_1D*density_factor + "density" : DEFAULT_POINT_DENSITY_1D*density_factor, + "point_thickness" : 2*DEFAULT_POINT_THICKNESS } if not use_faded_lines: config["x_faded_line_frequency"] = None @@ -153,20 +168,22 @@ class TransformScene2D(Scene): self.add(self.number_plane) def add_background(self): - self.paint_into_background( - NumberPlane(color = "grey").add_coordinates() - ) + grey_plane = NumberPlane(color = "grey") + num_mobs = grey_plane.get_coordinate_labels() + self.paint_into_background(grey_plane, *num_mobs) def add_x_y_arrows(self): self.x_arrow = Arrow( ORIGIN, self.number_plane.num_pair_to_point((1, 0)), - color = "green" + color = "lightgreen", + **ARROW_CONFIG ) self.y_arrow = Arrow( ORIGIN, self.number_plane.num_pair_to_point((0, 1)), - color = "red" + color = LIGHT_RED, + **ARROW_CONFIG ) self.add(self.x_arrow, self.y_arrow) self.number_plane.filter_out( @@ -180,6 +197,8 @@ class TransformScene2D(Scene): class ShowMatrixTransform(TransformScene2D): args_list = [ + ([[1, 3], [-2, 0]], False, False), + ([[1, 3], [-2, 0]], True, False), ([[1, 0.5], [0.5, 1]], True, False), ([[2, 0], [0, 2]], True, False), ([[0.5, 0], [0, 0.5]], True, False), @@ -223,7 +242,8 @@ class ShowMatrixTransform(TransformScene2D): new_arrow = Arrow( ORIGIN, self.number_plane.num_pair_to_point(matrix[:,index]), - color = arrow.get_color() + color = arrow.get_color(), + **ARROW_CONFIG ) arrow.remove_tip() new_arrow.remove_tip() @@ -233,8 +253,6 @@ class ShowMatrixTransform(TransformScene2D): anims.append(Transform(arrow, new_arrow, **kwargs)) self.play(*anims) self.dither() - self.set_name(str(self) + self.args_to_string(matrix, with_background, show_matrix)) - self.save_image(os.path.join(MOVIE_DIR, MOVIE_PREFIX, "images")) def get_density_factor(self, matrix): max_norm = max([ @@ -288,7 +306,7 @@ class ExamplesOfNonlinearTwoDimensionalTransformations(Scene): def shift_zero((x, y, z)): return (2*x + 3*y + 4, -1*x+y+2, z) self.nonlinear = text_mobject("Nonlinear Transform") - self.nonlinear.highlight("red").to_edge(UP) + self.nonlinear.highlight(LIGHT_RED).to_edge(UP, buff = 1.5) pairs = [ (squiggle, "lines to not remain straight"), (shift_zero, "the origin does not remain fixed") @@ -305,7 +323,7 @@ class ExamplesOfNonlinearTwoDimensionalTransformations(Scene): density = 3*DEFAULT_POINT_DENSITY_1D, ) numbers = number_plane.get_coordinate_labels() - words = text_mobject(explanation).highlight("red") + words = text_mobject(explanation).highlight(LIGHT_RED) words.next_to(self.nonlinear, DOWN, buff = 0.5) self.add(number_plane, self.nonlinear, words, *numbers) @@ -330,8 +348,8 @@ class TrickyExamplesOfNonlinearTwoDimensionalTransformations(Scene): phrase1, phrase2 = text_mobject([ "These might look like they keep lines straight...", "but diagonal lines get curved" - ]).to_edge(UP).split() - phrase2.highlight("red") + ]).to_edge(UP, buff = 1.5).split() + phrase2.highlight(LIGHT_RED) diagonal = Line( DOWN*SPACE_HEIGHT+LEFT*SPACE_WIDTH, UP*SPACE_HEIGHT+RIGHT*SPACE_WIDTH @@ -363,7 +381,7 @@ class TrickyExamplesOfNonlinearTwoDimensionalTransformations(Scene): ############# HORRIBLE! ########################## -class ShowMatrixTransformHack(TransformScene2D): +class ShowMatrixTransformWithDot(TransformScene2D): args_list = [ ([[1, 3], [-2, 0]], True, False), ] @@ -392,7 +410,7 @@ class ShowMatrixTransformHack(TransformScene2D): return mobject dot = Dot((-1, 2, 0), color = "yellow") x_arrow_copy = deepcopy(self.x_arrow) - y_arrow_copy = Arrow(LEFT, LEFT+2*UP, color = "red") + y_arrow_copy = Arrow(LEFT, LEFT+2*UP, color = LIGHT_RED, **ARROW_CONFIG) self.number_plane.add(dot) self.play(ApplyMethod(x_arrow_copy.rotate, np.pi)) @@ -409,7 +427,8 @@ class ShowMatrixTransformHack(TransformScene2D): new_arrow = Arrow( ORIGIN, self.number_plane.num_pair_to_point(matrix[:,index]), - color = arrow.get_color() + color = arrow.get_color(), + **ARROW_CONFIG ) arrow.remove_tip() new_arrow.remove_tip() @@ -421,7 +440,7 @@ class ShowMatrixTransformHack(TransformScene2D): self.dither() x_arrow_copy = deepcopy(self.x_arrow) - y_arrow_copy = Arrow(LEFT+2*UP, 5*RIGHT+2*UP, color = "red") + y_arrow_copy = Arrow(LEFT+2*UP, 5*RIGHT+2*UP, color = LIGHT_RED, **ARROW_CONFIG) self.play(ApplyMethod(x_arrow_copy.rotate, np.pi)) self.play(ShowCreation(y_arrow_copy)) self.remove(x_arrow_copy, y_arrow_copy)