Want more features on Pastebin? Sign Up, it's FREE!
Guest

VideoTools

By: ali_m on Aug 4th, 2012  |  syntax: Python  |  size: 9.35 KB  |  views: 123  |  expires: Never
download  |  raw  |  embed  |  report abuse  |  print
Text below is selected. Please press Ctrl+C to copy to your clipboard. (⌘+C on Mac)
  1. import subprocess
  2. import threading
  3. from os import path,remove,devnull
  4. from sys import stdin,stdout
  5. import numpy as np
  6.  
  7. class VideoSink(object) :
  8.         """
  9.         VideoSink: write numpy array to a movie file
  10.         ------------------------------------------------------------------------
  11.  
  12.         Requires mencoder:
  13.         <http://www.mplayerhq.hu/design7/dload.html>
  14.  
  15.         Parameters:
  16.         ----------------
  17.         filename        string  The path/name of the output file
  18.  
  19.         size            tuple   The row/column dimensions of the output movie
  20.  
  21.         rate            scalar  The framerate (fps)
  22.  
  23.         colorspace      string  The color space of the output frames, 8-bit RGB
  24.                                 by default ('$ mencoder -vf format=fmt=help'
  25.                                 for a list of valid color spaces)
  26.  
  27.         codec           string  The codec to use, libavcodecs by default
  28.                                 ('$ mencoder -ovc -h' for a list of valid codecs)
  29.  
  30.         Methods:
  31.         ----------------
  32.         VideoSink(array)        Write the input array to the specified .avi file
  33.                                 - must be in the form [rows,columns,(channels)],
  34.                                 and (rows,colums) must match the 'size'
  35.                                 specified when initialising the VideoSink.
  36.  
  37.         VideoSink.close()       Close the .avi file. The file *must* be closed
  38.                                 after writing, or the header information isn't
  39.                                 written correctly
  40.  
  41.         Example useage:
  42.         ----------------
  43.         frames = np.random.random_integers(0,255,size=50,100,200,3).astype('uint8')
  44.         vsnk = VideoSink('mymovie.avi',size=frames.shape[1:3],colorspace='rgb24')
  45.         for frame in frames:
  46.                 vsnk(frame)
  47.         vsnk.close()
  48.  
  49.         Alistair Muldal, Aug 2012
  50.  
  51.         Credit to VokkiCodder for this idea
  52.         <http://vokicodder.blogspot.co.uk/2011/02/numpy-arrays-to-video.html>
  53.  
  54.         """
  55.         def __init__( self, filename='output.avi', size=(512,512), rate=25, colorspace='rgb24',codec='lavc'):
  56.  
  57.                 # row/col --> x/y by swapping order
  58.                 self.size = size[::-1]
  59.  
  60.                 cmdstring  = (  'mencoder',
  61.                                 '/dev/stdin',
  62.                                 '-demuxer', 'rawvideo',
  63.                                 '-rawvideo', 'w=%i:h=%i'%self.size+':fps=%i:format=%s'%(rate,colorspace),
  64.                                 '-o', filename,
  65.                                 '-ovc', codec,
  66.                                 '-nosound',
  67.                                 '-really-quiet'
  68.                                 )
  69.                 self.p = subprocess.Popen(cmdstring, stdin=subprocess.PIPE, shell=False)
  70.  
  71.         def __call__(self, frame) :
  72.                 assert frame.shape[0:2][::-1] == self.size
  73.                 # frame.tofile(self.p.stdin) # should be faster but it is indeed slower
  74.                 self.p.stdin.write(frame.tostring())
  75.         def close(self) :
  76.                 self.p.stdin.close()
  77.                 self.p.terminate()
  78.  
  79. class VideoSource(object):
  80.         """
  81.         VideoSource: create numpy arrays from frames in a movie file
  82.         ------------------------------------------------------------------------
  83.  
  84.         Requires mencoder:
  85.         <http://www.mplayerhq.hu/design7/dload.html>
  86.  
  87.         Parameters:
  88.         ----------------
  89.         filename        string  The path/name of the output file.
  90.  
  91.         colorspace      string  The color space of the output frames, 8-bit RGB
  92.                                 by default*.
  93.  
  94.         cache           string  Temporary cache for the raw frames,
  95.                                 '/tmp/tmpframes_<filename>' by default. This
  96.                                 file is deleted when .close() is called.
  97.  
  98.         blocking        bool    If True, VideoSource will block control until
  99.                                 all of the source frames have been decoded
  100.                                 (useful if VideoSource is being created within a
  101.                                 function). By default it will return control
  102.                                 immediately, but frames will not be accessible
  103.                                 until decoding is complete.
  104.        
  105.         *Only the following colorspaces are currently supported:
  106.         8-bit monochrome:       'y800','y8'
  107.         8-bit RGB:              'rgb24','bgr24'
  108.         8-bit RGBA:             'rgba','argb','rgb32','bgra','abgr','bgr32'
  109.  
  110.         Methods:
  111.         ----------------
  112.         VideoSource[1:10:2]     Grab frames from the source movie as numpy
  113.                                 arrays. Currently only a single set of indices
  114.                                 is supported. It is also possible to iterate
  115.                                 over frames.
  116.  
  117.         VideoSource.close()     Closes and deletes the temporary frame cache
  118.  
  119.         VideoSink.decode()      Decode the frames again and create a new cache.
  120.                                 Do this if you called .close() and want to
  121.                                 access the frames again.
  122.  
  123.         Example useage:
  124.         ----------------
  125.         vsrc = VideoSource('mymovie.avi',blocking=True)
  126.         oddframes = vsrc[1::2]                  # get item(s)
  127.         framelist = [frame for frame in vsrc]   # iterate over frames
  128.         vsrc.close()                            # close and delete cache
  129.  
  130.         Alistair Muldal, Aug 2012
  131.  
  132.         Inspired by VokkiCodder's VideoSink class
  133.         <http://vokicodder.blogspot.co.uk/2011/02/numpy-arrays-to-video.html>
  134.  
  135.         """
  136.         def __init__(self,filename,colorspace='rgb24',cache=None,blocking=False):
  137.  
  138.                 # check that the file exists
  139.                 try:
  140.                         open(filename)
  141.                 except IOError:
  142.                         raise IOError('Movie "%s" does not exist!' %filename)
  143.  
  144.                 # default cache location
  145.                 if cache == None:
  146.                         cache = '/tmp/%s_tmpframes' %path.basename(filename)
  147.  
  148.                 # get the format of the movie
  149.                 format = getformat(filename)
  150.                 nrows = int(format['ID_VIDEO_HEIGHT'])
  151.                 ncols = int(format['ID_VIDEO_WIDTH'])
  152.                 seconds = float(format['ID_LENGTH'])
  153.                 framerate = float(format['ID_VIDEO_FPS'])
  154.                 nframes = int(seconds*framerate)
  155.  
  156.                 self.colorspace = colorspace
  157.  
  158.                 # get the number of color channels (so we can know the size of a
  159.                 # single frame)
  160.                 if colorspace in ['y800','y8']:
  161.                         nchan = 1
  162.                         bitdepth = 'uint8'
  163.                         atomsize = 1
  164.                 elif colorspace in ['rgb16','rgb24','bgr16','bgr24']:
  165.                         nchan = 3
  166.                         bitdepth = 'uint8'
  167.                         atomsize = 1
  168.                 elif colorspace in ['rgba','argb','rgb32','bgra','abgr','bgr32']:
  169.                         nchan = 4
  170.                         atomsize = 1
  171.                         bitdepth = 'uint8'
  172.                 else:
  173.                         raise Exception('Sorry - "%s" is not a currently supported colorspace',colorspace)
  174.  
  175.                 self.blocking = blocking
  176.                 self.info = format
  177.                 self.filename = filename
  178.                 self.framerate = framerate
  179.                 self.duration = seconds
  180.                 self.shape = (nframes,nrows,ncols,nchan)
  181.                 self.bitdepth = bitdepth
  182.  
  183.                 self._framesize = nrows*ncols*nchan*atomsize
  184.                 self._offsets = self._getoffsets()
  185.                 self._cache_p = cache
  186.  
  187.                 # make temporary file containing the raw frames
  188.                 self.frames = self.decode()
  189.  
  190.         def __getitem__(self,key):
  191.                 """
  192.                 Grab frames from the source movie as numpy arrays
  193.                 """
  194.  
  195.                 nframes,nrows,ncols,nchan = self.shape
  196.                 bits = self.bitdepth
  197.                 rcount = self._framesize
  198.                 offsets = self._offsets
  199.                 raw = self._cache_h
  200.                 thread = self._thread
  201.  
  202.                 if thread.isAlive():
  203.                         print 'Not done decoding frames yet'
  204.                         return
  205.  
  206.                 if raw.closed:
  207.                         print 'Temporary frame cache is closed. Call .decode() to create a new frame cache.'
  208.                         return
  209.  
  210.                 if isinstance(key,tuple):
  211.                         raise IndexError('Too many indices')
  212.                 elif isinstance(key,int):
  213.                         # -ve indices read from the end
  214.                         if key < 0:
  215.                                 key = nframes+key
  216.                         indices = (key,)
  217.                 elif isinstance(key,slice):
  218.                         indices = np.arange(*key.indices(nframes),dtype='uint16')
  219.  
  220.                 outsize = (len(indices),nrows,ncols,nchan)
  221.                 framesout = np.empty(outsize,dtype=bits)
  222.  
  223.                 for ii in xrange(len(indices)):
  224.                         oo = offsets[indices[ii]]
  225.                         raw.seek(oo,0)
  226.                         frame = np.fromfile(raw,count=rcount,dtype=bits)
  227.                         framesout[ii] = frame.reshape(self.shape[1:])
  228.  
  229.                 return framesout.squeeze()
  230.  
  231.         def __iter__(self):
  232.                 """
  233.                 Iterate over the source movie, return frames as numpy arrays
  234.                 """
  235.                 nframes = self.shape[0]
  236.                 bits = self.bitdepth
  237.                 rcount = self._framesize
  238.                 raw = self._cache_h
  239.                 thread = self._thread
  240.  
  241.                 if thread.isAlive():
  242.                         print 'Not done decoding frames yet'
  243.                         return
  244.  
  245.                 if raw.closed:
  246.                         print 'Temporary frame cache is closed. Call .decode() to create a new frame cache.'
  247.                         return
  248.  
  249.                 raw.seek(0,0)
  250.                 for ii in xrange(nframes):
  251.                         frame = np.fromfile(raw,count=rcount,dtype=bits)
  252.                         yield frame.reshape(self.shape[1:])
  253.  
  254.         def close(self):
  255.                 """
  256.                 Close the temporary frame cache and delete it from the disk
  257.                 """
  258.                 print "Deleting temporary frame cache"
  259.                 self._cache_h.close()
  260.                 remove(self._cache_p)
  261.  
  262.         def decode(self):
  263.  
  264.                 # wipe any existing data in the cache
  265.                 file(self._cache_p,'w')
  266.  
  267.                 # read the .avi, dump the output in a temporary file
  268.                 cmdstring = (   'mencoder',
  269.                                 self.filename,
  270.                                 '-ac','none',
  271.                                 '-ovc','raw',
  272.                                 '-nosound',
  273.                                 '-vf','format=%s' %self.colorspace,
  274.                                 '-of','rawvideo',
  275.                                 '-o',self._cache_p,
  276.                                 '-really-quiet'
  277.                                 )
  278.  
  279.                 def _notifyOnExit(onExit, popenArgs,blocking):
  280.                         def runInThread(onExit, popenArgs):
  281.                                 proc = subprocess.Popen(popenArgs[0],**popenArgs[1])
  282.                                 proc.wait()
  283.                                 onExit()
  284.                                 return
  285.                         thread = threading.Thread(target=runInThread,args=(onExit,popenArgs))
  286.                         thread.start()
  287.                         if blocking:
  288.                                 # wait until thread completes
  289.                                 thread.join()
  290.                         return thread
  291.                 def onExit():
  292.                         print "Finished decoding frames in %s" %self.filename
  293.                         return
  294.  
  295.                 with open(devnull,"w") as fnull:
  296.                         print "Decoding frames in %s to temporary cache. Be patient..." %self.filename
  297.                         self._thread = _notifyOnExit(onExit,(cmdstring,dict(shell=False)),self.blocking)
  298.                         # self._getframes = subprocess.Popen(cmdstring,stdout=fnull,stderr=fnull,shell=False)
  299.  
  300.                 # open the temporary file for reading
  301.                 self._cache_h = file(self._cache_p,'r')
  302.  
  303.         def _getoffsets(self):
  304.                 offsets = []
  305.                 for index in xrange(self.shape[0]):
  306.                         offsets.append(self._framesize*index)
  307.                 return offsets
  308.  
  309.  
  310. def getformat(filename):
  311.         """
  312.         Grab the header information from a movie file using mplayer, return it
  313.         as a dict
  314.         """
  315.  
  316.         # use mplayer to grab the header information
  317.         cmdstring = (   'mplayer',
  318.                         '-vo','null',
  319.                         '-ao','null',
  320.                         '-frames','0',
  321.                         '-really-quiet',
  322.                         '-identify',
  323.                         filename
  324.                         )
  325.  
  326.         # suppress socket error messages
  327.         with open(devnull, "w") as fnull:
  328.                 formatstr = subprocess.check_output(cmdstring,stderr=fnull,shell=False)
  329.         lines = formatstr.splitlines()
  330.         format = {}
  331.  
  332.         for line in lines:
  333.                 name,value = line.split('=')
  334.                 format.update({name:value})
  335.  
  336.         return format
clone this paste RAW Paste Data