File size: 3,976 Bytes
d2f661a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | from numba import njit, prange
import numpy as np
from scipy.signal import convolve
def log_scale_with_zero(range, n=65536, dtype=np.float32):
scale = np.linspace(np.log10(range[0]), np.log10(range[1]), n-1)
scale = np.hstack((0, 10**scale)).astype(dtype)
return scale
def log_quantize_with_zero(x, range, n=65536, dtype=np.uint16):
scale = log_scale_with_zero(range, n=n, dtype=x.dtype)
y = np.empty_like(x, dtype=dtype)
log_quant_with_zeros(x, y, np.log10(scale[1:]))
return (y, scale)
# optimized helper function for the above
@njit(parallel=True)
def log_quant_with_zeros(x, y, scale):
x = x.ravel()
y = y.ravel()
min_val = 10**scale[0]
for i in prange(x.shape[0]):
# map small values to 0
if x[i] < min_val:
y[i] = 0
continue
lx = np.log10(x[i])
if lx >= scale[-1]:
# map too big values to max of scale
y[i] = len(scale)
else:
# binary search for the rest
k0 = 0
k1 = len(scale)
while k1-k0 > 1:
km = k0 + (k1-k0)//2
if lx < scale[km]:
k1 = km
else:
k0 = km
if k0 == len(scale)-1:
q = k0
elif k0 == 0:
q = 0
else:
d0 = abs(lx-scale[k0])
d1 = abs(lx-scale[k1])
if d0 < d1:
q = k0
else:
q = k1
y[i] = q+1 # add 1 to leave space for zero
@njit(parallel=True)
def average_pool(x, factor=2, missing=65535):
y = np.empty((x.shape[0]//factor, x.shape[1]//factor), dtype=x.dtype)
N = factor**2
N_thresh = N//2
for iy in prange(y.shape[0]):
ix0 = iy * factor
ix1 = ix0 + factor
for jy in range(y.shape[1]):
jx0 = jy * factor
jx1 = jx0 + factor
v = float(0.0)
num_valid = 0
for ix in range(ix0, ix1):
for jx in range(jx0, jx1):
if x[ix,jx] != missing:
v += x[ix,jx]
num_valid += 1
if num_valid >= N_thresh:
y[iy,jy] = v/num_valid
else:
y[iy,jy] = missing
return y
@njit(parallel=True)
def mode_pool(x, num_values=256, factor=2):
y = np.empty((x.shape[0]//factor, x.shape[1]//factor), dtype=x.dtype)
for iy in prange(y.shape[0]):
v = np.empty(num_values, dtype=np.int64)
ix0 = iy * factor
ix1 = ix0 + factor
for jy in range(y.shape[1]):
jx0 = jy * factor
jx1 = jx0 + factor
v[:] = 0
for ix in range(ix0, ix1):
for jx in range(jx0, jx1):
v[x[ix,jx]] += 1
y[iy,jy] = v.argmax()
return y
def fill_holes(missing=65535, rad=1):
def fill(x):
# identify mask of points to fill
o = np.ones((2*rad+1,2*rad+1), dtype=np.uint16)
n = np.prod(o.shape)
valid = (x != missing)
num_valid_neighbors = convolve(valid, o, mode='same', method='direct')
mask = ~valid & (num_valid_neighbors > 0)
# compute mean of valid points around each fillable point
fx = x.copy()
fx[~valid] = 0
mx = convolve(fx, o.astype(np.float64), mode='same', method='direct')
mx = mx[mask] / num_valid_neighbors[mask]
if np.issubdtype(x.dtype, np.integer):
mx = mx.round().astype(x.dtype)
# fill holes with mean
fx = x.copy()
fx[mask] = mx
return fx
return fill
|