Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import subprocess
- import threading
- from os import path,remove,devnull
- from sys import stdin,stdout
- import numpy as np
- class VideoSink(object) :
- """
- VideoSink: write numpy array to a movie file
- ------------------------------------------------------------------------
- Requires mencoder:
- <http://www.mplayerhq.hu/design7/dload.html>
- Parameters:
- ----------------
- filename string The path/name of the output file
- size tuple The row/column dimensions of the output movie
- rate scalar The framerate (fps)
- colorspace string The color space of the output frames, 8-bit RGB
- by default ('$ mencoder -vf format=fmt=help'
- for a list of valid color spaces)
- codec string The codec to use, libavcodecs by default
- ('$ mencoder -ovc -h' for a list of valid codecs)
- Methods:
- ----------------
- VideoSink(array) Write the input array to the specified .avi file
- - must be in the form [rows,columns,(channels)],
- and (rows,colums) must match the 'size'
- specified when initialising the VideoSink.
- VideoSink.close() Close the .avi file. The file *must* be closed
- after writing, or the header information isn't
- written correctly
- Example useage:
- ----------------
- frames = np.random.random_integers(0,255,size=50,100,200,3).astype('uint8')
- vsnk = VideoSink('mymovie.avi',size=frames.shape[1:3],colorspace='rgb24')
- for frame in frames:
- vsnk(frame)
- vsnk.close()
- Alistair Muldal, Aug 2012
- Credit to VokkiCodder for this idea
- <http://vokicodder.blogspot.co.uk/2011/02/numpy-arrays-to-video.html>
- """
- def __init__( self, filename='output.avi', size=(512,512), rate=25, colorspace='rgb24',codec='lavc'):
- # row/col --> x/y by swapping order
- self.size = size[::-1]
- cmdstring = ( 'mencoder',
- '/dev/stdin',
- '-demuxer', 'rawvideo',
- '-rawvideo', 'w=%i:h=%i'%self.size+':fps=%i:format=%s'%(rate,colorspace),
- '-o', filename,
- '-ovc', codec,
- '-nosound',
- '-really-quiet'
- )
- self.p = subprocess.Popen(cmdstring, stdin=subprocess.PIPE, shell=False)
- def __call__(self, frame) :
- assert frame.shape[0:2][::-1] == self.size
- # frame.tofile(self.p.stdin) # should be faster but it is indeed slower
- self.p.stdin.write(frame.tostring())
- def close(self) :
- self.p.stdin.close()
- self.p.terminate()
- class VideoSource(object):
- """
- VideoSource: create numpy arrays from frames in a movie file
- ------------------------------------------------------------------------
- Requires mencoder:
- <http://www.mplayerhq.hu/design7/dload.html>
- Parameters:
- ----------------
- filename string The path/name of the output file.
- colorspace string The color space of the output frames, 8-bit RGB
- by default*.
- cache string Temporary cache for the raw frames,
- '/tmp/tmpframes_<filename>' by default. This
- file is deleted when .close() is called.
- blocking bool If True, VideoSource will block control until
- all of the source frames have been decoded
- (useful if VideoSource is being created within a
- function). By default it will return control
- immediately, but frames will not be accessible
- until decoding is complete.
- *Only the following colorspaces are currently supported:
- 8-bit monochrome: 'y800','y8'
- 8-bit RGB: 'rgb24','bgr24'
- 8-bit RGBA: 'rgba','argb','rgb32','bgra','abgr','bgr32'
- Methods:
- ----------------
- VideoSource[1:10:2] Grab frames from the source movie as numpy
- arrays. Currently only a single set of indices
- is supported. It is also possible to iterate
- over frames.
- VideoSource.close() Closes and deletes the temporary frame cache
- VideoSink.decode() Decode the frames again and create a new cache.
- Do this if you called .close() and want to
- access the frames again.
- Example useage:
- ----------------
- vsrc = VideoSource('mymovie.avi',blocking=True)
- oddframes = vsrc[1::2] # get item(s)
- framelist = [frame for frame in vsrc] # iterate over frames
- vsrc.close() # close and delete cache
- Alistair Muldal, Aug 2012
- Inspired by VokkiCodder's VideoSink class
- <http://vokicodder.blogspot.co.uk/2011/02/numpy-arrays-to-video.html>
- """
- def __init__(self,filename,colorspace='rgb24',cache=None,blocking=False):
- # check that the file exists
- try:
- open(filename)
- except IOError:
- raise IOError('Movie "%s" does not exist!' %filename)
- # default cache location
- if cache == None:
- cache = '/tmp/%s_tmpframes' %path.basename(filename)
- # get the format of the movie
- format = getformat(filename)
- nrows = int(format['ID_VIDEO_HEIGHT'])
- ncols = int(format['ID_VIDEO_WIDTH'])
- seconds = float(format['ID_LENGTH'])
- framerate = float(format['ID_VIDEO_FPS'])
- nframes = int(seconds*framerate)
- self.colorspace = colorspace
- # get the number of color channels (so we can know the size of a
- # single frame)
- if colorspace in ['y800','y8']:
- nchan = 1
- bitdepth = 'uint8'
- atomsize = 1
- elif colorspace in ['rgb16','rgb24','bgr16','bgr24']:
- nchan = 3
- bitdepth = 'uint8'
- atomsize = 1
- elif colorspace in ['rgba','argb','rgb32','bgra','abgr','bgr32']:
- nchan = 4
- atomsize = 1
- bitdepth = 'uint8'
- else:
- raise Exception('Sorry - "%s" is not a currently supported colorspace',colorspace)
- self.blocking = blocking
- self.info = format
- self.filename = filename
- self.framerate = framerate
- self.duration = seconds
- self.shape = (nframes,nrows,ncols,nchan)
- self.bitdepth = bitdepth
- self._framesize = nrows*ncols*nchan*atomsize
- self._offsets = self._getoffsets()
- self._cache_p = cache
- # make temporary file containing the raw frames
- self.frames = self.decode()
- def __getitem__(self,key):
- """
- Grab frames from the source movie as numpy arrays
- """
- nframes,nrows,ncols,nchan = self.shape
- bits = self.bitdepth
- rcount = self._framesize
- offsets = self._offsets
- raw = self._cache_h
- thread = self._thread
- if thread.isAlive():
- print 'Not done decoding frames yet'
- return
- if raw.closed:
- print 'Temporary frame cache is closed. Call .decode() to create a new frame cache.'
- return
- if isinstance(key,tuple):
- raise IndexError('Too many indices')
- elif isinstance(key,int):
- # -ve indices read from the end
- if key < 0:
- key = nframes+key
- indices = (key,)
- elif isinstance(key,slice):
- indices = np.arange(*key.indices(nframes),dtype='uint16')
- outsize = (len(indices),nrows,ncols,nchan)
- framesout = np.empty(outsize,dtype=bits)
- for ii in xrange(len(indices)):
- oo = offsets[indices[ii]]
- raw.seek(oo,0)
- frame = np.fromfile(raw,count=rcount,dtype=bits)
- framesout[ii] = frame.reshape(self.shape[1:])
- return framesout.squeeze()
- def __iter__(self):
- """
- Iterate over the source movie, return frames as numpy arrays
- """
- nframes = self.shape[0]
- bits = self.bitdepth
- rcount = self._framesize
- raw = self._cache_h
- thread = self._thread
- if thread.isAlive():
- print 'Not done decoding frames yet'
- return
- if raw.closed:
- print 'Temporary frame cache is closed. Call .decode() to create a new frame cache.'
- return
- raw.seek(0,0)
- for ii in xrange(nframes):
- frame = np.fromfile(raw,count=rcount,dtype=bits)
- yield frame.reshape(self.shape[1:])
- def close(self):
- """
- Close the temporary frame cache and delete it from the disk
- """
- print "Deleting temporary frame cache"
- self._cache_h.close()
- remove(self._cache_p)
- def decode(self):
- # wipe any existing data in the cache
- file(self._cache_p,'w')
- # read the .avi, dump the output in a temporary file
- cmdstring = ( 'mencoder',
- self.filename,
- '-ac','none',
- '-ovc','raw',
- '-nosound',
- '-vf','format=%s' %self.colorspace,
- '-of','rawvideo',
- '-o',self._cache_p,
- '-really-quiet'
- )
- def _notifyOnExit(onExit, popenArgs,blocking):
- def runInThread(onExit, popenArgs):
- proc = subprocess.Popen(popenArgs[0],**popenArgs[1])
- proc.wait()
- onExit()
- return
- thread = threading.Thread(target=runInThread,args=(onExit,popenArgs))
- thread.start()
- if blocking:
- # wait until thread completes
- thread.join()
- return thread
- def onExit():
- print "Finished decoding frames in %s" %self.filename
- return
- with open(devnull,"w") as fnull:
- print "Decoding frames in %s to temporary cache. Be patient..." %self.filename
- self._thread = _notifyOnExit(onExit,(cmdstring,dict(shell=False)),self.blocking)
- # self._getframes = subprocess.Popen(cmdstring,stdout=fnull,stderr=fnull,shell=False)
- # open the temporary file for reading
- self._cache_h = file(self._cache_p,'r')
- def _getoffsets(self):
- offsets = []
- for index in xrange(self.shape[0]):
- offsets.append(self._framesize*index)
- return offsets
- def getformat(filename):
- """
- Grab the header information from a movie file using mplayer, return it
- as a dict
- """
- # use mplayer to grab the header information
- cmdstring = ( 'mplayer',
- '-vo','null',
- '-ao','null',
- '-frames','0',
- '-really-quiet',
- '-identify',
- filename
- )
- # suppress socket error messages
- with open(devnull, "w") as fnull:
- formatstr = subprocess.check_output(cmdstring,stderr=fnull,shell=False)
- lines = formatstr.splitlines()
- format = {}
- for line in lines:
- name,value = line.split('=')
- format.update({name:value})
- return format
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement