forked from lbrindze/angle_grinder
-
Notifications
You must be signed in to change notification settings - Fork 0
/
image_utils.py
88 lines (71 loc) · 2.39 KB
/
image_utils.py
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
import io
import numpy as np
import xarray as xr
import cv2
from PIL import Image
from typing import Iterable, Callable
import colormap
from helpers import run_async
def resize(data_tile: xr.DataArray) -> np.ndarray:
img = Image.fromarray(data_tile.values).resize((512, 512))
return np.array(img)
def normalize(
data_tile: xr.DataArray, max_val=None, min_val=None
) -> np.ndarray:
if max_val is None:
max_val = data_tile.max().values
if min_val is None:
min_val = data_tile.min().values
diff_from_min = data_tile - max_val
data_tile_range = max_val - min_val
return np.floor(255 * diff_from_min / data_tile_range).astype(np.uint8)
def apply_colormap(color_map: Iterable[int]) -> Callable[[int], int]:
num_colors = len(color_map)
vals_per_tier = 255 // (num_colors - 1)
gradients = np.concatenate(
[
create_color_gradient(color_map[i - 1], color, vals_per_tier)
for i, color in enumerate(color_map)
if i > 0
]
)
def mapper(vals):
canvas = np.ones(512 * 512 * 3)
for i, val in enumerate(vals.reshape(512 * 512)):
try:
color_val = gradients[val]
except IndexError:
color_val = gradients[-1]
idx = 3 * i
stride = idx + 3
canvas[idx:stride] = color_val
return canvas.reshape(512, 512, 3)
return mapper
@run_async
def encode_as_png(vals: np.ndarray, colorize=None) -> io.BytesIO:
if colorize is None or colorize is "":
pixels = vals.astype(np.uint8)
else:
try:
mapping = getattr(colormap, colorize)
pixels = (
apply_colormap(mapping)(vals)
.reshape(512, 512, 3)
.astype(np.uint8)
)
except AttributeError:
print(f"No colormap with name {colorize} exists")
pixels = vals.astype(np.uint8)
is_success, img_bytes = cv2.imencode(".png", pixels)
if not is_success:
raise Exception("Could not encode!")
io_buffer = io.BytesIO(img_bytes)
return io_buffer
def create_color_gradient(from_color, to_color, steps) -> np.ndarray:
c = np.linspace(0, 1, steps)[:, None]
x = np.ones((steps, 3))
x[:, 0:3] = from_color
y = np.ones((steps, 3))
y[:, 0:3] = to_color
gradient = x + (y - x) * c
return gradient.astype(np.uint8)