2424class Colorizer ():
2525 """
2626 Class that holds the data to color pipeline
27- accessible via `.to_rgba (A)` and executed via
28- the `.norm` and `.cmap` attributes.
27+ accessible via `Colorizer.__call__ (A)` and executed via
28+ the `Colorizer .norm` and `Colorizer .cmap` attributes.
2929 """
3030 def __init__ (self , cmap = None , norm = None ):
3131
@@ -94,7 +94,7 @@ def norm(self, norm):
9494 if not in_init :
9595 self .changed ()
9696
97- def to_rgba (self , x , alpha = None , bytes = False , norm = True ):
97+ def __call__ (self , x , alpha = None , bytes = False , norm = True ):
9898 """
9999 Return a normalized RGBA array corresponding to *x*.
100100
@@ -125,56 +125,59 @@ def to_rgba(self, x, alpha=None, bytes=False, norm=True):
125125
126126 """
127127 # First check for special case, image input:
128- # First check for special case, image input:
129- try :
130- if x .ndim == 3 :
131- if x .shape [2 ] == 3 :
132- if alpha is None :
133- alpha = 1
134- if x .dtype == np .uint8 :
135- alpha = np .uint8 (alpha * 255 )
136- m , n = x .shape [:2 ]
137- xx = np .empty (shape = (m , n , 4 ), dtype = x .dtype )
138- xx [:, :, :3 ] = x
139- xx [:, :, 3 ] = alpha
140- elif x .shape [2 ] == 4 :
141- xx = x
142- else :
143- raise ValueError ("Third dimension must be 3 or 4" )
144- if xx .dtype .kind == 'f' :
145- # If any of R, G, B, or A is nan, set to 0
146- if np .any (nans := np .isnan (x )):
147- if x .shape [2 ] == 4 :
148- xx = xx .copy ()
149- xx [np .any (nans , axis = 2 ), :] = 0
150-
151- if norm and (xx .max () > 1 or xx .min () < 0 ):
152- raise ValueError ("Floating point image RGB values "
153- "must be in the 0..1 range." )
154- if bytes :
155- xx = (xx * 255 ).astype (np .uint8 )
156- elif xx .dtype == np .uint8 :
157- if not bytes :
158- xx = xx .astype (np .float32 ) / 255
159- else :
160- raise ValueError ("Image RGB array must be uint8 or "
161- "floating point; found %s" % xx .dtype )
162- # Account for any masked entries in the original array
163- # If any of R, G, B, or A are masked for an entry, we set alpha to 0
164- if np .ma .is_masked (x ):
165- xx [np .any (np .ma .getmaskarray (x ), axis = 2 ), 3 ] = 0
166- return xx
167- except AttributeError :
168- # e.g., x is not an ndarray; so try mapping it
169- pass
170-
171- # This is the normal case, mapping a scalar array:
128+ if isinstance (x , np .ndarray ) and x .ndim == 3 :
129+ return self ._pass_image_data (x , alpha , bytes , norm )
130+
131+ # Otherwise run norm -> colormap pipeline
172132 x = ma .asarray (x )
173133 if norm :
174134 x = self .norm (x )
175135 rgba = self .cmap (x , alpha = alpha , bytes = bytes )
176136 return rgba
177137
138+ @staticmethod
139+ def _pass_image_data (x , alpha = None , bytes = False , norm = True ):
140+ """
141+ Helper function to pass ndarray of shape (...,3) or (..., 4)
142+ through `__call__()`, see `__call__()` for docstring.
143+ """
144+ if x .shape [2 ] == 3 :
145+ if alpha is None :
146+ alpha = 1
147+ if x .dtype == np .uint8 :
148+ alpha = np .uint8 (alpha * 255 )
149+ m , n = x .shape [:2 ]
150+ xx = np .empty (shape = (m , n , 4 ), dtype = x .dtype )
151+ xx [:, :, :3 ] = x
152+ xx [:, :, 3 ] = alpha
153+ elif x .shape [2 ] == 4 :
154+ xx = x
155+ else :
156+ raise ValueError ("Third dimension must be 3 or 4" )
157+ if xx .dtype .kind == 'f' :
158+ # If any of R, G, B, or A is nan, set to 0
159+ if np .any (nans := np .isnan (x )):
160+ if x .shape [2 ] == 4 :
161+ xx = xx .copy ()
162+ xx [np .any (nans , axis = 2 ), :] = 0
163+
164+ if norm and (xx .max () > 1 or xx .min () < 0 ):
165+ raise ValueError ("Floating point image RGB values "
166+ "must be in the 0..1 range." )
167+ if bytes :
168+ xx = (xx * 255 ).astype (np .uint8 )
169+ elif xx .dtype == np .uint8 :
170+ if not bytes :
171+ xx = xx .astype (np .float32 ) / 255
172+ else :
173+ raise ValueError ("Image RGB array must be uint8 or "
174+ "floating point; found %s" % xx .dtype )
175+ # Account for any masked entries in the original array
176+ # If any of R, G, B, or A are masked for an entry, we set alpha to 0
177+ if np .ma .is_masked (x ):
178+ xx [np .any (np .ma .getmaskarray (x ), axis = 2 ), 3 ] = 0
179+ return xx
180+
178181 def normalize (self , x ):
179182 """
180183 Normalize the data in x.
@@ -350,7 +353,7 @@ def to_rgba(self, x, alpha=None, bytes=False, norm=True):
350353 performed, and it is assumed to be in the range (0-1).
351354
352355 """
353- return self .colorizer . to_rgba (x , alpha = alpha , bytes = bytes , norm = norm )
356+ return self .colorizer (x , alpha = alpha , bytes = bytes , norm = norm )
354357
355358 def get_clim (self ):
356359 """
0 commit comments