<s> from __future__ import absolute_import , division , print_function <EOL> import warnings <EOL> from collections import Iterator <EOL> from io import BytesIO <EOL> import numpy as np <EOL> import numba as nb <EOL> import toolz as tz <EOL> import xarray as xr <EOL> from PIL . Image import fromarray <EOL> from . colors import rgb <EOL> from . composite import composite_op_lookup , over <EOL> from . utils import ngjit <EOL> __all__ = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' ] <EOL> class Image ( xr . DataArray ) : <EOL> __array_priority__ = <NUM_LIT> <EOL> def _repr_png_ ( self ) : <EOL> return self . to_pil ( ) . _repr_png_ ( ) <EOL> def to_pil ( self , origin = '<STR_LIT>' ) : <EOL> arr = np . flipud ( self . data ) if origin == '<STR_LIT>' else self . data <EOL> return fromarray ( arr , '<STR_LIT>' ) <EOL> def to_bytesio ( self , format = '<STR_LIT>' , origin = '<STR_LIT>' ) : <EOL> fp = BytesIO ( ) <EOL> self . to_pil ( origin ) . save ( fp , format ) <EOL> fp . seek ( <NUM_LIT:0> ) <EOL> return fp <EOL> def stack ( * imgs , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if not imgs : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> for i in imgs : <EOL> if not isinstance ( i , Image ) : <EOL> raise TypeError ( "<STR_LIT>" . format ( type ( i ) ) ) <EOL> op = composite_op_lookup [ kwargs . get ( '<STR_LIT>' , '<STR_LIT>' ) ] <EOL> if len ( imgs ) == <NUM_LIT:1> : <EOL> return imgs [ <NUM_LIT:0> ] <EOL> imgs = xr . align ( * imgs , copy = False , join = '<STR_LIT>' ) <EOL> out = tz . reduce ( tz . flip ( op ) , [ i . data for i in imgs ] ) <EOL> return Image ( out , coords = imgs [ <NUM_LIT:0> ] . coords , dims = imgs [ <NUM_LIT:0> ] . dims ) <EOL> def eq_hist ( data , mask = None , nbins = <NUM_LIT> * <NUM_LIT> ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( data , np . ndarray ) : <EOL> raise TypeError ( "<STR_LIT>" ) <EOL> data2 = data if mask is None else data [ ~ mask ] <EOL> if np . issubdtype ( data2 . dtype , np . integer ) : <EOL> hist = np . bincount ( data2 . ravel ( ) ) <EOL> bin_centers = np . arange ( len ( hist ) ) <EOL> idx = np . nonzero ( hist ) [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] <EOL> hist , bin_centers = hist [ idx : ] , bin_centers [ idx : ] <EOL> else : <EOL> hist , bin_edges = np . histogram ( data2 , bins = nbins ) <EOL> bin_centers = ( bin_edges [ : - <NUM_LIT:1> ] + bin_edges [ <NUM_LIT:1> : ] ) / <NUM_LIT:2> <EOL> cdf = hist . cumsum ( ) <EOL> cdf = cdf / float ( cdf [ - <NUM_LIT:1> ] ) <EOL> out = np . interp ( data . flat , bin_centers , cdf ) . reshape ( data . shape ) <EOL> return out if mask is None else np . where ( mask , np . nan , out ) <EOL> _interpolate_lookup = { '<STR_LIT>' : lambda d , m : np . log1p ( np . where ( m , np . nan , d ) ) , <EOL> '<STR_LIT>' : lambda d , m : np . where ( m , np . nan , d ) ** ( <NUM_LIT:1> / <NUM_LIT> ) , <EOL> '<STR_LIT>' : lambda d , m : np . where ( m , np . nan , d ) , <EOL> '<STR_LIT>' : eq_hist } <EOL> def _normalize_interpolate_how ( how ) : <EOL> if callable ( how ) : <EOL> return how <EOL> elif how in _interpolate_lookup : <EOL> return _interpolate_lookup [ how ] <EOL> raise ValueError ( "<STR_LIT>" . format ( how ) ) <EOL> def interpolate ( agg , low = None , high = None , cmap = None , how = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( agg , xr . DataArray ) : <EOL> raise TypeError ( "<STR_LIT>" ) <EOL> if agg . ndim != <NUM_LIT:2> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if cmap is not None : <EOL> if low or high : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> else : <EOL> cmap = [ '<STR_LIT>' , '<STR_LIT>' ] <EOL> if low or high : <EOL> with warnings . catch_warnings ( ) : <EOL> warnings . simplefilter ( '<STR_LIT>' , DeprecationWarning ) <EOL> w = DeprecationWarning ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> warnings . warn ( w ) <EOL> cmap = [ low or cmap [ <NUM_LIT:0> ] , high or cmap [ <NUM_LIT:1> ] ] <EOL> how = _normalize_interpolate_how ( how ) <EOL> data = agg . data <EOL> if np . issubdtype ( data . dtype , np . bool_ ) : <EOL> mask = ~ data <EOL> interp = data <EOL> else : <EOL> if np . issubdtype ( data . dtype , np . integer ) : <EOL> mask = data == <NUM_LIT:0> <EOL> else : <EOL> mask = np . isnan ( data ) <EOL> offset = data [ ~ mask ] . min ( ) <EOL> interp = data - offset <EOL> data = how ( interp , mask ) <EOL> span = [ np . nanmin ( data ) , np . nanmax ( data ) ] <EOL> if isinstance ( cmap , Iterator ) : <EOL> cmap = list ( cmap ) <EOL> if isinstance ( cmap , list ) : <EOL> rspan , gspan , bspan = np . array ( list ( zip ( * map ( rgb , cmap ) ) ) ) <EOL> span = np . linspace ( span [ <NUM_LIT:0> ] , span [ <NUM_LIT:1> ] , len ( cmap ) ) <EOL> r = np . interp ( data , span , rspan , left = <NUM_LIT:255> ) . astype ( np . uint8 ) <EOL> g = np . interp ( data , span , gspan , left = <NUM_LIT:255> ) . astype ( np . uint8 ) <EOL> b = np . interp ( data , span , bspan , left = <NUM_LIT:255> ) . astype ( np . uint8 ) <EOL> a = np . where ( np . isnan ( data ) , <NUM_LIT:0> , <NUM_LIT:255> ) . astype ( np . uint8 ) <EOL> img = np . dstack ( [ r , g , b , a ] ) . view ( np . uint32 ) . reshape ( a . shape ) <EOL> elif callable ( cmap ) : <EOL> img = cmap ( ( data - span [ <NUM_LIT:0> ] ) / ( span [ <NUM_LIT:1> ] - span [ <NUM_LIT:0> ] ) , bytes = True ) <EOL> img [ : , : , <NUM_LIT:3> ] = np . where ( np . isnan ( data ) , <NUM_LIT:0> , <NUM_LIT:255> ) . astype ( np . uint8 ) <EOL> img = img . view ( np . uint32 ) . reshape ( data . shape ) <EOL> else : <EOL> raise TypeError ( "<STR_LIT>" <EOL> "<STR_LIT>" . format ( type ( cmap ) ) ) <EOL> return Image ( img , coords = agg . coords , dims = agg . dims ) <EOL> def colorize ( agg , color_key , how = '<STR_LIT>' , min_alpha = <NUM_LIT:20> ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( agg , xr . DataArray ) : <EOL> raise TypeError ( "<STR_LIT>" ) <EOL> if not agg . ndim == <NUM_LIT:3> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> cats = agg . indexes [ agg . dims [ - <NUM_LIT:1> ] ] <EOL> if not isinstance ( color_key , dict ) : <EOL> color_key = dict ( zip ( cats , color_key ) ) <EOL> if len ( color_key ) != len ( cats ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not ( <NUM_LIT:0> <= min_alpha <= <NUM_LIT:255> ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> colors = [ rgb ( color_key [ c ] ) for c in cats ] <EOL> rs , gs , bs = map ( np . array , zip ( * colors ) ) <EOL> data = agg . data <EOL> total = data . sum ( axis = <NUM_LIT:2> ) <EOL> r = ( data . dot ( rs ) / total ) . astype ( np . uint8 ) <EOL> g = ( data . dot ( gs ) / total ) . astype ( np . uint8 ) <EOL> b = ( data . dot ( bs ) / total ) . astype ( np . uint8 ) <EOL> offset = total . min ( ) <EOL> mask = np . isnan ( total ) <EOL> if offset == <NUM_LIT:0> : <EOL> mask = mask | ( total <= <NUM_LIT:0> ) <EOL> offset = total [ total > <NUM_LIT:0> ] . min ( ) <EOL> a = _normalize_interpolate_how ( how ) ( total - offset , mask ) <EOL> a = np . interp ( a , [ np . nanmin ( a ) , np . nanmax ( a ) ] , <EOL> [ min_alpha , <NUM_LIT:255> ] , left = <NUM_LIT:0> , right = <NUM_LIT:255> ) . astype ( np . uint8 ) <EOL> r [ mask ] = g [ mask ] = b [ mask ] = <NUM_LIT:255> <EOL> return Image ( np . dstack ( [ r , g , b , a ] ) . view ( np . uint32 ) . reshape ( a . shape ) , <EOL> dims = agg . dims [ : - <NUM_LIT:1> ] , coords = list ( agg . coords . values ( ) ) [ : - <NUM_LIT:1> ] ) <EOL> def set_background ( img , color = None ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( img , Image ) : <EOL> raise TypeError ( "<STR_LIT>" . format ( type ( img ) ) ) <EOL> if color is None : <EOL> return img <EOL> background = np . uint8 ( rgb ( color ) + ( <NUM_LIT:255> , ) ) . view ( '<STR_LIT>' ) [ <NUM_LIT:0> ] <EOL> data = over ( img . data , background ) <EOL> return Image ( data , coords = img . coords , dims = img . dims ) <EOL> def spread ( img , px = <NUM_LIT:1> , shape = '<STR_LIT>' , how = '<STR_LIT>' , mask = None ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( img , Image ) : <EOL> raise TypeError ( "<STR_LIT>" . format ( type ( img ) ) ) <EOL> if mask is None : <EOL> if not isinstance ( px , int ) or px < <NUM_LIT:0> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if px == <NUM_LIT:0> : <EOL> return img <EOL> mask = _mask_lookup [ shape ] ( px ) <EOL> elif not ( isinstance ( mask , np . ndarray ) and mask . ndim == <NUM_LIT:2> and <EOL> mask . shape [ <NUM_LIT:0> ] == mask . shape [ <NUM_LIT:1> ] and mask . shape [ <NUM_LIT:0> ] % <NUM_LIT:2> == <NUM_LIT:1> ) : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> mask = mask if mask . dtype == '<STR_LIT:bool>' else mask . astype ( '<STR_LIT:bool>' ) <EOL> kernel = _build_spread_kernel ( how ) <EOL> w = mask . shape [ <NUM_LIT:0> ] <EOL> extra = w // <NUM_LIT:2> <EOL> M , N = img . shape <EOL> buf = np . zeros ( ( M + <NUM_LIT:2> * extra , N + <NUM_LIT:2> * extra ) , dtype = '<STR_LIT>' ) <EOL> kernel ( img . data , mask , buf ) <EOL> out = buf [ extra : - extra , extra : - extra ] . copy ( ) <EOL> return Image ( out , dims = img . dims , coords = img . coords ) <EOL> @ tz . memoize <EOL> def _build_spread_kernel ( how ) : <EOL> """<STR_LIT>""" <EOL> op = composite_op_lookup [ how ] <EOL> @ ngjit <EOL> def kernel ( arr , mask , out ) : <EOL> M , N = arr . shape <EOL> w = mask . shape [ <NUM_LIT:0> ] <EOL> for y in range ( M ) : <EOL> for x in range ( N ) : <EOL> el = arr [ y , x ] <EOL> if ( el >> <NUM_LIT> ) & <NUM_LIT:255> : <EOL> for i in range ( w ) : <EOL> for j in range ( w ) : <EOL> if mask [ i , j ] : <EOL> out [ i + y , j + x ] = op ( el , out [ i + y , j + x ] ) <EOL> return kernel <EOL> def _square_mask ( px ) : <EOL> """<STR_LIT>""" <EOL> px = int ( px ) <EOL> w = <NUM_LIT:2> * px + <NUM_LIT:1> <EOL> return np . ones ( ( w , w ) , dtype = '<STR_LIT:bool>' ) <EOL> def _circle_mask ( r ) : <EOL> """<STR_LIT>""" <EOL> x = np . arange ( - r , r + <NUM_LIT:1> , dtype = '<STR_LIT>' ) <EOL> return np . where ( np . sqrt ( x ** <NUM_LIT:2> + x [ : , None ] ** <NUM_LIT:2> ) <= r + <NUM_LIT:0.5> , True , False ) <EOL> _mask_lookup = { '<STR_LIT>' : _square_mask , <EOL> '<STR_LIT>' : _circle_mask } <EOL> def dynspread ( img , threshold = <NUM_LIT:0.5> , max_px = <NUM_LIT:3> , shape = '<STR_LIT>' , how = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> if not <NUM_LIT:0> <= threshold <= <NUM_LIT:1> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not isinstance ( max_px , int ) or max_px < <NUM_LIT:0> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> for px in range ( max_px + <NUM_LIT:1> ) : <EOL> out = spread ( img , px , shape = shape , how = how ) <EOL> if _density ( out . data ) >= threshold : <EOL> break <EOL> return out <EOL> @ nb . jit ( nopython = True , nogil = True , cache = True ) <EOL> def _density ( arr ) : <EOL> """<STR_LIT>""" <EOL> M , N = arr . shape <EOL> cnt = total = <NUM_LIT:0> <EOL> for y in range ( <NUM_LIT:1> , M - <NUM_LIT:1> ) : <EOL> for x in range ( <NUM_LIT:1> , N - <NUM_LIT:1> ) : <EOL> if ( arr [ y , x ] >> <NUM_LIT> ) & <NUM_LIT:255> : <EOL> cnt += <NUM_LIT:1> <EOL> for i in range ( y - <NUM_LIT:1> , y + <NUM_LIT:2> ) : <EOL> for j in range ( x - <NUM_LIT:1> , x + <NUM_LIT:2> ) : <EOL> if ( arr [ i , j ] >> <NUM_LIT> ) & <NUM_LIT:255> : <EOL> total += <NUM_LIT:1> <EOL> return ( total - cnt ) / ( cnt * <NUM_LIT:8> ) if cnt else np . inf </s>
