corrections and cleanups for functions.makeARGB

added unit test coverage

pyqtgraph/functions.py

@@ -902,23 +902,35 @@ def makeARGB(data, lut=None, levels=None, scale=None, useRGBA=False):
     """
     profile = debug.Profiler()
 
+    if data.ndim not in (2, 3):
+        raise TypeError("data must be 2D or 3D")
+    if data.ndim == 3 and data.shape[2] > 4:
+        raise TypeError("data.shape[2] must be <= 4")
+    
     if lut is not None and not isinstance(lut, np.ndarray):
         lut = np.array(lut)
-    if levels is not None and not isinstance(levels, np.ndarray):
-        levels = np.array(levels)
     
-    if levels is not None:
+    if levels is None:
-        if levels.ndim == 1:
+        # automatically decide levels based on data dtype
-            if len(levels) != 2:
+        if data.dtype.kind == 'u':
-                raise Exception('levels argument must have length 2')
+            levels = np.array([0, 2**(data.itemsize*8)-1])
-        elif levels.ndim == 2:
+        elif data.dtype.kind == 'i':
-            if lut is not None and lut.ndim > 1:
+            s = 2**(data.itemsize*8 - 1)
-                raise Exception('Cannot make ARGB data when both levels and lut have ndim > 2')
+            levels = np.array([-s, s-1])
-            if levels.shape != (data.shape[-1], 2):
-                raise Exception('levels must have shape (data.shape[-1], 2)')
         else:
-            print(levels)
+            raise Exception('levels argument is required for float input types')
-            raise Exception("levels argument must be 1D or 2D.")
+    if not isinstance(levels, np.ndarray):
+        levels = np.array(levels)
+    if levels.ndim == 1:
+        if levels.shape[0] != 2:
+            raise Exception('levels argument must have length 2')
+    elif levels.ndim == 2:
+        if lut is not None and lut.ndim > 1:
+            raise Exception('Cannot make ARGB data when both levels and lut have ndim > 2')
+        if levels.shape != (data.shape[-1], 2):
+            raise Exception('levels must have shape (data.shape[-1], 2)')
+    else:
+        raise Exception("levels argument must be 1D or 2D.")
 
     profile()
 
@@ -935,11 +947,10 @@ def makeARGB(data, lut=None, levels=None, scale=None, useRGBA=False):
     else:
         dtype = np.min_scalar_type(lut.shape[0]-1)
             
-    ## Apply levels if given
+    # Apply levels if given
     if levels is not None:
-        
         if isinstance(levels, np.ndarray) and levels.ndim == 2:
-            ## we are going to rescale each channel independently
+            # we are going to rescale each channel independently
             if levels.shape[0] != data.shape[-1]:
                 raise Exception("When rescaling multi-channel data, there must be the same number of levels as channels (data.shape[-1] == levels.shape[0])")
             newData = np.empty(data.shape, dtype=int)
@@ -950,14 +961,17 @@ def makeARGB(data, lut=None, levels=None, scale=None, useRGBA=False):
                 newData[...,i] = rescaleData(data[...,i], scale/(maxVal-minVal), minVal, dtype=dtype)
             data = newData
         else:
+            # Apply level scaling unless it would have no effect on the data
             minVal, maxVal = levels
-            if minVal == maxVal:
+            if minVal != 0 or maxVal != scale:
-                maxVal += 1e-16
+                if minVal == maxVal:
-            data = rescaleData(data, scale/(maxVal-minVal), minVal, dtype=dtype)
+                    maxVal += 1e-16
+                data = rescaleData(data, scale/(maxVal-minVal), minVal, dtype=dtype)
+            
 
     profile()
 
-    ## apply LUT if given
+    # apply LUT if given
     if lut is not None:
         data = applyLookupTable(data, lut)
     else:
@@ -966,16 +980,18 @@ def makeARGB(data, lut=None, levels=None, scale=None, useRGBA=False):
 
     profile()
 
-    ## copy data into ARGB ordered array
+    # this will be the final image array
     imgData = np.empty(data.shape[:2]+(4,), dtype=np.ubyte)
 
     profile()
 
+    # decide channel order
     if useRGBA:
-        order = [0,1,2,3] ## array comes out RGBA
+        order = [0,1,2,3] # array comes out RGBA
     else:
-        order = [2,1,0,3] ## for some reason, the colors line up as BGR in the final image.
+        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:
         # This is tempting:
         #   imgData[..., :3] = data[..., np.newaxis]
@@ -991,6 +1007,7 @@ def makeARGB(data, lut=None, levels=None, scale=None, useRGBA=False):
         
     profile()
     
+    # add opaque alpha channel if needed
     if data.ndim == 2 or data.shape[2] == 3:
         alpha = False
         imgData[..., 3] = 255

pyqtgraph/tests/test_functions.py

@@ -132,48 +132,161 @@ def test_rescaleData():
 
 
 def test_makeARGB():
+    # Many parameters to test here:
+    #  * data dtype (ubyte, uint16, float, others)
+    #  * data ndim (2 or 3)
+    #  * levels (None, 1D, or 2D)
+    #  * lut dtype
+    #  * lut size
+    #  * lut ndim (1 or 2)
+    #  * useRGBA argument
+    # Need to check that all input values map to the correct output values, especially
+    # at and beyond the edges of the level range.
+
+    def checkArrays(a, b):
+        # because py.test output is difficult to read for arrays
+        if not np.all(a == b):
+            comp = []
+            for i in range(a.shape[0]):
+                if a.shape[1] > 1:
+                    comp.append('[')
+                for j in range(a.shape[1]):
+                    m = a[i,j] == b[i,j]
+                    comp.append('%d,%d  %s %s  %s%s' % 
+                                (i, j, str(a[i,j]).ljust(15), str(b[i,j]).ljust(15),
+                                 m, ' ********' if not np.all(m) else ''))
+                if a.shape[1] > 1:
+                    comp.append(']')
+            raise Exception("arrays do not match:\n%s" % '\n'.join(comp))
+    
+    def checkImage(img, check, alpha, alphaCheck):
+        assert img.dtype == np.ubyte
+        assert alpha is alphaCheck
+        if alpha is False:
+            checkArrays(img[..., 3], 255)
+        
+        if np.isscalar(check) or check.ndim == 3:
+            checkArrays(img[..., :3], check)
+        elif check.ndim == 2:
+            checkArrays(img[..., :3], check[..., np.newaxis])
+        elif check.ndim == 1:
+            checkArrays(img[..., :3], check[..., np.newaxis, np.newaxis])
+        else:
+            raise Exception('invalid check array ndim')
         
     # uint8 data tests
     
-    im1 = np.array([[1,2,3], [4,5,8]], dtype='ubyte')
+    im1 = np.arange(256).astype('ubyte').reshape(256, 1)
-    im2, alpha = pg.makeARGB(im1, levels=(0, 6))
+    im2, alpha = pg.makeARGB(im1, levels=(0, 255))
-    assert im2.dtype == np.ubyte
+    checkImage(im2, im1, alpha, False)
-    assert alpha == False
-    assert np.all(im2[...,3] == 255)
-    assert np.all(im2[...,:3] == np.array([[42, 85, 127], [170, 212, 255]], dtype=np.ubyte)[...,np.newaxis])
     
-    im3, alpha = pg.makeARGB(im1, levels=(0.0, 6.0))
+    im3, alpha = pg.makeARGB(im1, levels=(0.0, 255.0))
-    assert im3.dtype == np.ubyte
+    checkImage(im3, im1, alpha, False)
-    assert alpha == False
-    assert np.all(im3 == im2)
 
-    im2, alpha = pg.makeARGB(im1, levels=(2, 10))
+    im4, alpha = pg.makeARGB(im1, levels=(255, 0))
-    assert im2.dtype == np.ubyte
+    checkImage(im4, 255-im1, alpha, False)
-    assert alpha == False
-    assert np.all(im2[...,3] == 255)
-    assert np.all(im2[...,:3] == np.array([[0, 0, 31], [63, 95, 191]], dtype=np.ubyte)[...,np.newaxis])
     
-    im2, alpha = pg.makeARGB(im1, levels=(2, 10), scale=1.0)
+    im5, alpha = pg.makeARGB(np.concatenate([im1]*3, axis=1), levels=[(0, 255), (0.0, 255.0), (255, 0)])
-    assert im2.dtype == np.float
+    checkImage(im5, np.concatenate([im1, im1, 255-im1], axis=1), alpha, False)
-    assert alpha == False
-    assert np.all(im2[...,3] == 1.0)
-    assert np.all(im2[...,:3] == np.array([[0, 0, 31], [63, 95, 191]], dtype=np.ubyte)[...,np.newaxis])
     
-    # uint8 input + uint8 LUT
+
-    lut = np.arange(512).astype(np.ubyte)[::2][::-1]
+    im2, alpha = pg.makeARGB(im1, levels=(128,383))
-    im2, alpha = pg.makeARGB(im1, lut=lut, levels=(2, 10))
+    checkImage(im2[:128], 0, alpha, False)
-    assert im2.dtype == np.ubyte
+    checkImage(im2[128:], im1[:128], alpha, False)
-    assert alpha == False
+    
-    assert np.all(im2[...,3] == 255)
+
-    assert np.all(im2[...,:3] == np.array([[0, 0, 31], [63, 95, 191]], dtype=np.ubyte)[...,np.newaxis])
+    # uint8 data + uint8 LUT
+    lut = np.arange(256)[::-1].astype(np.uint8)
+    im2, alpha = pg.makeARGB(im1, lut=lut)
+    checkImage(im2, lut, alpha, False)
+    
+    # lut larger than maxint
+    lut = np.arange(511).astype(np.uint8)
+    im2, alpha = pg.makeARGB(im1, lut=lut)
+    checkImage(im2, lut[::2], alpha, False)
+    
+    # lut smaller than maxint
+    lut = np.arange(128).astype(np.uint8)
+    im2, alpha = pg.makeARGB(im1, lut=lut)
+    checkImage(im2, np.linspace(0, 127, 256).astype('ubyte'), alpha, False)
+
+    # lut + levels
+    lut = np.arange(256)[::-1].astype(np.uint8)
+    im2, alpha = pg.makeARGB(im1, lut=lut, levels=[-128, 384])
+    checkImage(im2, np.linspace(192, 65.5, 256).astype('ubyte'), alpha, False)
+    
+    im2, alpha = pg.makeARGB(im1, lut=lut, levels=[64, 192])
+    checkImage(im2, np.clip(np.linspace(385.5, -126.5, 256), 0, 255).astype('ubyte'), alpha, False)
 
     # uint8 data + uint16 LUT
+    lut = np.arange(4096)[::-1].astype(np.uint16) // 16
+    im2, alpha = pg.makeARGB(im1, lut=lut)
+    checkImage(im2, np.arange(256)[::-1].astype('ubyte'), alpha, False)
 
     # uint8 data + float LUT
+    lut = np.linspace(10., 137., 256)
+    im2, alpha = pg.makeARGB(im1, lut=lut)
+    checkImage(im2, lut.astype('ubyte'), alpha, False)
+
+    # uint8 data + 2D LUT
+    lut = np.zeros((256, 3), dtype='ubyte')
+    lut[:,0] = np.arange(256)
+    lut[:,1] = np.arange(256)[::-1]
+    lut[:,2] = 7
+    im2, alpha = pg.makeARGB(im1, lut=lut)
+    checkImage(im2, lut[:,None,::-1], alpha, False)
+    
+    # check useRGBA
+    im2, alpha = pg.makeARGB(im1, lut=lut, useRGBA=True)
+    checkImage(im2, lut[:,None,:], alpha, False)
+
     
     # uint16 data tests
+    im1 = np.arange(0, 2**16, 256).astype('uint16')[:, None]
+    im2, alpha = pg.makeARGB(im1, levels=(512, 2**16))
+    checkImage(im2, np.clip(np.linspace(-2, 253, 256), 0, 255).astype('ubyte'), alpha, False)
 
-    im1 = np.array([[1,2,3], [4,5,8]], dtype='ubyte')
+    lut = (np.arange(512, 2**16)[::-1] // 256).astype('ubyte')
+    im2, alpha = pg.makeARGB(im1, lut=lut, levels=(512, 2**16-256))
+    checkImage(im2, np.clip(np.linspace(257, 2, 256), 0, 255).astype('ubyte'), alpha, False)
+    
+    
+    # float data tests
+    im1 = np.linspace(1.0, 17.0, 256)[:, None]
+    im2, alpha = pg.makeARGB(im1, levels=(5.0, 13.0))
+    checkImage(im2, np.clip(np.linspace(-128, 383, 256), 0, 255).astype('ubyte'), alpha, False)
+    
+    lut = (np.arange(1280)[::-1] // 10).astype('ubyte')
+    im2, alpha = pg.makeARGB(im1, lut=lut, levels=(1, 17))
+    checkImage(im2, np.linspace(127.5, 0, 256).astype('ubyte'), alpha, False)
+
+
+    # test sanity checks
+    class AssertExc(object):
+        def __init__(self, exc=Exception):
+            self.exc = exc
+        def __enter__(self):
+            return self
+        def __exit__(self, *args):
+            assert args[0] is self.exc, "Should have raised %s (got %s)" % (self.exc, args[0])
+            return True
+    
+    with AssertExc(TypeError):  # invalid image shape
+        pg.makeARGB(np.zeros((2,), dtype='float'))
+    with AssertExc(TypeError):  # invalid image shape
+        pg.makeARGB(np.zeros((2,2,7), dtype='float'))
+    with AssertExc():  # float images require levels arg
+        pg.makeARGB(np.zeros((2,2), dtype='float'))
+    with AssertExc():  # bad levels arg
+        pg.makeARGB(np.zeros((2,2), dtype='float'), levels=[1])
+    with AssertExc():  # bad levels arg
+        pg.makeARGB(np.zeros((2,2), dtype='float'), levels=[1,2,3])
+    with AssertExc():  # can't mix 3-channel levels and LUT
+        pg.makeARGB(np.zeros((2,2)), lut=np.zeros((10,3), dtype='ubyte'), levels=[(0,1)]*3)
+    with AssertExc():  # multichannel levels must have same number of channels as image
+        pg.makeARGB(np.zeros((2,2,3), dtype='float'), levels=[(1,2)]*4)
+    with AssertExc():  # 3d levels not allowed
+        pg.makeARGB(np.zeros((2,2,3), dtype='float'), levels=np.zeros([3, 2, 2]))
 
     
 if __name__ == '__main__':