video workaround for visdom

visdom_video.py

from tqdm import tqdm
import base64
import tempfile
import subprocess

def encode(tensor):
    L = tensor.size(0)
    H = tensor.size(1)
    W = tensor.size(2)

    t = tempfile.NamedTemporaryFile(suffix='.mp4')

    command = [ 'ffmpeg',
        '-loglevel', 'error',
        '-f', 'rawvideo',
        '-vcodec','rawvideo',
        '-s', '{}x{}'.format(W, H), # size of one frame
        '-pix_fmt', 'rgb24',
        '-r', '24', # frames per second
        '-i', '-', # The imput comes from a pipe
        '-pix_fmt', 'yuv420p',
        '-an', # Tells FFMPEG not to expect any audio
        '-vcodec', 'h264',
        '-f', 'mp4',
        '-y', # overwrite
        t.name
        ]

    proc = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE)

    output = bytes()

    frame = 0

    print("Encoding...")

    with tqdm(total=L) as bar:
        while frame < L:
            state = proc.poll()
            if state is not None:
                print('Could not call ffmpeg (see above)')
                raise IOError

            read_ready, write_ready, _ = select.select([proc.stdout], [proc.stdin], [])

            if proc.stdout in read_ready:
                buf = proc.stdout.read1(1024 * 1024)
                output += buf

            if proc.stdin in write_ready:
                proc.stdin.write(tensor[frame].numpy().tobytes())
                frame += 1
                bar.update()

        remaining_output, _ = proc.communicate()
        output += remaining_output

    data = open(t.name, 'rb').read()
    t.close()

    return data



output = encode(video)
print('Got output: {} bytes'.format(len(output)))

videodata = """
    <video controls>
        <source type="video/mp4" src="data:video/mp4;base64,{}">
        Your browser does not support the video tag.
    </video>
""".format(base64.b64encode(output).decode('utf-8'))
vis.text(text=videodata, opts=dict(title='Sequence {}'.format(i)))