implement np2video
[numm.git] / numm / video.py
1 import subprocess
2 import numpy as np
3
4 from numm import FFMPEG
5
6 def video_frames(path, width=320, height=240, fps=30):
7     cmd = [FFMPEG, '-i', path, 
8            '-vf', 'scale=%d:%d'%(width,height),
9            '-r', str(fps),
10            '-an',
11            '-c:v', 'rawvideo', '-f', 'rawvideo',
12            '-pix_fmt', 'rgb24',
13            '-']
14     p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
15     while True:
16         arr = np.fromstring(p.stdout.read(width*height*3), dtype=np.uint8)
17         if len(arr) == 0:
18             p.wait()
19             return
20             
21         yield arr.reshape((height, width, 3))
22
23 def video2np(path, **kw):
24     "Load video data from a file."
25
26     return np.array([x for x in video_frames(path, **kw)])
27
28 def frames_to_video(generator, path, fps=30, ffopts=[]):
29     p = None 
30     for fr in generator:
31         if p is None:
32             cmd =[FFMPEG, '-y', '-s', '%dx%d' % (fr.shape[1], fr.shape[0]),
33                   '-r', str(fps), 
34                   '-an',
35                   '-c:v', 'rawvideo', '-f', 'rawvideo', 
36                   '-pix_fmt', 'rgb24',
37                   '-i', '-'] + ffopts + [path]
38             p = subprocess.Popen(cmd, stdin=subprocess.PIPE)
39         p.stdin.write(fr.tostring())
40     p.stdin.close()
41     print 'done generating video'
42     p.wait()
43
44 def np2video(np, path, fps=30, ffopts=[]):
45     """
46     Save video data to a file.
47     """
48     def framegen():
49         for fr in np:
50             yield fr
51     frames_to_video(framegen(), path, fps=fps, ffopts=ffopts)
52
53 if __name__=='__main__':
54     import sys
55     frames_to_video(video_frames(sys.argv[1]), sys.argv[2])