Coverage for /home/travis/build/RobotLocomotion/director/build/src/director-build/lib/python2.7/dist-packages/director/screengrabberpanel : 0%

import PythonQt 

from PythonQt import QtCore, QtGui, QtUiTools 

from director import applogic as app 

from director.timercallback import TimerCallback 

from director.simpletimer import FPSCounter 

from director import ioUtils as io 

import director.vtkAll as vtk 

import os 

import glob 

import time 

import datetime 

import itertools 

def addWidgetsToDict(widgets, d): 

    for widget in widgets: 

        if widget.objectName: 

            d[str(widget.objectName)] = widget 

        addWidgetsToDict(widget.children(), d) 

class WidgetDict(object): 

    def __init__(self, widgets): 

        addWidgetsToDict(widgets, self.__dict__) 

class ScreenGrabberPanel(object): 

    def __init__(self, view): 

        self.view = view 

        loader = QtUiTools.QUiLoader() 

        uifile = QtCore.QFile(':/ui/ddScreenGrabber.ui') 

        assert uifile.open(uifile.ReadOnly) 

        self.frameCount = 0 

        self.widget = loader.load(uifile) 

        self.ui = WidgetDict(self.widget.children()) 

        self.ui.lockViewSizeCheck.connect('clicked()', self.onLockViewSize) 

        self.ui.screenshotOutputBrowseButton.connect('clicked()', self.onChooseScreenshotOutputDir) 

        self.ui.movieOutputBrowseButton.connect('clicked()', self.onChooseMovieOutputDir) 

        self.ui.saveScreenshotButton.connect('clicked()', self.onSaveScreenshot) 

        self.ui.recordMovieButton.connect('clicked()', self.onRecordMovie) 

        self.ui.viewSizeCombo.connect('currentIndexChanged(const QString&)', self.updateViewSize) 

        self.ui.viewHeightSpin.connect('valueChanged(int)', self.onViewSizeChanged) 

        self.ui.viewWidthSpin.connect('valueChanged(int)', self.onViewSizeChanged) 

        self.updateViewSize() 

        self.onLockViewSize() 

        self.recordTimer = QtCore.QTimer() 

        self.recordTimer.connect('timeout()', self.onRecordTimer) 

        self.fpsCounter = FPSCounter() 

        self.eventFilter = PythonQt.dd.ddPythonEventFilter() 

        self.ui.scrollArea.installEventFilter(self.eventFilter) 

        self.eventFilter.addFilteredEventType(QtCore.QEvent.Resize) 

        self.eventFilter.connect('handleEvent(QObject*, QEvent*)', self.onEvent) 

    def onEvent(self, obj, event): 

        minSize = self.ui.scrollArea.widget().minimumSizeHint.width() + self.ui.scrollArea.verticalScrollBar().width 

        self.ui.scrollArea.setMinimumWidth(minSize) 

    def dateTimeString(self): 

        return datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d_%H:%M:%S') 

    def movieOutputDirectory(self): 

        return os.path.expanduser(self.ui.movieOutputDirectory.text) 

    def screenshotOutputDirectory(self): 

        return os.path.expanduser(self.ui.screenshotOutputDirectory.text) 

    def captureRate(self): 

        return self.ui.captureRateSpin.value 

    def chooseDirectory(self): 

        return QtGui.QFileDialog.getExistingDirectory(app.getMainWindow(), "Choose directory...", self.movieOutputDirectory()) 

    def ensureDirectoryIsWritable(self, dirname): 

        if not os.path.isdir(dirname): 

            try: 

                os.makedirs(dirname) 

            except OSError: 

                app.showErrorMessage('Error creating directory: %s' % dirname) 

                return False 

        if not os.access(dirname, os.W_OK | os.X_OK): 

                app.showErrorMessage('Directory is not writable: %s' % dirname) 

                return False 

        return True 

    def onChooseScreenshotOutputDir(self): 

        newDir = self.chooseDirectory() 

        if newDir: 

            self.ui.screenshotOutputDirectory.text = newDir 

    def onChooseMovieOutputDir(self): 

        newDir = self.chooseDirectory() 

        if newDir: 

            self.ui.movieOutputDirectory.text = newDir 

    def onSaveScreenshot(self): 

        outDir = self.screenshotOutputDirectory() 

        if not self.ensureDirectoryIsWritable(outDir): 

            return 

        filename = os.path.join(outDir, 'Screenshot-' + self.dateTimeString() + '.png') 

        saveScreenshot(self.view, filename) 

        app.getMainWindow().statusBar().showMessage('Saved: ' + filename, 2000) 

    def nextMovieFileName(self): 

        filename = os.path.join(self.movieOutputDirectory(), 'frame_%07d.tiff' % self.frameCount) 

        self.frameCount += 1 

        return filename 

    def updateRecordingStats(self): 

        isRecordMode = self.ui.recordMovieButton.checked 

        currentRate = 0.0 

        writeQueue = 0.0 

        self.ui.currentRateValueLabel.setText('%.1f' % currentRate) 

        self.ui.writeQueueValueLabel.setText('%.1f' % currentRate) 

    def isRecordMode(self): 

        return self.ui.recordMovieButton.checked 

    def updateRecordingButtons(self): 

        isRecordMode = self.isRecordMode() 

        self.ui.movieOutputDirectory.setEnabled(not isRecordMode) 

        self.ui.movieOutputBrowseButton.setEnabled(not isRecordMode) 

        self.ui.captureRateSpin.setEnabled(not isRecordMode) 

        self.ui.captureRateLabel.setEnabled(not isRecordMode) 

        self.ui.moveOutputDirectoryLabel.setEnabled(not isRecordMode) 

        self.ui.currentRateLabel.setEnabled(isRecordMode) 

        self.ui.currentRateValueLabel.setEnabled(isRecordMode) 

        self.ui.writeQueueLabel.setEnabled(isRecordMode) 

        self.ui.writeQueueValueLabel.setEnabled(isRecordMode) 

        self.updateRecordingStats() 

    def onRecordMovie(self): 

        # Enforce even width number, otherwise avconv will fail 

        _width = (self.view.width if self.view.width % 2 == 0 else self.view.width + 1) 

        _height = (self.view.height if self.view.height % 2 == 0 else self.view.height + 1) 

        self.view.setFixedSize(_width, _height) 

        if self.isRecordMode(): 

            self.startRecording() 

        else: 

            self.stopRecording() 

        self.updateRecordingButtons() 

    def startRecording(self): 

        self.frameCount = 0 

        if not self.ensureDirectoryIsWritable(self.movieOutputDirectory()): 

            self.ui.recordMovieButton.checked = False 

            return 

        existingFiles = glob.glob(os.path.join(self.movieOutputDirectory(), '*.tiff')) 

        if len(existingFiles): 

            choice = QtGui.QMessageBox.question(app.getMainWindow(), 'Continue?', 

              'There are existing image files in the output directory.  They will be deleted prior to recording.  Continue?', 

              QtGui.QMessageBox.Yes | QtGui.QMessageBox.No, 

              QtGui.QMessageBox.No) 

            if choice == QtGui.QMessageBox.No: 

                self.ui.recordMovieButton.checked = False 

                return 

        for fileToRemove in existingFiles: 

            os.remove(fileToRemove) 

        self.fpsCounter.tick() 

        self.startT = time.time() 

        interval = int(round(1000.0 / self.captureRate())) 

        self.recordTimer.setInterval(interval) 

        self.recordTimer.start() 

    def stopRecording(self): 

        self.recordTimer.stop() 

        if self.frameCount > 0: 

            self.showEncodingDialog() 

    def showEncodingDialog(self): 

        msg = 'Recorded %d frames.  For encoding, use this command line:\n\n\n' % self.frameCount 

        msg += '    cd "%s"\n\n' % self.movieOutputDirectory() 

        msg += '    avconv -r %d -i frame_%%07d.tiff \\\n' % self.captureRate() 

        msg += '           -vcodec libx264 \\\n' 

        msg += '           -preset slow \\\n' 

        msg += '           -crf 18 \\\n' 

        msg += '           output.mp4\n\n\n' 

        app.showInfoMessage(msg, title='Recording Stopped') 

    def updateViewSize(self): 

        current = str(self.ui.viewSizeCombo.currentText) 

        useCustom = (current == 'Custom') 

        self.ui.viewWidthSpin.setEnabled(useCustom) 

        self.ui.viewHeightSpin.setEnabled(useCustom) 

        if useCustom: 

            return 

        else: 

            viewSize = [int(value) for value in current.split(' ')[0].split('x')] 

            self.ui.viewWidthSpin.value = viewSize[0] 

            self.ui.viewHeightSpin.value = viewSize[1] 

    def onViewSizeChanged(self): 

        self.onLockViewSize() 

    def viewSize(self): 

        return self.ui.viewWidthSpin.value, self.ui.viewHeightSpin.value 

    def lockViewSize(self): 

        self.view.setFixedSize(*self.viewSize()) 

        self.ui.viewSizeFrame.setEnabled(True) 

    def unlockViewSize(self): 

        self.ui.viewSizeFrame.setEnabled(False) 

        qtwidgetMaxViewSize = 16777215 

        self.view.setFixedSize(qtwidgetMaxViewSize, qtwidgetMaxViewSize) 

    def onLockViewSize(self): 

        if self.ui.lockViewSizeCheck.checked: 

            self.lockViewSize() 

        else: 

            self.unlockViewSize() 

    def onRecordTimer(self): 

        saveScreenshot(self.view, self.nextMovieFileName(), shouldRender=False) 

        self.fpsCounter.tick() 

        tNow = time.time() 

        if tNow - self.startT > 1.0: 

            self.startT = tNow 

            self.ui.currentRateValueLabel.text = '%.1f' % self.fpsCounter.getAverageFPS() 

