Spaces:
Running on Zero
Running on Zero
File size: 2,344 Bytes
bac69fc | 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 | import matplotlib.cm as cm
import numpy as np
from PIL import Image as PImage, ImageDraw as PImageDraw
from scipy.interpolate import RBFInterpolator
def scale_2d_array(array, size, sampling=PImage.Resampling.BILINEAR):
if len(array.shape) == 1:
dim = int(array.shape[0] ** 0.5)
array = array.reshape(dim, dim)
return np.array(PImage.fromarray(array).resize(size, resample=sampling))
def scale_2d_array_rbf(array, size, kernel="thin_plate_spline"):
x0 = np.linspace(0, size[0], array.shape[0])
y0 = np.linspace(0, size[1], array.shape[0])
xy0 = np.array([[x,y] for y in y0 for x in x0])
z0 = array.reshape(-1)
rbf = RBFInterpolator(xy0, z0, kernel=kernel)
X1 = np.arange(0, size[0])
Y1 = np.arange(0, size[1])
XY1 = np.asarray(np.meshgrid(X1, Y1, indexing="xy"))
XY1_flat = XY1.reshape(2, -1).T
Z1_flat = rbf(XY1_flat)
return Z1_flat.reshape(size[1], size[0])
def mask_image(img, mask, sampling=PImage.Resampling.BILINEAR):
img_np = np.array(img)
if len(mask.shape) < 2 or mask.shape[0] != img_np.shape[0] or mask.shape[1] != img_np.shape[1]:
mask = scale_2d_array(mask, img.size, sampling=sampling)
if len(mask.shape) == 2:
mask = mask[:, :, None]
return PImage.fromarray((mask * img_np).astype(np.uint8))
# map := [ 'viridis', 'plasma', 'inferno', 'magma' ]
def heatmap_image(data, *, size=None, cmap="inferno", sampling=PImage.Resampling.BILINEAR):
if size:
data = scale_2d_array(data, size, sampling=sampling)
map_fun_np = np.vectorize(cm.get_cmap(cmap))
rgba_np = 255 * np.stack(map_fun_np(data)[:3], axis=-1)
himg = PImage.fromarray(rgba_np.astype(np.uint8))
if size:
himg = himg.resize(size)
return himg
def heatmap_image_rbf(data, *, size=None, cmap="inferno", kernel="thin_plate_spline"):
if size:
data = scale_2d_array_rbf(data, size, kernel=kernel)
map_fun_np = np.vectorize(cm.get_cmap(cmap))
rgba_np = 255 * np.stack(map_fun_np(data)[:3], axis=-1)
himg = PImage.fromarray(rgba_np.astype(np.uint8))
if size:
himg = himg.resize(size)
return himg
def draw_results(img, objs):
img = img.convert("RGB").copy()
draw = PImageDraw.Draw(img)
for o in objs:
draw.rectangle(tuple(o["box"].values()),
outline=(10, 220, 10),
width=(min(img.size) // 128))
img.thumbnail((512, 512))
return img
|