second_seek_example.py

"""
Note this example only really works accurately on constant frame rate media. 
"""
from PyQt4 import QtGui
from PyQt4 import QtCore
from PyQt4.QtCore import Qt

import sys
import av


AV_TIME_BASE = 1000000

def pts_to_frame(pts, time_base, frame_rate, start_time):
    return int(pts * time_base * frame_rate) - int(start_time * time_base * frame_rate)

def get_frame_rate(stream):
    
    if stream.average_rate.denominator and stream.average_rate.numerator:
        return float(stream.average_rate)
    if stream.time_base.denominator and stream.time_base.numerator:
        return 1.0/float(stream.time_base)
    else:
        raise ValueError("Unable to determine FPS")
    
def get_frame_count(f, stream):
    
    if stream.frames:
        return stream.frames
    elif stream.duration:
        return pts_to_frame(stream.duration, float(stream.time_base), get_frame_rate(stream), 0)
    elif f.duration:
        return pts_to_frame(f.duration, 1/float(AV_TIME_BASE), get_frame_rate(stream), 0)
        
    else:
        raise ValueError("Unable to determine number for frames")

class FrameGrabber(QtCore.QObject):
    
    frame_ready = QtCore.pyqtSignal(object, object)
    update_frame_range = QtCore.pyqtSignal(object, object)
    
    def __init__(self, parent =None):
        super(FrameGrabber, self).__init__(parent)
        self.file = None
        self.stream = None
        self.frame = None
        self.active_time = None
        self.start_time = 0
        self.pts_seen = False
        self.nb_frames = None
        
        self.rate = None
        self.time_base = None
        
        self.pts_map = {}
        
    def next_frame(self):
        
        frame_index = None
        
        rate = self.rate
        time_base = self.time_base
        
        self.pts_seen = False
        
        for packet in self.file.demux(self.stream):
            #print "    pkt", packet.pts, packet.dts, packet
            if packet.pts:
                self.pts_seen = True
            
            for frame in packet.decode():
                
                if frame_index is None:
                    
                    if self.pts_seen:
                        pts = frame.pts
                    else:
                        pts = frame.dts
                    
                    if not pts is None:
                        frame_index = pts_to_frame(pts, time_base, rate, self.start_time)
                        
                elif not frame_index is None:
                    frame_index += 1
                
                if not frame.dts in self.pts_map:
                    secs  = None
                    
                    if not pts is None:
                        secs = pts * time_base
                    
                    self.pts_map[frame.dts] = secs
                    
                
                #if frame.pts == None:
                    
                
                
                yield frame_index, frame
                
                
    
    @QtCore.pyqtSlot(object)
    def request_time(self, second):
        
        frame = self.get_frame(second)
        if not frame:
            return
        
        rgba = frame.reformat(frame.width, frame.height, "rgb24", 'itu709')
        #print rgba.to_image().save("test.png")
        # could use the buffer interface here instead, some versions of PyQt don't support it for some reason
        # need to track down which version they added support for it
        self.frame = bytearray(rgba.planes[0])
        bytesPerPixel  =3 
        img = QtGui.QImage(self.frame, rgba.width, rgba.height, rgba.width * bytesPerPixel, QtGui.QImage.Format_RGB888)
        
        #img = QtGui.QImage(rgba.planes[0], rgba.width, rgba.height, QtGui.QImage.Format_RGB888)

        #pixmap = QtGui.QPixmap.fromImage(img)
        self.frame_ready.emit(img, second)
        
    def get_frame(self, target_sec):

        if target_sec != self.active_time:
            return
        print 'seeking to', target_sec
        
        rate = self.rate
        time_base = self.time_base
        
        target_pts = int(target_sec / time_base) + self.start_time
        seek_pts = target_pts
        
        
        self.stream.seek(seek_pts)
        
        #frame_cache = []
        
        last_frame= None
        
        for i, (frame_index, frame) in enumerate(self.next_frame()):
            
            
            if target_sec != self.active_time:
                return
            
            pts =frame.dts
            if self.pts_seen:
                pts = frame.pts
            
            if pts > target_pts:
                break
            
            print frame.pts, seek_pts
            last_frame = frame
            
        if last_frame:
            
            return last_frame
        
        
    def get_frame_old(self, target_frame):

        if target_frame != self.active_frame:
            return
        print 'seeking to', target_frame
        
        seek_frame = target_frame

        rate = self.rate
        time_base = self.time_base
        
        frame = None
        reseek = 250
        
        original_target_frame_pts = None
        
        while reseek >= 0:
            
            # convert seek_frame to pts
            target_sec = seek_frame * 1/rate
            target_pts = int(target_sec / time_base) + self.start_time
            
            if original_target_frame_pts is None:
                original_target_frame_pts = target_pts
            
            self.stream.seek(int(target_pts))
            
            frame_index = None
            
            frame_cache = []
            
            for i, (frame_index, frame) in enumerate(self.next_frame()):
                
                # optimization if the time slider has changed, the requested frame no longer valid
                if target_frame != self.active_frame:
                    return
                
                print "   ", i, "at frame", frame_index, "at ts:", frame.pts,frame.dts,"target:", target_pts, 'orig', original_target_frame_pts

                if frame_index is None:
                    pass
                
                elif frame_index >= target_frame:
                    break
                
                frame_cache.append(frame)
                
            # Check if we over seeked, if we over seekd we need to seek to a earlier time
            # but still looking for the target frame
            if frame_index != target_frame:
                
                if frame_index is None:
                    over_seek = '?'
                else:
                    over_seek = frame_index - target_frame
                    if frame_index > target_frame:
                        
                        print over_seek, frame_cache
                        if over_seek <= len(frame_cache):
                            print "over seeked by %i, using cache" % over_seek
                            frame = frame_cache[-over_seek]
                            break

                    
                seek_frame -= 1
                reseek -= 1
                print "over seeked by %s, backtracking.. seeking: %i target: %i retry: %i" % (str(over_seek),  seek_frame, target_frame, reseek)

            else:
                break
        
        if reseek < 0:
            raise ValueError("seeking failed %i" % frame_index)
            
        # frame at this point should be the correct frame
        
        if frame:
            
            return frame
        
        else:
            raise ValueError("seeking failed %i" % target_frame)
        
    def get_frame_count(self):
        
        frame_count = None
        
        if self.stream.frames:
            frame_count = self.stream.frames
        elif self.stream.duration:
            frame_count =  pts_to_frame(self.stream.duration, float(self.stream.time_base), get_frame_rate(self.stream), 0)
        elif self.file.duration:
            frame_count = pts_to_frame(self.file.duration, 1/float(AV_TIME_BASE), get_frame_rate(self.stream), 0)
        else:
            raise ValueError("Unable to determine number for frames")
        
        seek_frame = frame_count
        
        retry = 100
        
        while retry:
            target_sec = seek_frame * 1/ self.rate
            target_pts = int(target_sec / self.time_base) + self.start_time
            
            self.stream.seek(int(target_pts))
            
            frame_index = None
            
            for frame_index, frame in self.next_frame():
                print frame_index, frame
                continue
            
            if not frame_index is None:
                break
            else:
                seek_frame -= 1 
                retry -= 1
                
        
        print "frame count seeked", frame_index, "container frame count", frame_count
        
        return frame_index or frame_count
    
    @QtCore.pyqtSlot(object)
    def set_file(self, path):
        self.file = av.open(path)
        self.stream = next(s for s in self.file.streams if s.type == b'video')
        self.rate = get_frame_rate(self.stream)
        self.time_base = float(self.stream.time_base)
        

        index, first_frame = next(self.next_frame())
        self.stream.seek(self.stream.start_time)

        # find the pts of the first frame
        index, first_frame = next(self.next_frame())

        if self.pts_seen:
            pts = first_frame.pts
        else:
            pts = first_frame.dts
 
        self.start_time = pts or first_frame.dts
            
        print "First pts", pts, self.stream.start_time, first_frame

        #self.nb_frames = get_frame_count(self.file, self.stream)
        self.nb_frames = self.get_frame_count()
        
        dur = None
        
        if self.stream.duration:
            dur = self.stream.duration * self.time_base
        else:
            dur = self.file.duration * 1.0 / float(AV_TIME_BASE)
        
        self.update_frame_range.emit(dur, self.rate)
        
        
        
        
        
