oxtimeline/oxgst/video.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
# vi:si:et:sw=4:sts=4:ts=4
# GPL 2008
import gobject
gobject.threads_init()

import pygst
pygst.require("0.10")
import gst

import Image
import time

from singledecodebin import SingleDecodeBin
from imagesink import ImageSink

class Video(gst.Pipeline):
    #def log(self, msg):
    #    print msg

    def __init__(self, uri, framerate='25/1'):
        gst.Pipeline.__init__(self)
        self.framerate = framerate
        self.duration = -1
        # queue of timestamps
        self._queue = []
        # queue callbacks
        self.callback = {}
        # extracted frames
        self._frames = {}

        # true only if we are prerolled
        self._ready = False

        uri = 'file://' + uri
        self.log("uri : %s" % uri)

        self.uri = uri

        self.sbin = SingleDecodeBin(caps=gst.Caps("video/x-raw-rgb;video/x-raw-yuv"),
                                    uri=self.uri)
        self.csp = gst.element_factory_make("ffmpegcolorspace")
        self.scale = gst.element_factory_make("videoscale")
        self.rate = gst.element_factory_make("videorate")
        self.queue = gst.element_factory_make("queue")

        self.sink = ImageSink()
        self.sink.connect('frame', self._frameCb)

        self.add(self.sbin, self.csp, self.queue, self.scale, self.rate, self.sink)

        self.sbin.connect('pad-added', self._sbinPadAddedCb)

        self.set_state(gst.STATE_PAUSED)
        self.get_state()
        self.width = self.sink.width
        self.height = self.sink.height
        self.framerate = self.sink.framerate
        self.getDuration()
        self.frames = int((float(self.duration) / gst.SECOND) * float(self.framerate))

    def _sbinPadAddedCb(self, unused_sbin, pad):
        self.log("pad : %s" % pad)
        pad.link(self.queue.get_pad("sink"))
        caps = pad.get_caps()
        width = caps[0]["width"]
        height = caps[0]["height"]
        
        if width % 4:
            width += 4 - width % 4
        if height % 4:
            height += 4 - height % 4
        
        self.queue.link(self.scale)
        self.scale.link(self.rate)
        self.rate.link(self.csp, gst.Caps("video/x-raw-rgb;video/x-raw-yuv,framerate=%s,width=%s,height=%s" % (self.framerate, width, height)))
        self.csp.link(self.sink)


    def _frameCb(self, unused_thsink, frame, timestamp):
        self.log("image:%s, timestamp:%s" % (frame, gst.TIME_ARGS(timestamp)))

        if not self._ready:
            # we know we're prerolled when we get the initial thumbnail
            self._ready = True
        elif timestamp in self.callback and self.callback[timestamp]:
            self.callback[timestamp](frame, timestamp)
            del self.callback[timestamp]

        if timestamp in self._queue:
            self._queue.remove(timestamp)

        if self._queue:
            # still some more thumbnails to process
            gobject.idle_add(self._getFrame, self._queue.pop(0))

    def getDuration(self):
        if self.duration < 0:
            pads = self.sink.sink_pads()
            q = gst.query_new_duration(gst.FORMAT_TIME)
            for pad in pads:
                if pad.get_peer() and pad.get_peer().query(q):
                  format, self.duration = q.parse_duration()
        return self.duration

    def getFrame(self, timestamp, callback):
        """
            Queue a frame request for the given timestamp,
            callback is called once frame is extracted.
            returns False if timestamp > duration
        """
        
        self.log("timestamp %s" % gst.TIME_ARGS(timestamp))
        if self.duration < 0:
            self.getDuration()
        
        if timestamp > self.duration:
            self.log("timestamp %s > duration %s" % (timestamp, self.duration))
            return False

        self.callback[timestamp] = callback

        if self._queue or not self._ready:
            self.log('ready')
            self._queue.append(timestamp)
        else:

            self._queue.append(timestamp)
            self._getFrame(timestamp)
        return True

    def _getFrame(self, timestamp):
        if not self._ready:
            return
        self.log("timestamp : %s" % gst.TIME_ARGS(timestamp))
        self.seek(1.0, gst.FORMAT_TIME, gst.SEEK_FLAG_FLUSH | gst.SEEK_FLAG_ACCURATE,
                  gst.SEEK_TYPE_SET, timestamp,
                  gst.SEEK_TYPE_NONE, -1)
        return False

    def frame(self, timestamp):
        self.mainloop = gobject.MainLoop()
        self._frames[timestamp] = None
        def callback(frame, timestamp):
            self._frames[timestamp] = frame
            self.mainloop.quit()
        self._quit = False
        def quit():
            if self._quit:
                self.mainloop.quit()
                return False
            else:
                self._quit = True
                return True
        gobject.timeout_add(1000, quit)
        if self.getFrame(timestamp, callback):
            self.mainloop.run()
        frame = self._frames[timestamp]
        del self._frames[timestamp]
        return frame