Skip to content

Instantly share code, notes, and snippets.

@xqms
Last active June 25, 2018 12:38
Show Gist options
  • Save xqms/443adc29b9a53d5621a481cdf3ce590d to your computer and use it in GitHub Desktop.
Save xqms/443adc29b9a53d5621a481cdf3ce590d to your computer and use it in GitHub Desktop.
video workaround for visdom
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)))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment