/src/echonest/video.py

  1#!/usr/bin/env python
  2# encoding: utf=8
  3"""
  4video.py
  5
  6Framework that turns video into silly putty.
  7
  8Created by Robert Ochshorn on 2008-5-30.
  9Refactored by Ben Lacker on 2009-6-18.
 10Copyright (c) 2008 The Echo Nest Corporation. All rights reserved.
 11"""
 12from numpy import *
 13import os
 14import re
 15import shutil
 16import subprocess
 17import sys
 18import tempfile
 19
 20from echonest import audio
 21from pyechonest import config
 22
 23
 24class ImageSequence():
 25    def __init__(self, sequence=None, settings=None):
 26        "builds sequence from a filelist, or another ImageSequence object"
 27        self.files, self.settings = [], VideoSettings()
 28        if isinstance(sequence, ImageSequence) or issubclass(sequence.__class__, ImageSequence): #from ImageSequence
 29           self.settings, self.files = sequence.settings, sequence.files
 30        if isinstance(sequence, list): #from filelist
 31            self.files = sequence
 32        if settings is not None:
 33            self.settings = settings
 34        self._init()
 35
 36    def _init(self):
 37        "extra init settings/options (can override...)"
 38        return
 39        
 40    def __len__(self):
 41        "how many frames are in this sequence?"
 42        return len(self.files)
 43
 44    def __getitem__(self, index):
 45        index = self.indexvoodo(index)
 46        if isinstance(index, slice):
 47            return self.getslice(index)
 48        else:
 49            raise TypeError("must provide an argument of type 'slice'")
 50
 51    def getslice(self, index):
 52        "returns a slice of the frames as a new instance"
 53        if isinstance(index.start, float):
 54            index = slice(int(index.start*self.settings.fps), int(index.stop*self.settings.fps), index.step)
 55        return self.__class__(self.files[index], self.settings)
 56
 57    def indexvoodo(self, index):
 58        "converts index to frame from a variety of forms"
 59        if isinstance(index, float):
 60            return int(index*self.settings.fps)
 61        return self._indexvoodoo(index)
 62
 63    def _indexvoodoo(self, index):
 64        #obj to slice
 65        if not isinstance(index, slice) and hasattr(index, "start") and hasattr(index, "duration"):
 66            sl =  slice(index.start, index.start+index.duration)
 67            return sl
 68        #slice of objs: return slice(start.start, start.end.start+start.end.duration)
 69        if isinstance(index, slice):
 70            if hasattr(index.start, "start") and hasattr(index.stop, "duration") and hasattr(index.stop, "start"):
 71                sl = slice(index.start.start, index.stop.start+index.stop.duration)
 72                return sl
 73        return index
 74
 75    def __add__(self, imseq2):
 76        """returns an ImageSequence with the second seq appended to this
 77        one. uses settings of the self."""
 78        self.render()
 79        imseq2.render() #todo: should the render be applied here? is it destructive? can it render in the new sequence?
 80        return self.__class__(self.files + imseq2.files, self.settings)
 81
 82    def duration(self):
 83        "duration of a clip in seconds"
 84        return len(self) / float(self.settings.fps)
 85
 86    def frametoim(self, index):
 87        "return a PIL image"
 88        return self.__getitem__(index)
 89
 90    def renderframe(self, index, dest=None, replacefileinseq=True):
 91        "renders frame to destination directory. can update sequence with rendered image (default)"
 92        if dest is None:
 93            #handle, dest = tempfile.mkstemp()
 94            dest = tempfile.NamedTemporaryFile().name
 95        #copy file without loading
 96        shutil.copyfile(self.files[index], dest)
 97        #symlink file...
 98        #os.symlink(self.files[index], dest)
 99        if replacefileinseq:
