3b1b-manim/displayer.py
2015-03-22 16:15:29 -06:00

125 lines
3.2 KiB
Python

import numpy as np
import itertools as it
import os
from PIL import Image
import cv2
from animate import *
def get_image(points, rgbs):
return Image.fromarray(get_pixels(points, rgbs))
def get_pixels(points, rgbs):
#TODO, Let z add a depth componenet?
points = points[:, :2]
#Flips y-axis
points[:,1] *= -1
#Map points to pixel space, then create pixel array first in terms
#of its flattened version
points += np.array(
[SPACE_WIDTH, SPACE_HEIGHT]*points.shape[0]
).reshape(points.shape)
points *= np.array(
[HEIGHT / (2.0 * SPACE_HEIGHT), WIDTH / (2.0 * SPACE_WIDTH)]*\
points.shape[0]
).reshape(points.shape)
points = points.astype('int')
flattener = np.array([1, WIDTH], dtype = 'int').reshape((2, 1))
indices = np.dot(points, flattener)
indices = indices.reshape(indices.size)
admissibles = (indices < HEIGHT * WIDTH) * (indices > 0)
indices = indices[admissibles]
rgbs = rgbs[admissibles]
rgbs = (rgbs * 255).astype(int)
pixels = np.zeros((HEIGHT * WIDTH, 3))
pixels[indices] = rgbs
return pixels.reshape((HEIGHT, WIDTH, 3)).astype('uint8')
def write_to_gif(scene, name):
#TODO, find better means of compression
if not name.endswith(".gif"):
name += ".gif"
filepath = os.path.join(GIF_DIR, name)
temppath = os.path.join(GIF_DIR, "Temp.gif")
print "Writing " + name + "..."
writeGif(temppath, scene.frames, scene.frame_duration)
print "Compressing..."
os.system("gifsicle -O " + temppath + " > " + filepath)
os.system("rm " + temppath)
def write_to_movie(scene, name):
#TODO, incorporate pause time
frames = scene.frames
progress_bar = progressbar.ProgressBar(maxval=len(frames))
progress_bar.start()
print "writing " + name + "..."
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))
frame.save(tmp_stem + suffix%count)
image_files.append(tmp_stem + suffix%count)
filepath = os.path.join(MOVIE_DIR, name + ".mp4")
filedir = "/".join(filepath.split("/")[:-1])
if not os.path.exists(filedir):
os.makedirs(filedir)
commands = [
"ffmpeg",
"-y",
"-loglevel",
"error",
"-i",
tmp_stem + suffix,
"-c:v",
"libx264",
"-vf",
"fps=%d,format=yuv420p"%int(1/scene.frame_duration),
filepath
]
os.system(" ".join(commands))
for image_file in image_files:
os.remove(image_file)
progress_bar.finish()
# filepath = os.path.join(MOVIE_DIR, name + ".mov")
# fourcc = cv2.cv.FOURCC(*"8bps")
# out = cv2.VideoWriter(
# filepath, fourcc, 1.0/animation.frame_duration, (WIDTH, 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()