140 lines
4.4 KiB
Python
140 lines
4.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
import numpy as np
|
|
|
|
import pyqtgraph as pg
|
|
|
|
try:
|
|
import cupy as cp
|
|
|
|
pg.setConfigOption("useCupy", True)
|
|
except ImportError:
|
|
cp = None
|
|
|
|
try:
|
|
import numba
|
|
except ImportError:
|
|
numba = None
|
|
|
|
|
|
def renderQImage(*args, **kwargs):
|
|
imgitem = pg.ImageItem(axisOrder='row-major')
|
|
if 'autoLevels' not in kwargs:
|
|
kwargs['autoLevels'] = False
|
|
imgitem.setImage(*args, **kwargs)
|
|
imgitem.render()
|
|
|
|
|
|
def prime_numba():
|
|
shape = (64, 64)
|
|
lut_small = np.random.randint(256, size=(256, 3), dtype=np.uint8)
|
|
lut_big = np.random.randint(256, size=(512, 3), dtype=np.uint8)
|
|
for lut in [lut_small, lut_big]:
|
|
renderQImage(np.zeros(shape, dtype=np.uint8), levels=(20, 220), lut=lut)
|
|
renderQImage(np.zeros(shape, dtype=np.uint16), levels=(250, 3000), lut=lut)
|
|
renderQImage(np.zeros(shape, dtype=np.float32), levels=(-4.0, 4.0), lut=lut)
|
|
|
|
|
|
class _TimeSuite(object):
|
|
def __init__(self):
|
|
super(_TimeSuite, self).__init__()
|
|
self.size = None
|
|
self.float_data = None
|
|
self.uint8_data = None
|
|
self.uint8_lut = None
|
|
self.uint16_data = None
|
|
self.uint16_lut = None
|
|
self.cupy_uint16_lut = None
|
|
self.cupy_uint8_lut = None
|
|
|
|
def setup(self):
|
|
size = (self.size, self.size)
|
|
self.float_data, self.uint16_data, self.uint8_data, self.uint16_lut, self.uint8_lut = self._create_data(
|
|
size, np
|
|
)
|
|
if numba is not None:
|
|
# ensure JIT compilation
|
|
pg.setConfigOption("useNumba", True)
|
|
prime_numba()
|
|
pg.setConfigOption("useNumba", False)
|
|
if cp:
|
|
_d1, _d2, _d3, self.cupy_uint16_lut, self.cupy_uint8_lut = self._create_data(size, cp)
|
|
renderQImage(cp.asarray(self.uint16_data["data"])) # prime the gpu
|
|
|
|
@property
|
|
def numba_uint16_lut(self):
|
|
return self.uint16_lut
|
|
|
|
@property
|
|
def numba_uint8_lut(self):
|
|
return self.uint8_lut
|
|
|
|
@property
|
|
def numpy_uint16_lut(self):
|
|
return self.uint16_lut
|
|
|
|
@property
|
|
def numpy_uint8_lut(self):
|
|
return self.uint8_lut
|
|
|
|
@staticmethod
|
|
def _create_data(size, xp):
|
|
float_data = {
|
|
"data": xp.random.normal(size=size).astype("float32"),
|
|
"levels": [-4.0, 4.0],
|
|
}
|
|
uint16_data = {
|
|
"data": xp.random.randint(100, 4500, size=size).astype("uint16"),
|
|
"levels": [250, 3000],
|
|
}
|
|
uint8_data = {
|
|
"data": xp.random.randint(0, 255, size=size).astype("ubyte"),
|
|
"levels": [20, 220],
|
|
}
|
|
c_map = xp.array([[-500.0, 255.0], [-255.0, 255.0], [0.0, 500.0]])
|
|
uint8_lut = xp.zeros((256, 4), dtype="ubyte")
|
|
for i in range(3):
|
|
uint8_lut[:, i] = xp.clip(xp.linspace(c_map[i][0], c_map[i][1], 256), 0, 255)
|
|
uint8_lut[:, 3] = 255
|
|
uint16_lut = xp.zeros((2 ** 16, 4), dtype="ubyte")
|
|
for i in range(3):
|
|
uint16_lut[:, i] = xp.clip(xp.linspace(c_map[i][0], c_map[i][1], 2 ** 16), 0, 255)
|
|
uint16_lut[:, 3] = 255
|
|
return float_data, uint16_data, uint8_data, uint16_lut, uint8_lut
|
|
|
|
|
|
def make_test(dtype, kind, use_levels, lut_name, func_name):
|
|
def time_test(self):
|
|
data = getattr(self, dtype + "_data")
|
|
levels = data["levels"] if use_levels else None
|
|
lut = getattr(self, f"{kind}_{lut_name}_lut", None) if lut_name is not None else None
|
|
pg.setConfigOption("useNumba", kind == "numba")
|
|
img_data = data["data"]
|
|
if kind == "cupy":
|
|
img_data = cp.asarray(img_data)
|
|
renderQImage(img_data, lut=lut, levels=levels)
|
|
|
|
time_test.__name__ = func_name
|
|
return time_test
|
|
|
|
|
|
for option in ["cupy", "numba", "numpy"]:
|
|
if option == "cupy" and cp is None:
|
|
continue
|
|
if option == "numba" and numba is None:
|
|
continue
|
|
for data_type in ["float", "uint16", "uint8"]:
|
|
for lvls in [True, False]:
|
|
if data_type == "float" and not lvls:
|
|
continue
|
|
for lutname in [None, "uint8", "uint16"]:
|
|
name = (
|
|
f'time_1x_renderImageItem_{option}_{data_type}_{"" if lvls else "no"}levels_{lutname or "no"}lut'
|
|
)
|
|
setattr(_TimeSuite, name, make_test(data_type, option, lvls, lutname, name))
|
|
|
|
|
|
class Time4096Suite(_TimeSuite):
|
|
def __init__(self):
|
|
super(Time4096Suite, self).__init__()
|
|
self.size = 4096
|