3b1b-manim/displayer.py

import numpy as np
import itertools as it
import os
from PIL import Image
import subprocess
import cv2
from colour import Color
import progressbar

from mobject import *
from constants import *

def get_pixels(image_array): #TODO, FIX WIDTH/HEIGHT PROBLEM HERE
    if image_array is None:
        return np.zeros(
            (DEFAULT_HEIGHT, DEFAULT_WIDTH, 3), 
            dtype = 'uint8'
        )
    else:
        pixels = np.array(image_array).astype('uint8')
        assert len(pixels.shape) == 3 and pixels.shape[2] == 3
        return pixels

def paint_region(region, image_array = None, color = None):
    pixels = get_pixels(image_array)
    assert region.shape == pixels.shape[:2]
    if color is None:
        #Random dark color
        rgb = 0.5 * np.random.random(3)
    else:
        rgb = np.array(Color(color).get_rgb()) 
    pixels[region.bool_grid] = (255*rgb).astype('uint8')
    return pixels

def paint_mobject(mobject, image_array = None):
    return paint_mobjects([mobject], image_array)

def paint_mobjects(mobjects, image_array = None):
    pixels = get_pixels(image_array)
    height = pixels.shape[0]
    width  = pixels.shape[1]
    space_height = SPACE_HEIGHT
    space_width  = SPACE_HEIGHT * width / height
    pixels = pixels.reshape((pixels.size/3, 3))

    for mobject in mobjects:
        if mobject.get_num_points() == 0:
            continue
        points, rgbs = place_on_screen(mobject.points, mobject.rgbs, 
                                       space_width, space_height)
        #Map points to pixel space, which requires rescaling and shifting
        #Remember, 2*space_height -> height
        points[:,0] = points[:,0]*width/space_width/2 + width/2
        #Flip on y-axis as you go
        points[:,1] = -1*points[:,1]*height/space_height/2 + height/2
        points, rgbs = add_thickness(
            points.astype('int'), rgbs, 
            mobject.point_thickness, 
            width, height
        )

        flattener = np.array([[1], [width]], dtype = 'int')
        indices = np.dot(points, flattener)
        pixels[indices[:,0]] = (255*rgbs).astype(int)

    pixels = pixels.reshape((height, width, 3)).astype('uint8')
    return pixels

def add_thickness(pixel_indices, rgbs, thickness, width, height):
    """
    Imagine dragging each pixel around like a paintbrush in
    a square of pixels tickness x thickness big surrounding it
    """
    nudge_values = range(-thickness/2+1, thickness/2+1)
    original = np.array(pixel_indices)
    original_rgbs = np.array(rgbs)
    for x, y in it.product(nudge_values, nudge_values):
        if x == 0 and y == 0:
            continue
        pixel_indices = np.append(
            pixel_indices, 
            original+[x, y], 
            axis = 0
        )
        rgbs = np.append(rgbs, original_rgbs, axis = 0)
    admissibles = (pixel_indices[:,0] >= 0) & \
                  (pixel_indices[:,0] < width) & \
                  (pixel_indices[:,1] >= 0) & \
                  (pixel_indices[:,1] < height)
    return pixel_indices[admissibles], rgbs[admissibles]

    

def place_on_screen(points, rgbs, space_width, space_height):
    """
    Projects points to 2d space and remove those outside a
    the space constraints
    """
    # camera_distance = 10
    points = np.array(points[:, :2])
    # for i in range(2):
    #     points[:,i] *= camera_distance/(camera_distance-mobject.points[:,2])
    rgbs   = np.array(rgbs)
    
    #Removes points out of space
    to_keep = (abs(points[:,0]) < space_width) & \
              (abs(points[:,1]) < space_height)
    return points[to_keep], rgbs[to_keep]

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 + "..."
    images = [Image.fromarray(frame) for frame in scene.frames]
    writeGif(temppath, images, scene.frame_duration)
    print "Compressing..."
    os.system("gifsicle -O " + temppath + " > " + filepath)
    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)
    filedir = "/".join(filepath.split("/")[:-1])
    if not os.path.exists(filedir):
        os.makedirs(filedir)
    rate = int(1/scene.frame_duration)

    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"
    ]
    os.system(" ".join(commands))
    for image_file in image_files:
        os.remove(image_file)
    progress_bar.finish()


    # 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()