Merge pull request #1617 from pijyoi/fastpath_argb

Optimize makeARGB for ubyte images
This commit is contained in:
Ogi Moore 2021-03-26 14:22:56 -07:00 committed by GitHub
commit 4dc76ed90e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 74 additions and 17 deletions

View File

@ -18,6 +18,8 @@ import pyqtgraph as pg
import pyqtgraph.ptime as ptime
from pyqtgraph.Qt import QtGui, QtCore, QT_LIB
pg.setConfigOption('imageAxisOrder', 'row-major')
import importlib
ui_template = importlib.import_module(f'VideoTemplate_{QT_LIB.lower()}')
@ -157,7 +159,7 @@ def mkData():
dt = xp.uint16
loc = 4096
scale = 1024
mx = 2**16
mx = 2**16 - 1
elif cacheKey[0] == 'float':
dt = xp.float32
loc = 1.0
@ -165,19 +167,25 @@ def mkData():
mx = 1.0
else:
raise ValueError(f"unable to handle dtype: {cacheKey[0]}")
chan_shape = (height, width)
if ui.rgbCheck.isChecked():
data = xp.random.normal(size=(frames,width,height,3), loc=loc, scale=scale)
data = pg.gaussianFilter(data, (0, 6, 6, 0))
frame_shape = chan_shape + (3,)
else:
data = xp.random.normal(size=(frames,width,height), loc=loc, scale=scale)
data = pg.gaussianFilter(data, (0, 6, 6))
if cacheKey[0] != 'float':
data = xp.clip(data, 0, mx)
data = data.astype(dt)
data[:, 10, 10:50] = mx
data[:, 9:12, 48] = mx
data[:, 8:13, 47] = mx
frame_shape = chan_shape
data = xp.empty((frames,) + frame_shape, dtype=dt)
view = data.reshape((-1,) + chan_shape)
for idx in range(view.shape[0]):
subdata = xp.random.normal(loc=loc, scale=scale, size=chan_shape)
# note: gaussian filtering has been removed as it slows down array
# creation greatly.
if cacheKey[0] != 'float':
xp.clip(subdata, 0, mx, out=subdata)
view[idx] = subdata
data[:, 10:50, 10] = mx
data[:, 48, 9:12] = mx
data[:, 47, 8:13] = mx
cache = {cacheKey: data} # clear to save memory (but keep one to prevent unnecessary regeneration)
data = cache[cacheKey]

View File

@ -1247,7 +1247,11 @@ def makeARGB(data, lut=None, levels=None, scale=None, useRGBA=False, output=None
order = [2,1,0,3] # for some reason, the colors line up as BGR in the final image.
# copy data into image array
if data.ndim == 2:
fastpath = try_fastpath_argb(xp, data, imgData, useRGBA)
if fastpath:
pass
elif data.ndim == 2:
# This is tempting:
# imgData[..., :3] = data[..., xp.newaxis]
# ..but it turns out this is faster:
@ -1263,11 +1267,12 @@ def makeARGB(data, lut=None, levels=None, scale=None, useRGBA=False, output=None
profile('reorder channels')
# add opaque alpha channel if needed
if data.ndim == 2 or data.shape[2] == 3:
alpha = False
imgData[..., 3] = 255
else:
if data.ndim == 3 and data.shape[2] == 4:
alpha = True
else:
alpha = False
if not fastpath: # fastpath has already filled it in
imgData[..., 3] = 255
# apply nan mask through alpha channel
if nanMask is not None:
@ -1282,6 +1287,50 @@ def makeARGB(data, lut=None, levels=None, scale=None, useRGBA=False, output=None
return imgData, alpha
def try_fastpath_argb(xp, ain, aout, useRGBA):
# we only optimize for certain cases
# return False if we did not handle it
can_handle = xp is np and ain.dtype == xp.ubyte and ain.flags['C_CONTIGUOUS']
if not can_handle:
return False
nrows, ncols = ain.shape[:2]
nchans = 1 if ain.ndim == 2 else ain.shape[2]
Format = QtGui.QImage.Format
if nchans == 1:
in_fmt = Format.Format_Grayscale8
elif nchans == 3:
in_fmt = Format.Format_RGB888
else:
in_fmt = Format.Format_RGBA8888
if useRGBA:
out_fmt = Format.Format_RGBA8888
else:
out_fmt = Format.Format_ARGB32
if in_fmt == out_fmt:
aout[:] = ain
return True
npixels_chunk = 512*1024
batch = int(npixels_chunk / ncols / nchans)
batch = max(1, batch)
row_beg = 0
while row_beg < nrows:
row_end = min(row_beg + batch, nrows)
ain_view = ain[row_beg:row_end, ...]
aout_view = aout[row_beg:row_end, ...]
qimg = QtGui.QImage(ain_view, ncols, ain_view.shape[0], ain.strides[0], in_fmt)
qimg = qimg.convertToFormat(out_fmt)
aout_view[:] = imageToArray(qimg, copy=False, transpose=False)
row_beg = row_end
return True
def makeQImage(imgData, alpha=None, copy=True, transpose=True):
"""
Turn an ARGB array into QImage.