Aluode/PerceptionLabPortable

1"""2Tests the speed of image updates for an ImageItem and RawImageWidget.3The speed will generally depend on the type of data being shown, whether4it is being scaled and/or converted by lookup table, and whether OpenGL5is used by the view widget6"""7 8import argparse9import itertools10import sys11 12import numpy as np13from utils import FrameCounter14 15import pyqtgraph as pg16from pyqtgraph.Qt import QtCore, QtGui, QtWidgets17 18pg.setConfigOption('imageAxisOrder', 'row-major')19 20import VideoTemplate_generic as ui_template21 22try:23    import cupy as cp24    pg.setConfigOption("useCupy", True)25    _has_cupy = True26except ImportError:27    cp = None28    _has_cupy = False29 30try:31    import numba32    _has_numba = True33except ImportError:34    numba = None35    _has_numba = False36 37try:38    from pyqtgraph.widgets.RawImageWidget import RawImageGLWidget39except ImportError:40    RawImageGLWidget = None41 42parser = argparse.ArgumentParser(description="Benchmark for testing video performance")43parser.add_argument('--cuda', default=False, action='store_true', help="Use CUDA to process on the GPU", dest="cuda")44parser.add_argument('--dtype', default='uint8', choices=['uint8', 'uint16', 'float'], help="Image dtype (uint8, uint16, or float)")45parser.add_argument('--frames', default=3, type=int, help="Number of image frames to generate (default=3)")46parser.add_argument('--image-mode', default='mono', choices=['mono', 'rgb'], help="Image data mode (mono or rgb)", dest='image_mode')47parser.add_argument('--levels', default=None, type=lambda s: tuple([float(x) for x in s.split(',')]), help="min,max levels to scale monochromatic image dynamic range, or rmin,rmax,gmin,gmax,bmin,bmax to scale rgb")48parser.add_argument('--lut', default=False, action='store_true', help="Use color lookup table")49parser.add_argument('--lut-alpha', default=False, action='store_true', help="Use alpha color lookup table", dest='lut_alpha')50parser.add_argument('--size', default='512x512', type=lambda s: tuple([int(x) for x in s.split('x')]), help="WxH image dimensions default='512x512'")51parser.add_argument('--iterations', default=float('inf'), type=float,52    help="Number of iterations to run before exiting"53)54args = parser.parse_args(sys.argv[1:])55iterations_counter = itertools.count()56 57if RawImageGLWidget is not None:58    # don't limit frame rate to vsync59    sfmt = QtGui.QSurfaceFormat()60    sfmt.setSwapInterval(0)61    QtGui.QSurfaceFormat.setDefaultFormat(sfmt)62 63app = pg.mkQApp("Video Speed Test Example")64 65win = QtWidgets.QMainWindow()66win.setWindowTitle('pyqtgraph example: VideoSpeedTest')67ui = ui_template.Ui_MainWindow()68ui.setupUi(win)69win.show()70 71if RawImageGLWidget is None:72    ui.rawGLRadio.setEnabled(False)73    ui.rawGLRadio.setText(ui.rawGLRadio.text() + " (OpenGL not available)")74else:75    ui.rawGLImg = RawImageGLWidget()76    ui.stack.addWidget(ui.rawGLImg)77    win.destroyed.connect(ui.rawGLImg.cleanup)78 79# read in CLI args80ui.cudaCheck.setChecked(args.cuda and _has_cupy)81ui.cudaCheck.setEnabled(_has_cupy)82ui.numbaCheck.setChecked(_has_numba and pg.getConfigOption("useNumba"))83ui.numbaCheck.setEnabled(_has_numba)84ui.framesSpin.setValue(args.frames)85ui.widthSpin.setValue(args.size[0])86ui.heightSpin.setValue(args.size[1])87ui.dtypeCombo.setCurrentText(args.dtype)88ui.rgbCheck.setChecked(args.image_mode=='rgb')89ui.maxSpin1.setOpts(value=255, step=1)90ui.minSpin1.setOpts(value=0, step=1)91levelSpins = [ui.minSpin1, ui.maxSpin1, ui.minSpin2, ui.maxSpin2, ui.minSpin3, ui.maxSpin3]92if args.cuda and _has_cupy:93    xp = cp94else:95    xp = np96if args.levels is None:97    ui.scaleCheck.setChecked(False)98    ui.rgbLevelsCheck.setChecked(False)99else:100    ui.scaleCheck.setChecked(True)101    if len(args.levels) == 2:102        ui.rgbLevelsCheck.setChecked(False)103        ui.minSpin1.setValue(args.levels[0])104        ui.maxSpin1.setValue(args.levels[1])105    elif len(args.levels) == 6:106        ui.rgbLevelsCheck.setChecked(True)107        for spin,val in zip(levelSpins, args.levels):108            spin.setValue(val)109    else:110        raise ValueError("levels argument must be 2 or 6 comma-separated values (got %r)" % (args.levels,))111ui.lutCheck.setChecked(args.lut)112ui.alphaCheck.setChecked(args.lut_alpha)113 114 115#ui.graphicsView.useOpenGL()  ## buggy, but you can try it if you need extra speed.116 117vb = pg.ViewBox()118ui.graphicsView.setCentralItem(vb)119vb.setAspectLocked()120img = pg.ImageItem()121vb.addItem(img)122 123 124 125LUT = None126def updateLUT():127    global LUT, ui128    dtype = ui.dtypeCombo.currentText()129    if dtype == 'uint8':130        n = 256131    else:132        n = 4096133    LUT = ui.gradient.getLookupTable(n, alpha=ui.alphaCheck.isChecked())134    if _has_cupy and xp == cp:135        LUT = cp.asarray(LUT)136ui.gradient.sigGradientChanged.connect(updateLUT)137updateLUT()138 139ui.alphaCheck.toggled.connect(updateLUT)140 141def updateScale():142    global ui, levelSpins143    if ui.rgbLevelsCheck.isChecked():144        for s in levelSpins[2:]:145            s.setEnabled(True)146    else:147        for s in levelSpins[2:]:148            s.setEnabled(False)149 150updateScale()151 152ui.rgbLevelsCheck.toggled.connect(updateScale)153 154cache = {}155def mkData():156    with pg.BusyCursor():157        global data, cache, ui, xp158        frames = ui.framesSpin.value()159        width = ui.widthSpin.value()160        height = ui.heightSpin.value()161        cacheKey = (ui.dtypeCombo.currentText(), ui.rgbCheck.isChecked(), frames, width, height)162        if cacheKey not in cache:163            if cacheKey[0] == 'uint8':164                dt = xp.uint8165                loc = 128166                scale = 64167                mx = 255168            elif cacheKey[0] == 'uint16':169                dt = xp.uint16170                loc = 4096171                scale = 1024172                mx = 2**16 - 1173            elif cacheKey[0] == 'float':174                dt = xp.float32175                loc = 1.0176                scale = 0.1177                mx = 1.0178            else:179                raise ValueError(f"unable to handle dtype: {cacheKey[0]}")180 181            chan_shape = (height, width)182            if ui.rgbCheck.isChecked():183                frame_shape = chan_shape + (3,)184            else:185                frame_shape = chan_shape186            data = xp.empty((frames,) + frame_shape, dtype=dt)187            view = data.reshape((-1,) + chan_shape)188            for idx in range(view.shape[0]):189                subdata = xp.random.normal(loc=loc, scale=scale, size=chan_shape)190                # note: gaussian filtering has been removed as it slows down array191                #       creation greatly.192                if cacheKey[0] != 'float':193                    xp.clip(subdata, 0, mx, out=subdata)194                view[idx] = subdata195 196            data[:, 10:50, 10] = mx197            data[:, 48, 9:12] = mx198            data[:, 47, 8:13] = mx199            cache = {cacheKey: data} # clear to save memory (but keep one to prevent unnecessary regeneration)200 201        data = cache[cacheKey]202        updateLUT()203        updateSize()204 205def updateSize():206    global ui, vb207    frames = ui.framesSpin.value()208    width = ui.widthSpin.value()209    height = ui.heightSpin.value()210    dtype = xp.dtype(str(ui.dtypeCombo.currentText()))211    rgb = 3 if ui.rgbCheck.isChecked() else 1212    ui.sizeLabel.setText('%d MB' % (frames * width * height * rgb * dtype.itemsize / 1e6))213    vb.setRange(QtCore.QRectF(0, 0, width, height))214 215 216def noticeCudaCheck():217    global xp, cache218    cache = {}219    if ui.cudaCheck.isChecked():220        if _has_cupy:221            xp = cp222        else:223            xp = np224            ui.cudaCheck.setChecked(False)225    else:226        xp = np227    mkData()228 229 230def noticeNumbaCheck():231    pg.setConfigOption('useNumba', _has_numba and ui.numbaCheck.isChecked())232 233 234mkData()235 236 237ui.dtypeCombo.currentIndexChanged.connect(mkData)238ui.rgbCheck.toggled.connect(mkData)239ui.widthSpin.editingFinished.connect(mkData)240ui.heightSpin.editingFinished.connect(mkData)241ui.framesSpin.editingFinished.connect(mkData)242 243ui.widthSpin.valueChanged.connect(updateSize)244ui.heightSpin.valueChanged.connect(updateSize)245ui.framesSpin.valueChanged.connect(updateSize)246ui.cudaCheck.toggled.connect(noticeCudaCheck)247ui.numbaCheck.toggled.connect(noticeNumbaCheck)248 249ptr = 0250def update():251    global ptr252    if next(iterations_counter) > args.iterations:253        # cleanly close down benchmark254        timer.stop()255        app.quit()256        return None257 258    if ui.lutCheck.isChecked():259        useLut = LUT260    else:261        useLut = None262 263    downsample = ui.downsampleCheck.isChecked()264 265    if ui.scaleCheck.isChecked():266        if ui.rgbLevelsCheck.isChecked():267            useScale = [268                [ui.minSpin1.value(), ui.maxSpin1.value()],269                [ui.minSpin2.value(), ui.maxSpin2.value()],270                [ui.minSpin3.value(), ui.maxSpin3.value()]]271        else:272            useScale = [ui.minSpin1.value(), ui.maxSpin1.value()]273    else:274        useScale = None275 276    if ui.rawRadio.isChecked():277        ui.rawImg.setImage(data[ptr%data.shape[0]], lut=useLut, levels=useScale)278        ui.stack.setCurrentIndex(1)279    elif ui.rawGLRadio.isChecked():280        ui.rawGLImg.setImage(data[ptr%data.shape[0]], lut=useLut, levels=useScale)281        ui.stack.setCurrentIndex(2)282    else:283        img.setImage(data[ptr%data.shape[0]], autoLevels=False, levels=useScale, lut=useLut, autoDownsample=downsample)284        ui.stack.setCurrentIndex(0)285        #img.setImage(data[ptr%data.shape[0]], autoRange=False)286 287    ptr += 1288    framecnt.update()289 290timer = QtCore.QTimer()291timer.timeout.connect(update)292timer.start(0)293 294framecnt = FrameCounter()295framecnt.sigFpsUpdate.connect(lambda fps: ui.fpsLabel.setText(f'{fps:.1f} fps'))296 297if __name__ == '__main__':298    pg.exec()299