class DisplayWidget(QtGui.QLabel):
    def __init__(self, parent=None):
        super(DisplayWidget, self).__init__(parent)
        #self.setScaledContents(True)
        self.setMinimumSize(1920/10, 1080/10)
        
        size_policy = QtGui.QSizePolicy(QtGui.QSizePolicy.Preferred, QtGui.QSizePolicy.Preferred)
        size_policy.setHeightForWidth(True)
        
        self.setSizePolicy(size_policy)
        
        self.setAlignment(Qt.AlignHCenter| Qt.AlignBottom)

        self.pixmap = None
        self.setMargin(10)
        
    def heightForWidth(self, width):
        return width * 9 / 16.0
    
    @QtCore.pyqtSlot(object, object)
    def setPixmap(self, img, index):
        #if index == self.current_index:
        self.pixmap = QtGui.QPixmap.fromImage(img)
        
        #super(DisplayWidget, self).setPixmap(self.pixmap)
        super(DisplayWidget, self).setPixmap(self.pixmap.scaled(self.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))
    
    def sizeHint(self):
        width = self.width()
        return QtCore.QSize(width, self.heightForWidth(width))
    
    def resizeEvent(self, event):
        if self.pixmap:
            super(DisplayWidget, self).setPixmap(self.pixmap.scaled(self.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))
        
    def sizeHint(self):
        return QtCore.QSize(1920/2.5,1080/2.5)
        

