Search code examples
pythonpyqtpyqt5qmediaplayerqvideowidget

PyQt5 Access Frames with QmediaPlayer


I am creating a program to play videos and then process them. I am able to play the videos with QMediaPlayer. How do I access specific frames as images or something similar. My end goal would be to format the video into a 4-d tensor of size [Num of frames, width_of_video, height_of_video, channels].

Here is the code that loads my video.:

self.clear_layout(self.vlayout)
videoItem = QVideoWidget()

self.mediaPlayer = QMediaPlayer(None, QMediaPlayer.VideoSurface)
self.mediaPlayer.durationChanged.connect(self.update_duration)     
self.mediaPlayer.positionChanged.connect(self.update_slider_position)

self.vlayout.addWidget(videoItem)
self.mediaPlayer.setVideoOutput(videoItem)
local = QUrl.fromLocalFile(self.video_paths)
media = QMediaContent(local)
self.mediaPlayer.setMedia(media)
self.play_video()

Solution

  • Here is a working example that I converted to Python from the C++ version available in this question: How to save a frame using QMediaPlayer?

    import sys
    import uuid
    
    import PyQt5
    from PyQt5 import QtCore, QtWidgets
    from PyQt5.QtCore import Qt, QObject, QUrl, QRect, pyqtSignal, QPoint
    from PyQt5.QtGui import QPainter, QImage
    from PyQt5.QtWidgets import QWidget, QApplication, QMainWindow, QGridLayout, QToolBar, QAction
    from PyQt5.QtMultimedia import QMediaPlayer, QMediaContent, QAbstractVideoBuffer, \
        QVideoFrame, QVideoSurfaceFormat, QAbstractVideoSurface
    from PyQt5.QtMultimediaWidgets import QVideoWidget
    
    
    class VideoFrameGrabber(QAbstractVideoSurface):
        frameAvailable = pyqtSignal(QImage)
    
        def __init__(self, widget: QWidget, parent: QObject):
            super().__init__(parent)
    
            self.widget = widget
    
        def supportedPixelFormats(self, handleType):
            return [QVideoFrame.Format_ARGB32, QVideoFrame.Format_ARGB32_Premultiplied,
                    QVideoFrame.Format_RGB32, QVideoFrame.Format_RGB24, QVideoFrame.Format_RGB565,
                    QVideoFrame.Format_RGB555, QVideoFrame.Format_ARGB8565_Premultiplied,
                    QVideoFrame.Format_BGRA32, QVideoFrame.Format_BGRA32_Premultiplied, QVideoFrame.Format_BGR32,
                    QVideoFrame.Format_BGR24, QVideoFrame.Format_BGR565, QVideoFrame.Format_BGR555,
                    QVideoFrame.Format_BGRA5658_Premultiplied, QVideoFrame.Format_AYUV444,
                    QVideoFrame.Format_AYUV444_Premultiplied, QVideoFrame.Format_YUV444,
                    QVideoFrame.Format_YUV420P, QVideoFrame.Format_YV12, QVideoFrame.Format_UYVY,
                    QVideoFrame.Format_YUYV, QVideoFrame.Format_NV12, QVideoFrame.Format_NV21,
                    QVideoFrame.Format_IMC1, QVideoFrame.Format_IMC2, QVideoFrame.Format_IMC3,
                    QVideoFrame.Format_IMC4, QVideoFrame.Format_Y8, QVideoFrame.Format_Y16,
                    QVideoFrame.Format_Jpeg, QVideoFrame.Format_CameraRaw, QVideoFrame.Format_AdobeDng]
    
        def isFormatSupported(self, format):
            imageFormat = QVideoFrame.imageFormatFromPixelFormat(format.pixelFormat())
            size = format.frameSize()
    
            return imageFormat != QImage.Format_Invalid and not size.isEmpty() and \
                   format.handleType() == QAbstractVideoBuffer.NoHandle
    
        def start(self, format: QVideoSurfaceFormat):
            imageFormat = QVideoFrame.imageFormatFromPixelFormat(format.pixelFormat())
            size = format.frameSize()
    
            if imageFormat != QImage.Format_Invalid and not size.isEmpty():
                self.imageFormat = imageFormat
                self.imageSize = size
                self.sourceRect = format.viewport()
    
                super().start(format)
    
                self.widget.updateGeometry()
                self.updateVideoRect()
    
                return True
            else:
                return False
    
        def stop(self):
            self.currentFrame = QVideoFrame()
            self.targetRect = QRect()
    
            super().stop()
    
            self.widget.update()
    
        def present(self, frame):
            if frame.isValid():
                cloneFrame = QVideoFrame(frame)
                cloneFrame.map(QAbstractVideoBuffer.ReadOnly)
                image = QImage(cloneFrame.bits(), cloneFrame.width(), cloneFrame.height(),
                               QVideoFrame.imageFormatFromPixelFormat(cloneFrame.pixelFormat()))
                self.frameAvailable.emit(image)  # this is very important
                cloneFrame.unmap()
    
            if self.surfaceFormat().pixelFormat() != frame.pixelFormat() or \
                    self.surfaceFormat().frameSize() != frame.size():
                self.setError(QAbstractVideoSurface.IncorrectFormatError)
                self.stop()
    
                return False
            else:
                self.currentFrame = frame
    
                self.widget.repaint(self.targetRect)
    
                return True
    
        def updateVideoRect(self):
            size = self.surfaceFormat().sizeHint()
            size.scale(self.widget.size().boundedTo(size), Qt.KeepAspectRatio)
    
            self.targetRect = QRect(QPoint(0, 0), size)
            self.targetRect.moveCenter(self.widget.rect().center())
    
        def paint(self, painter):
            if self.currentFrame.map(QAbstractVideoBuffer.ReadOnly):
                oldTransform = self.painter.transform()
    
            if self.surfaceFormat().scanLineDirection() == QVideoSurfaceFormat.BottomToTop:
                self.painter.scale(1, -1)
                self.painter.translate(0, -self.widget.height())
    
            image = QImage(self.currentFrame.bits(), self.currentFrame.width(), self.currentFrame.height(),
                           self.currentFrame.bytesPerLine(), self.imageFormat)
    
            self.painter.drawImage(self.targetRect, image, self.sourceRect)
    
            self.painter.setTransform(oldTransform)
    
            self.currentFrame.unmap()
    
    
    class App(QApplication):
        def __init__(self, sys_argv):
            super().__init__(sys_argv)
    
            # Show main window
            self.view = QMainWindow()
    
            self.centralWidget = QWidget(self.view)
    
            self.gridLayout = QGridLayout(self.centralWidget)
            self.gridLayout.setContentsMargins(0, 0, 0, 0)
            self.gridLayout.setSpacing(0)
    
            self.video_item = QVideoWidget()
    
            self.gridLayout.addWidget(self.video_item)
    
            self.view.setCentralWidget(self.centralWidget)
    
            self.mediaPlayer = QMediaPlayer(None, QMediaPlayer.VideoSurface)
    
            self.grabber = VideoFrameGrabber(self.video_item, self)
            self.mediaPlayer.setVideoOutput(self.grabber)
    
            self.grabber.frameAvailable.connect(self.process_frame)
    
            self.mediaPlayer.durationChanged.connect(self.update_duration)
            self.mediaPlayer.positionChanged.connect(self.update_slider_position)
    
            local = QUrl.fromLocalFile('c:/temp/lorem.mp4')
            media = QMediaContent(local)
            self.mediaPlayer.setMedia(media)
            self.mediaPlayer.play()
    
            self.view.show()
    
        def process_frame(self, image):
            # Save image here
            image.save('c:/temp/{}.jpg'.format(str(uuid.uuid4())))
    
        def update_duration(self):
            pass
    
        def update_slider_position(self):
            pass
    
    
    if __name__ == '__main__':
        def except_hook(cls, exception, traceback):
            sys.__excepthook__(cls, exception, traceback)
    
    
        if hasattr(QtCore.Qt, 'AA_EnableHighDpiScaling'):
            PyQt5.QtWidgets.QApplication.setAttribute(QtCore.Qt.AA_EnableHighDpiScaling, True)
    
        if hasattr(QtCore.Qt, 'AA_UseHighDpiPixmaps'):
            PyQt5.QtWidgets.QApplication.setAttribute(QtCore.Qt.AA_UseHighDpiPixmaps, True)
    
        app = App(sys.argv)
        app.setAttribute(QtCore.Qt.AA_EnableHighDpiScaling)
    
        sys.excepthook = except_hook
        sys.exit(app.exec_())