Skip to content
Snippets Groups Projects
make_gif.py 3.04 KiB
import matplotlib
import matplotlib.pyplot as plt
from contextlib import contextmanager
import subprocess
import io
import PIL.Image
import os
import IPython.display
from pathlib import Path
import numpy as np
import sys


def render_rounded(out_fp):
    fp = io.BytesIO()
    matplotlib.pyplot.savefig(fp, format='png', bbox_inches='tight', transparent=False)
    fp.seek(0)
    im = PIL.Image.open(fp)
    width, height = im.size
    width -= width % 8
    height -= height % 8
    im = im.resize((width,height))
    im.save(out_fp, format='png')
    
    

def render_fast(fig, out_fp):
    canvas = fig.canvas
    #fp = io.BytesIO()
    #savefig(fp, format='png', bbox_inches='tight', transparent=False)
    im = PIL.Image.fromarray(np.asarray(fig.canvas.buffer_rgba()))
    #fp.seek(0)
    #im = PIL.Image.open(fp)
    width, height = im.size
    width -= width % 8
    height -= height % 8
    im = im.resize((width,height))
    im.save(out_fp, format='png')



FFMPEG = Path(sys.executable).parent.parent.parent / 'ffmpeg/bin/ffmpeg'
#FFMPEG = '/data/cphillips/miniconda3/envs/ffmpeg/bin/ffmpeg'

@contextmanager
def make_video(out='output.mkv', framerate=10, debug=False, ffmpeg=FFMPEG, bitrate="4000k"):
    if os.path.isfile(out):
        os.unlink(out)
        
    palette = Path('palette.png').absolute()
    tmp = Path('tmp.mkv').absolute()
    if tmp.is_file():
        tmp.unlink()
    if palette.is_file():
        palette.unlink()
    args_mk_video = [ffmpeg,"-framerate",f"{framerate}",
            "-f","image2pipe","-i","-","-b:v",bitrate, "-pix_fmt","yuv420p", str(tmp)]
    args_palettegen = [ffmpeg, "-i",str(tmp),'-filter_complex','[0:v] palettegen',str(palette)]
    args_mk_gif = [ffmpeg,
            "-i",str(tmp),
            "-i",str(palette),'-filter_complex', '[0:v][1:v] paletteuse',
             str(out)]
    if debug:
        print(' '.join(args_mk_video))
        p = subprocess.Popen(args_mk_video, stdin=subprocess.PIPE)
    else:
        p = subprocess.Popen(args_mk_video, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    try:
        yield p.stdin
    finally:
        p.stdin.close()
        p.wait()
        
    if debug:
        print(' '.join(args_palettegen))
        p = subprocess.Popen(args_palettegen)
        p.wait()
        print(' '.join(args_mk_gif))
        p = subprocess.Popen(args_mk_gif)
        p.wait()
    else:
        p = subprocess.Popen(args_palettegen, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        p.wait()
        p = subprocess.Popen(args_mk_gif, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        p.wait()
        if palette.is_file():
            palette.unlink()
        if tmp.is_file():
            tmp.unlink()
            

def play_video(out, add_controls=True):
    if out.name.endswith('.gif'):
        return play_gif(out)
    if add_controls:
        controls = 'controls'
    else:
        controls = ''
    return IPython.display.HTML(data=f'<video autoplay loop {controls}><source src="{out}"/></video>')



def play_gif(out):
    return IPython.display.Image(out)