def saveScreenshot(view, filename, shouldRender=True, shouldWrite=True): 

    if shouldRender: 

        view.forceRender() 

    grabber = vtk.vtkWindowToImageFilter() 

    grabber.SetInput(view.renderWindow()) 

    grabber.SetInputBufferTypeToRGB() 

    grabber.ReadFrontBufferOff() 

    grabber.SetShouldRerender(False) 

    grabber.Update() 

    if shouldWrite: 

        io.writeImage(grabber.GetOutput(), filename) 

    return grabber.GetOutput() 

def test(n=30, height=1080, aspect=16/9.0, ext='tiff', shouldRender=True, shouldWrite=True): 

    view.resize(height*aspect, height) 

    tStart = time.time() 

    tPrev = tStart 

    orbitTime = 5.0 

    speed = 360.0 / orbitTime 

    for i in xrange(n): 

        tNow = time.time() 

        elapsed = tNow - tPrev 

        tPrev = tNow 

        camera.Azimuth(elapsed * speed) 

        saveScreenshot('out_%04d.%s' % (i,ext), shouldRender, shouldWrite ) 

    elapsed = time.time() - tStart 

    print n, 'frames' 

    print '%.3f' % elapsed, 'seconds' 

    print '%.2f' % (n/elapsed), 'fps' 

def _getAction(): 

    return app.getToolBarActions()['ActionScreenGrabberPanel'] 

def init(view): 

    global panel 

    global dock 

    panel = ScreenGrabberPanel(view) 

    dock = app.addWidgetToDock(panel.widget, action=_getAction()) 

    dock.hide() 

    return panel