263 lines
9.2 KiB
Python
263 lines
9.2 KiB
Python
|
import numpy as np
|
||
|
|
import scipy.interpolate
|
||
|
|
from pyqtgraph.Qt import QtGui, QtCore
|
||
|
|
|
||
|
|
class ColorMap(object):
|
||
|
|
|
||
|
|
## color interpolation modes
|
||
|
|
RGB = 1
|
||
|
|
HSV_POS = 2
|
||
|
|
HSV_NEG = 3
|
||
|
|
|
||
|
|
## boundary modes
|
||
|
|
CLIP = 1
|
||
|
|
REPEAT = 2
|
||
|
|
MIRROR = 3
|
||
|
|
|
||
|
|
## return types
|
||
|
|
BYTE = 1
|
||
|
|
FLOAT = 2
|
||
|
|
QCOLOR = 3
|
||
|
|
|
||
|
|
enumMap = {
|
||
|
|
'rgb': RGB,
|
||
|
|
'hsv+': HSV_POS,
|
||
|
|
'hsv-': HSV_NEG,
|
||
|
|
'clip': CLIP,
|
||
|
|
'repeat': REPEAT,
|
||
|
|
'mirror': MIRROR,
|
||
|
|
'byte': BYTE,
|
||
|
|
'float': FLOAT,
|
||
|
|
'qcolor': QCOLOR,
|
||
|
|
}
|
||
|
|
|
||
|
|
def __init__(self, pos, color, mode=None):
|
||
|
|
"""
|
||
|
|
========= ==============================================================
|
||
|
|
Arguments
|
||
|
|
pos Array of positions where each color is defined
|
||
|
|
color Array of RGBA colors.
|
||
|
|
Integer data types are interpreted as 0-255; float data types
|
||
|
|
are interpreted as 0.0-1.0
|
||
|
|
mode Array of color modes (ColorMap.RGB, HSV_POS, or HSV_NEG)
|
||
|
|
indicating the color space that should be used when
|
||
|
|
interpolating between stops. Note that the last mode value is
|
||
|
|
ignored. By default, the mode is entirely RGB.
|
||
|
|
========= ==============================================================
|
||
|
|
"""
|
||
|
|
self.pos = pos
|
||
|
|
self.color = color
|
||
|
|
if mode is None:
|
||
|
|
mode = np.ones(len(pos))
|
||
|
|
self.mode = mode
|
||
|
|
self.stopsCache = {}
|
||
|
|
|
||
|
|
def map(self, data, mode='byte'):
|
||
|
|
"""
|
||
|
|
Data must be either a scalar position or an array (any shape) of positions.
|
||
|
|
"""
|
||
|
|
if isinstance(mode, basestring):
|
||
|
|
mode = self.enumMap[mode.lower()]
|
||
|
|
|
||
|
|
if mode == self.QCOLOR:
|
||
|
|
pos, color = self.getStops(self.BYTE)
|
||
|
|
else:
|
||
|
|
pos, color = self.getStops(mode)
|
||
|
|
|
||
|
|
data = np.clip(data, pos.min(), pos.max())
|
||
|
|
|
||
|
|
if not isinstance(data, np.ndarray):
|
||
|
|
interp = scipy.interpolate.griddata(pos, color, np.array([data]))[0]
|
||
|
|
else:
|
||
|
|
interp = scipy.interpolate.griddata(pos, color, data)
|
||
|
|
|
||
|
|
if mode == self.QCOLOR:
|
||
|
|
if not isinstance(data, np.ndarray):
|
||
|
|
return QtGui.QColor(*interp)
|
||
|
|
else:
|
||
|
|
return [QtGui.QColor(*x) for x in interp]
|
||
|
|
else:
|
||
|
|
return interp
|
||
|
|
|
||
|
|
def mapToQColor(self, data):
|
||
|
|
return self.map(data, mode=self.QCOLOR)
|
||
|
|
|
||
|
|
def mapToByte(self, data):
|
||
|
|
return self.map(data, mode=self.BYTE)
|
||
|
|
|
||
|
|
def mapToFloat(self, data):
|
||
|
|
return self.map(data, mode=self.FLOAT)
|
||
|
|
|
||
|
|
def getGradient(self, p1=None, p2=None):
|
||
|
|
"""Return a QLinearGradient object."""
|
||
|
|
if p1 == None:
|
||
|
|
p1 = QtCore.QPointF(0,0)
|
||
|
|
if p2 == None:
|
||
|
|
p2 = QtCore.QPointF(self.pos.max()-self.pos.min(),0)
|
||
|
|
g = QtGui.QLinearGradient(p1, p2)
|
||
|
|
|
||
|
|
pos, color = self.getStops(mode=self.BYTE)
|
||
|
|
color = [QtGui.QColor(*x) for x in color]
|
||
|
|
g.setStops(zip(pos, color))
|
||
|
|
|
||
|
|
#if self.colorMode == 'rgb':
|
||
|
|
#ticks = self.listTicks()
|
||
|
|
#g.setStops([(x, QtGui.QColor(t.color)) for t,x in ticks])
|
||
|
|
#elif self.colorMode == 'hsv': ## HSV mode is approximated for display by interpolating 10 points between each stop
|
||
|
|
#ticks = self.listTicks()
|
||
|
|
#stops = []
|
||
|
|
#stops.append((ticks[0][1], ticks[0][0].color))
|
||
|
|
#for i in range(1,len(ticks)):
|
||
|
|
#x1 = ticks[i-1][1]
|
||
|
|
#x2 = ticks[i][1]
|
||
|
|
#dx = (x2-x1) / 10.
|
||
|
|
#for j in range(1,10):
|
||
|
|
#x = x1 + dx*j
|
||
|
|
#stops.append((x, self.getColor(x)))
|
||
|
|
#stops.append((x2, self.getColor(x2)))
|
||
|
|
#g.setStops(stops)
|
||
|
|
return g
|
||
|
|
|
||
|
|
def getColors(self, mode=None):
|
||
|
|
"""Return list of all colors converted to the specified mode.
|
||
|
|
If mode is None, then no conversion is done."""
|
||
|
|
if isinstance(mode, basestring):
|
||
|
|
mode = self.enumMap[mode.lower()]
|
||
|
|
|
||
|
|
color = self.color
|
||
|
|
if mode in [self.BYTE, self.QCOLOR] and color.dtype.kind == 'f':
|
||
|
|
color = (color * 255).astype(np.ubyte)
|
||
|
|
elif mode == self.FLOAT and color.dtype.kind != 'f':
|
||
|
|
color = color.astype(float) / 255.
|
||
|
|
|
||
|
|
if mode == self.QCOLOR:
|
||
|
|
color = [QtGui.QColor(*x) for x in color]
|
||
|
|
|
||
|
|
return color
|
||
|
|
|
||
|
|
def getStops(self, mode):
|
||
|
|
## Get fully-expanded set of RGBA stops in either float or byte mode.
|
||
|
|
if mode not in self.stopsCache:
|
||
|
|
color = self.color
|
||
|
|
if mode == self.BYTE and color.dtype.kind == 'f':
|
||
|
|
color = (color * 255).astype(np.ubyte)
|
||
|
|
elif mode == self.FLOAT and color.dtype.kind != 'f':
|
||
|
|
color = color.astype(float) / 255.
|
||
|
|
|
||
|
|
## to support HSV mode, we need to do a little more work..
|
||
|
|
#stops = []
|
||
|
|
#for i in range(len(self.pos)):
|
||
|
|
#pos = self.pos[i]
|
||
|
|
#color = color[i]
|
||
|
|
|
||
|
|
#imode = self.mode[i]
|
||
|
|
#if imode == self.RGB:
|
||
|
|
#stops.append((x,color))
|
||
|
|
#else:
|
||
|
|
#ns =
|
||
|
|
self.stopsCache[mode] = (self.pos, color)
|
||
|
|
return self.stopsCache[mode]
|
||
|
|
|
||
|
|
#def getColor(self, x, toQColor=True):
|
||
|
|
#"""
|
||
|
|
#Return a color for a given value.
|
||
|
|
|
||
|
|
#============= ==================================================================
|
||
|
|
#**Arguments**
|
||
|
|
#x Value (position on gradient) of requested color.
|
||
|
|
#toQColor If true, returns a QColor object, else returns a (r,g,b,a) tuple.
|
||
|
|
#============= ==================================================================
|
||
|
|
#"""
|
||
|
|
#ticks = self.listTicks()
|
||
|
|
#if x <= ticks[0][1]:
|
||
|
|
#c = ticks[0][0].color
|
||
|
|
#if toQColor:
|
||
|
|
#return QtGui.QColor(c) # always copy colors before handing them out
|
||
|
|
#else:
|
||
|
|
#return (c.red(), c.green(), c.blue(), c.alpha())
|
||
|
|
#if x >= ticks[-1][1]:
|
||
|
|
#c = ticks[-1][0].color
|
||
|
|
#if toQColor:
|
||
|
|
#return QtGui.QColor(c) # always copy colors before handing them out
|
||
|
|
#else:
|
||
|
|
#return (c.red(), c.green(), c.blue(), c.alpha())
|
||
|
|
|
||
|
|
#x2 = ticks[0][1]
|
||
|
|
#for i in range(1,len(ticks)):
|
||
|
|
#x1 = x2
|
||
|
|
#x2 = ticks[i][1]
|
||
|
|
#if x1 <= x and x2 >= x:
|
||
|
|
#break
|
||
|
|
|
||
|
|
#dx = (x2-x1)
|
||
|
|
#if dx == 0:
|
||
|
|
#f = 0.
|
||
|
|
#else:
|
||
|
|
#f = (x-x1) / dx
|
||
|
|
#c1 = ticks[i-1][0].color
|
||
|
|
#c2 = ticks[i][0].color
|
||
|
|
#if self.colorMode == 'rgb':
|
||
|
|
#r = c1.red() * (1.-f) + c2.red() * f
|
||
|
|
#g = c1.green() * (1.-f) + c2.green() * f
|
||
|
|
#b = c1.blue() * (1.-f) + c2.blue() * f
|
||
|
|
#a = c1.alpha() * (1.-f) + c2.alpha() * f
|
||
|
|
#if toQColor:
|
||
|
|
#return QtGui.QColor(int(r), int(g), int(b), int(a))
|
||
|
|
#else:
|
||
|
|
#return (r,g,b,a)
|
||
|
|
#elif self.colorMode == 'hsv':
|
||
|
|
#h1,s1,v1,_ = c1.getHsv()
|
||
|
|
#h2,s2,v2,_ = c2.getHsv()
|
||
|
|
#h = h1 * (1.-f) + h2 * f
|
||
|
|
#s = s1 * (1.-f) + s2 * f
|
||
|
|
#v = v1 * (1.-f) + v2 * f
|
||
|
|
#c = QtGui.QColor()
|
||
|
|
#c.setHsv(h,s,v)
|
||
|
|
#if toQColor:
|
||
|
|
#return c
|
||
|
|
#else:
|
||
|
|
#return (c.red(), c.green(), c.blue(), c.alpha())
|
||
|
|
|
||
|
|
def getLookupTable(self, start=0.0, stop=1.0, nPts=512, alpha=None, mode='byte'):
|
||
|
|
"""
|
||
|
|
Return an RGB(A) lookup table (ndarray).
|
||
|
|
|
||
|
|
============= ============================================================================
|
||
|
|
**Arguments**
|
||
|
|
nPts The number of points in the returned lookup table.
|
||
|
|
alpha True, False, or None - Specifies whether or not alpha values are included
|
||
|
|
in the table. If alpha is None, it will be automatically determined.
|
||
|
|
============= ============================================================================
|
||
|
|
"""
|
||
|
|
if isinstance(mode, basestring):
|
||
|
|
mode = self.enumMap[mode.lower()]
|
||
|
|
|
||
|
|
if alpha is None:
|
||
|
|
alpha = self.usesAlpha()
|
||
|
|
|
||
|
|
x = np.linspace(start, stop, nPts)
|
||
|
|
table = self.map(x, mode)
|
||
|
|
|
||
|
|
if not alpha:
|
||
|
|
return table[:,:3]
|
||
|
|
else:
|
||
|
|
return table
|
||
|
|
|
||
|
|
def usesAlpha(self):
|
||
|
|
"""Return True if any stops have an alpha < 255"""
|
||
|
|
max = 1.0 if self.color.dtype.kind == 'f' else 255
|
||
|
|
return np.any(self.color[:,3] != max)
|
||
|
|
|
||
|
|
def isMapTrivial(self):
|
||
|
|
"""Return True if the gradient has exactly two stops in it: black at 0.0 and white at 1.0"""
|
||
|
|
if len(self.pos) != 2:
|
||
|
|
return False
|
||
|
|
if self.pos[0] != 0.0 or self.pos[1] != 1.0:
|
||
|
|
return False
|
||
|
|
if self.color.dtype.kind == 'f':
|
||
|
|
return np.all(self.color == np.array([[0.,0.,0.,1.], [1.,1.,1.,1.]]))
|
||
|
|
else:
|
||
|
|
return np.all(self.color == np.array([[0,0,0,255], [255,255,255,255]]))
|
||
|
|
|
||
|
|
|