100            self.files[index] = dest
101        
102    def render(self, direc=None, pre="image", replacefiles=True):
103        "renders sequence to stills. can update sequence with rendered images (default)"
104        if direc is None: 
105            #nothing to render...
106            return
107        dest = None
108        for i in xrange(len(self.files)):
109            if direc is not None:
110                dest = os.path.join(direc, pre+'%(#)06d.' % {'#':i})+self.settings.imageformat()
111            self.renderframe(i, dest, replacefiles)
112
113
114class EditableFrames(ImageSequence):
115    "Collection of frames that can be easily edited"
116
117    def fadein(self, frames):
118        "linear fade in"
119        for i in xrange(frames):
120            self[i] *= (float(i)/frames) #todo: can i do this without floats?
121
122    def fadeout(self, frames):
123        "linear fade out"
124        for i in xrange(frames):
125            self[len(self)-i-1] *= (float(i)/frames)
126
127
128class VideoSettings():
129    "simple container for video settings"
130    def __init__(self):
131        self.fps = None #SS.MM
132        self.size = None #(w,h)
133        self.aspect = None #(x,y) -> x:y
134        self.bitrate = None #kb/s
135        self.uncompressed = False
136
137    def __str__(self):
138        "format as ffmpeg commandline settings"
139        cmd = ""
140        if self.bitrate is not None:
141            #bitrate
142            cmd += " -b "+str(self.bitrate)+"k"
143        if self.fps is not None:
144            #framerate
145            cmd += " -r "+str(self.fps)
146        if self.size is not None:
147            #size
148            cmd += " -s "+str(self.size[0])+"x"+str(self.size[1])
149        if self.aspect is not None:
150            #aspect
151            cmd += " -aspect "+str(self.aspect[0])+":"+str(self.aspect[1])
152        return cmd
153
154    def imageformat(self):
155        "return a string indicating to PIL the image format"
156        if self.uncompressed:
157            return "ppm"
158        else:
159            return "jpeg"
160
161
162class SynchronizedAV():
163    "SynchronizedAV has audio and video; cuts return new SynchronizedAV objects"
164
165    def __init__(self, audio=None, video=None):
166        self.audio = audio
167        self.video = video
168
169    def __getitem__(self, index):
170        "Returns a slice as synchronized AV"
171        if isinstance(index, slice):
172            return self.getslice(index)
173        else:
174            print >> sys.stderr, "WARNING: frame-based sampling not supported for synchronized AV"
175            return None
176    
177    def getslice(self, index):
178        return SynchronizedAV(audio=self.audio[index], video=self.video[index])
179    
180    def save(self, filename):
181        audio_filename = filename + '.wav'
182        audioout = self.audio.encode(audio_filename, mp3=False)
183        self.video.render()
184        res = sequencetomovie(filename, self.video, audioout)
185        os.remove(audio_filename)
186        return res
187    
188    def saveAsBundle(self, outdir):
189        videodir = os.path.join(outdir, "video")
190        videofile = os.path.join(outdir, "source.flv")
191        audiofile = os.path.join(outdir, "audio.wav")
192        os.makedirs(videodir)
193        # audio.wav
194        audioout = self.audio.encode(audiofile, mp3=False)
195        # video frames (some may be symlinked)
196        self.video.render(dir=videodir)
197        # video file
198        print sequencetomovie(videofile, self.video, audioout)
199
200
201def loadav(videofile, verbose=True):
202    foo, audio_file = tempfile.mkstemp(".mp3")        
203    cmd = "en-ffmpeg -y -i \"" + videofile + "\" " + audio_file
204    if verbose:
205        print >> sys.stderr, cmd
206    out = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
207    res = out.communicate()
208    ffmpeg_error_check(res[1])
209    a = audio.LocalAudioFile(audio_file)
210    v = sequencefrommov(videofile)
211    return SynchronizedAV(audio=a, video=v)
212
213
214def loadavfrombundle(dir):
215    # video
216    videopath = os.path.join(dir, "video")
217    videosettings = VideoSettings()
218    videosettings.fps = 25
219    videosettings.size = (320, 240)
220    video = sequencefromdir(videopath, settings=videosettings)
221    # audio
222    audiopath = os.path.join(dir, "audio.wav")
223    analysispath = os.path.join(dir, "analysis.xml")
224    myaudio = audio.LocalAudioFile(audiopath, analysis=analysispath, samplerate=22050, numchannels=1)
225    return SynchronizedAV(audio=myaudio, video=video)
226
227
228def loadavfromyoutube(url, verbose=True):
229    """returns an editable sequence from a youtube video"""
230    #todo: cache youtube videos?
231    foo, yt_file = tempfile.mkstemp()        
232    # http://bitbucket.org/rg3/youtube-dl
233    cmd = "youtube-dl -o " + yt_file + " " + url
234    if verbose:
235        print >> sys.stderr, cmd
236    out = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
237    res = out.communicate()
238    return loadav(yt_file)
239
240
241def youtubedl(url, verbose=True):
242    """downloads a video from youtube and returns the file object"""
243    foo, yt_file = tempfile.mkstemp()        
244    # http://bitbucket.org/rg3/youtube-dl
245    cmd = "youtube-dl -o " + yt_file + " " + url
246    if verbose:
247        print >> sys.stderr, cmd
248    out = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
249    res = out.communicate()
250    return yt_file
251
252
253def getpieces(video, segs):
254    a = audio.getpieces(video.audio, segs)
255    newv = EditableFrames(settings=video.video.settings)
256    for s in segs:
257        newv += video.video[s]
258    return SynchronizedAV(audio=a, video=newv)
259
260
261def sequencefromyoutube(url, settings=None, dir=None, pre="frame-", verbose=True):
262    """returns an editable sequence from a youtube video"""
263    #todo: cache youtube videos?
264    foo, yt_file = tempfile.mkstemp()
265    # http://bitbucket.org/rg3/youtube-dl
266    cmd = "youtube-dl -o " + yt_file + " " + url
267    if verbose:
268        print >> sys.stderr, cmd
269    out = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
270    out.communicate()
271    return sequencefrommov(yt_file, settings, dir, pre)
272
273
274def sequencefromdir(dir, ext=None, settings=None):
275    """returns an image sequence with lexographically-ordered images from
276    a directory"""
277    listing = os.listdir(dir)
278    #remove files without the chosen extension
279    if ext is not None:
280        listing = filter(lambda x: x.split(".")[-1]==ext, listing)
281    listing.sort()
282    #full file paths, please
283    listing = map(lambda x: os.path.join(dir, x), listing)
284    return EditableFrames(listing, settings)
285
286
287def sequencefrommov(mov, settings=None, direc=None, pre="frame-", verbose=True):
288    """full-quality video import from stills. will save frames to
289    tempspace if no directory is given"""
290    if direc is None:
291        #make directory for jpegs
292        direc = tempfile.mkdtemp()
293    format = "jpeg"
294    if settings is not None:
295        format = settings.imageformat()
296    cmd = "en-ffmpeg -i " + mov + " -an -sameq " + os.path.join(direc, pre + "%06d." + format)
297    if verbose:
298        print >> sys.stderr, cmd
299    out = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
300    res = out.communicate()
301    ffmpeg_error_check(res[1])
302    settings = settingsfromffmpeg(res[1])
303    seq =  sequencefromdir(direc, format, settings)
304    #parse ffmpeg output for to find framerate and image size
305    #todo: did this actually happen? errorcheck ffmpeg...
306    return seq
307
308
309def sequencetomovie(outfile, seq, audio=None, verbose=True):
310    "renders sequence to a movie file, perhaps with an audio track"
311    direc = tempfile.mkdtemp()
312    seq.render(direc, "image-", False)
313    cmd = "en-ffmpeg -y " + str(seq.settings) + " -i " + os.path.join(direc, "image-%06d." + seq.settings.imageformat())
314    if audio:
315        cmd += " -i " + audio
316        cmd += " -ab %dk " % getattr(config, 'MP3_BITRATE', '64')
317    cmd += " -sameq " + outfile
318    if verbose:
319        print >> sys.stderr, cmd
320    out = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
321    res = out.communicate()
322    ffmpeg_error_check(res[1])
323
324
325def convertmov(infile, outfile=None, settings=None, verbose=True):
326    """
327    Converts a movie file to a new movie file with different settings.
328    """
329    if settings is None:
330        settings = VideoSettings()
331        settings.fps = 29.97
332        settings.size = (320, 180)
333        settings.bitrate = 200
334    if not isinstance(settings, VideoSettings):
335        raise TypeError("settings arg must be a VideoSettings object")
336    if outfile is None:
337        foo, outfile = tempfile.mkstemp(".flv")
338    cmd = "en-ffmpeg -y -i " + infile + " " + str(settings) + " -sameq " + outfile
339    if verbose:
340        print >> sys.stderr, cmd
341    out = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
342    res = out.communicate()
343    ffmpeg_error_check(res[1])
344    return outfile
345
346
347def settingsfromffmpeg(parsestring):
348    """takes ffmpeg output and returns a VideoSettings object mimicking
349    the input video"""
350    settings = VideoSettings()
351    parse = parsestring.split('\n')
352    for line in parse:
353        if "Stream #0.0" in line and "Video" in line:
354            segs = line.split(", ")
355            for seg in segs:
356                if re.match("\d*x\d*", seg):
357                    #dimensions found
358                    settings.size = map(int, seg.split(" ")[0].split('x'))
359                    if "DAR " in seg:
360                        #display aspect ratio
361                        start = seg.index("DAR ")+4
362                        end = seg.index("]", start)
363                        settings.aspect = map(int, seg[start:end].split(":"))
364                elif re.match("(\d*\.)?\d+[\s]((fps)|(tbr)|(tbc)).*", seg):
365                    #fps found
366                    #todo: what's the difference between tb(r) and tb(c)?
367                    settings.fps = float(seg.split(' ')[0])
368                elif re.match("\d*.*kb.*s", seg):
369                    #bitrate found. assume we want the same bitrate
370                    settings.bitrate = int(seg[:seg.index(" ")])
371    return settings
372
373def ffmpeg_error_check(parsestring):
374    parse = parsestring.split('\n')
375    for num, line in enumerate(parse):
376        if "Unknown format" in line or "error occur" in line:
377            raise RuntimeError("en-ffmpeg conversion error:\n\t" + "\n\t".join(parse[num:]))
378