Spaces:
Running
Running
| import matplotlib.cm as cm | |
| import numpy as np | |
| from PIL import Image as PImage, ImageDraw as PImageDraw, ImageFont as PImageFont | |
| 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 | |
| def draw_text_boxes(img, box_info, thold=0.15): | |
| mdraw = PImageDraw.Draw(img) | |
| font = PImageFont.load_default(32) | |
| for b,t,s in box_info: | |
| if s < thold: continue | |
| (x0, y0), (x1, y1) = b[0], b[2] | |
| x0,x1 = sorted([x0,x1]) | |
| y0,y1 = sorted([y0,y1]) | |
| mdraw.rectangle((x0,y0,x1,y1), outline=(0,0,0), width=8) | |
| mdraw.rectangle((x0,y0-20,x0+22*len(t),y0+15), fill=(0,0,0,160)) | |
| mdraw.text((x0,y0-24), t, font=font) | |
| return img | |