class VideoPlayerWidget(QtGui.QWidget):
    
    request_time = QtCore.pyqtSignal(object)
    
    load_file = QtCore.pyqtSignal(object)
    
    def __init__(self, parent=None):
        super(VideoPlayerWidget, self).__init__(parent)
        
        self.rate = None
        
        self.display = DisplayWidget()
        self.timeline = QtGui.QScrollBar(Qt.Horizontal)
        self.timeline_base = 100000
        
        self.frame_grabber = FrameGrabber()
        
        self.frame_control = QtGui.QDoubleSpinBox()
        self.frame_control.setFixedWidth(100)

        self.timeline.valueChanged.connect(self.slider_changed)
        self.frame_control.valueChanged.connect(self.frame_changed)
        
        self.request_time.connect(self.frame_grabber.request_time)
        self.load_file.connect(self.frame_grabber.set_file)

        self.frame_grabber.frame_ready.connect(self.display.setPixmap)
        self.frame_grabber.update_frame_range.connect(self.set_frame_range)
        
        self.frame_grabber_thread = QtCore.QThread()
        
        self.frame_grabber.moveToThread(self.frame_grabber_thread)
        self.frame_grabber_thread.start()
        
        control_layout = QtGui.QHBoxLayout()
        control_layout.addWidget(self.frame_control)
        control_layout.addWidget(self.timeline)
        
        layout = QtGui.QVBoxLayout()
        layout.addWidget(self.display)
        layout.addLayout(control_layout)
        self.setLayout(layout)
        self.setAcceptDrops(True)
        
    def set_file(self, path):
        #self.frame_grabber.set_file(path)
        self.load_file.emit(path)
        self.frame_changed(0)
        
    @QtCore.pyqtSlot(object, object)
    def set_frame_range(self, maximum, rate):
        print "frame range =", maximum, rate, int(maximum * self.timeline_base)
        
        self.timeline.setMaximum( int(maximum * self.timeline_base))
        
        self.frame_control.setMaximum(maximum)
        self.frame_control.setSingleStep(1/rate)
        #self.timeline.setSingleStep( int(AV_TIME_BASE * 1/rate))
        self.rate = rate
        
    def slider_changed(self, value):
        print '..', value
        self.frame_changed(value * 1.0 / float(self.timeline_base))
    
    def frame_changed(self, value):
        self.timeline.blockSignals(True)
        self.frame_control.blockSignals(True)
        
        self.timeline.setValue(int(value * self.timeline_base))
        self.frame_control.setValue( value)
        
        self.timeline.blockSignals(False)
        self.frame_control.blockSignals(False)
        
        #self.display.current_index = value
        self.frame_grabber.active_time = value
        
        self.request_time.emit(value)
        
    def keyPressEvent(self, event):
        if event.key() in (Qt.Key_Right, Qt.Key_Left):
            direction = 1
            if event.key() == Qt.Key_Left:
                direction = -1
                
            if event.modifiers() == Qt.ShiftModifier:
                print 'shift'
                direction *= 10
            
            direction = direction * 1/self.rate
                
            self.frame_changed(self.frame_control.value() + direction)
                
        else:
            super(VideoPlayerWidget,self).keyPressEvent(event)
            
    def mousePressEvent(self, event):
        # clear focus of spinbox
        focused_widget = QtGui.QApplication.focusWidget()
        if focused_widget:
            focused_widget.clearFocus()
            
        super(VideoPlayerWidget,self).mousePressEvent(event)
        
    def dragEnterEvent(self, event):
        event.accept()
        
    def dropEvent(self, event):
        
        mime = event.mimeData()
        event.accept()
        
        
        if mime.hasUrls():
            path = str(mime.urls()[0].path())
            self.set_file(path)
    def closeEvent(self, event):
        
        self.frame_grabber.active_time = -1
        self.frame_grabber_thread.quit()
        self.frame_grabber_thread.wait()
        
        for key,value in sorted(self.frame_grabber.pts_map.items()):
            print key, '=', value
        
        event.accept()
        

if __name__ == "__main__":
    app = QtGui.QApplication(sys.argv)
    window = VideoPlayerWidget()
    test_file = sys.argv[1]
    window.set_file(test_file)                      
    window.show()
    sys.exit(app.exec_())