-
Notifications
You must be signed in to change notification settings - Fork 64
Expand file tree
/
Copy pathimage.py
More file actions
553 lines (439 loc) · 16.2 KB
/
image.py
File metadata and controls
553 lines (439 loc) · 16.2 KB
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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
import math
from typing import *
import numpy as np
import pygfx
from ..utils import quick_min_max
from ._base import Graphic
from .selectors import (
LinearSelector,
LinearRegionSelector,
RectangleSelector,
PolygonSelector,
)
from .features import (
TextureArray,
ImageCmap,
ImageVmin,
ImageVmax,
ImageInterpolation,
ImageCmapInterpolation,
)
def _format_value(value: float):
"""float -> rounded str, or str with scientific notation"""
abs_val = abs(value)
if abs_val < 0.01 or abs_val > 9_999:
return f"{value:.2e}"
else:
return f"{value:.4f}"
class _ImageTile(pygfx.Image):
"""
Similar to pygfx.Image, only difference is that it modifies the pick_info
by adding the data row start indices that correspond to this chunk of the big image
"""
def __init__(
self,
geometry,
material,
data_slice: tuple[slice, slice],
chunk_index: tuple[int, int],
**kwargs,
):
super().__init__(geometry, material, **kwargs)
self._data_slice = data_slice
self._chunk_index = chunk_index
def _wgpu_get_pick_info(self, pick_value):
pick_info = super()._wgpu_get_pick_info(pick_value)
data_row_start, data_col_start = (
self.data_slice[0].start,
self.data_slice[1].start,
)
# add the actual data row and col start indices
x, y = pick_info["index"]
x += data_col_start
y += data_row_start
pick_info["index"] = (x, y)
xp, yp = pick_info["pixel_coord"]
xp += data_col_start
yp += data_row_start
pick_info["pixel_coord"] = (xp, yp)
# add row chunk and col chunk index to pick_info dict
return {
**pick_info,
"data_slice": self.data_slice,
"chunk_index": self.chunk_index,
}
@property
def data_slice(self) -> tuple[slice, slice]:
return self._data_slice
@property
def chunk_index(self) -> tuple[int, int]:
return self._chunk_index
class ImageGraphic(Graphic):
_features = {
"data": TextureArray,
"cmap": ImageCmap,
"vmin": ImageVmin,
"vmax": ImageVmax,
"interpolation": ImageInterpolation,
"cmap_interpolation": ImageCmapInterpolation,
}
def __init__(
self,
data: Any,
vmin: float = None,
vmax: float = None,
cmap: str = "plasma",
interpolation: str = "nearest",
cmap_interpolation: str = "linear",
**kwargs,
):
"""
Create an Image Graphic
Parameters
----------
data: array-like
array-like, usually numpy.ndarray, must support ``memoryview()``
| shape must be ``[n_rows, n_cols]``, ``[n_rows, n_cols, 3]`` for RGB or ``[n_rows, n_cols, 4]`` for RGBA
vmin: float, optional
minimum value for color scaling, estimated from data if not provided
vmax: float, optional
maximum value for color scaling, estimated from data if not provided
cmap: str, optional, default "plasma"
colormap to use to display the data. For supported colormaps see the
``cmap`` library catalogue: https://cmap-docs.readthedocs.io/en/stable/catalog/
interpolation: str, optional, default "nearest"
interpolation filter, one of "nearest" or "linear"
cmap_interpolation: str, optional, default "linear"
colormap interpolation method, one of "nearest" or "linear"
kwargs:
additional keyword arguments passed to :class:`.Graphic`
"""
super().__init__(**kwargs)
group = pygfx.Group()
if isinstance(data, TextureArray):
# share buffer
self._data = data
else:
# create new texture array to manage buffer
# texture array that manages the multiple textures on the GPU that represent this image
self._data = TextureArray(data)
if (vmin is None) or (vmax is None):
_vmin, _vmax = quick_min_max(self.data.value)
if vmin is None:
vmin = _vmin
if vmax is None:
vmax = _vmax
# other graphic features
self._vmin = ImageVmin(vmin)
self._vmax = ImageVmax(vmax)
self._interpolation = ImageInterpolation(interpolation)
self._cmap_interpolation = ImageCmapInterpolation(cmap_interpolation)
# set map to None for RGB images
if self._data.value.ndim == 3:
self._cmap = None
_map = None
elif self._data.value.ndim == 2:
# use TextureMap for grayscale images
self._cmap = ImageCmap(cmap)
_map = pygfx.TextureMap(
self._cmap.texture,
filter=self._cmap_interpolation.value,
wrap="clamp-to-edge",
)
else:
raise ValueError(
f"ImageGraphic `data` must have 2 dimensions for grayscale images, or 3 dimensions for RGB(A) images.\n"
f"You have passed a a data array with: {self._data.value.ndim} dimensions, "
f"and of shape: {self._data.value.shape}"
)
# one common material is used for every Texture chunk
self._material = pygfx.ImageBasicMaterial(
clim=(vmin, vmax),
map=_map,
interpolation=self._interpolation.value,
pick_write=True,
)
# create the _ImageTile world objects, add to group
for tile in self._create_tiles():
group.add(tile)
self._set_world_object(group)
def _create_tiles(self) -> list[_ImageTile]:
tiles = list()
# iterate through each texture chunk and create
# an _ImageTile, offset the tile using the data indices
for texture, chunk_index, data_slice in self._data:
# create an ImageTile using the texture for this chunk
img = _ImageTile(
geometry=pygfx.Geometry(grid=texture),
material=self._material,
data_slice=data_slice, # used to parse pick_info
chunk_index=chunk_index,
)
# row and column start index for this chunk
data_row_start = data_slice[0].start
data_col_start = data_slice[1].start
# offset tile position using the indices from the big data array
# that correspond to this chunk
img.world.x = data_col_start
img.world.y = data_row_start
tiles.append(img)
return tiles
@property
def data(self) -> TextureArray:
"""
Get or set the image data.
Note that if the shape of the new data array does not equal the shape of
current data array, a new set of GPU Textures are automatically created.
This can have performance drawbacks when you have a ver large images.
This is usually fine as long as you don't need to do it hundreds of times
per second.
"""
return self._data
@data.setter
def data(self, data):
if isinstance(data, np.ndarray):
# check if a new buffer is required
if self._data.value.shape != data.shape:
# create new TextureArray
self._data = TextureArray(data)
# cmap based on if rgb or grayscale
if self._data.value.ndim > 2:
self._cmap = None
# must be None if RGB(A)
self._material.map = None
else:
if self.cmap is None: # have switched from RGBA -> grayscale image
# create default cmap
self._cmap = ImageCmap("plasma")
self._material.map = pygfx.TextureMap(
self._cmap.texture,
filter=self._cmap_interpolation.value,
wrap="clamp-to-edge",
)
# remove tiles from the WorldObject -> Graphic map
self._remove_group_graphic_map(self.world_object)
# clear image tiles
self.world_object.clear()
# create new tiles
for tile in self._create_tiles():
self.world_object.add(tile)
# add new tiles to WorldObject -> Graphic map
self._add_group_graphic_map(self.world_object)
return
self._data[:] = data
@property
def cmap(self) -> str | None:
"""
Get or set the colormap for grayscale images. Returns ``None`` if image is RGB(A).
For supported colormaps see the ``cmap`` library catalogue: https://cmap-docs.readthedocs.io/en/stable/catalog/
"""
if self._cmap is not None:
return self._cmap.value
@cmap.setter
def cmap(self, name: str):
if self.data.value.ndim > 2:
raise AttributeError("RGB(A) images do not have a colormap property")
self._cmap.set_value(self, name)
@property
def vmin(self) -> float:
"""lower contrast limit"""
return self._vmin.value
@vmin.setter
def vmin(self, value: float):
self._vmin.set_value(self, value)
@property
def vmax(self) -> float:
"""upper contrast limit"""
return self._vmax.value
@vmax.setter
def vmax(self, value: float):
self._vmax.set_value(self, value)
@property
def interpolation(self) -> str:
"""Data interpolation method"""
return self._interpolation.value
@interpolation.setter
def interpolation(self, value: str):
self._interpolation.set_value(self, value)
@property
def cmap_interpolation(self) -> str:
"""cmap interpolation method, 'linear' or 'nearest'. Used only for grayscale images"""
return self._cmap_interpolation.value
@cmap_interpolation.setter
def cmap_interpolation(self, value: str):
self._cmap_interpolation.set_value(self, value)
def reset_vmin_vmax(self):
"""
Reset the vmin, vmax by estimating it from the data by subsampling.
"""
vmin, vmax = quick_min_max(self._data.value)
self.vmin = vmin
self.vmax = vmax
def add_linear_selector(
self, selection: int = None, axis: str = "x", **kwargs
) -> LinearSelector:
"""
Adds a :class:`.LinearSelector`.
Selectors are just ``Graphic`` objects, so you can manage, remove, or delete them
from a plot area just like any other ``Graphic``.
Parameters
----------
selection: int, optional
initial position of the selector
kwargs:
passed to :class:`.LinearSelector`
Returns
-------
LinearSelector
"""
if axis == "x":
limits = (0, self._data.value.shape[1])
elif axis == "y":
limits = (0, self._data.value.shape[0])
else:
raise ValueError("`axis` must be one of 'x' | 'y'")
if selection is None:
selection = limits[0]
if selection < limits[0] or selection > limits[1]:
raise ValueError(
f"the passed selection: {selection} is beyond the limits: {limits}"
)
selector = LinearSelector(
selection=selection,
limits=limits,
axis=axis,
parent=self,
**kwargs,
)
self._plot_area.add_graphic(selector, center=False)
return selector
def add_linear_region_selector(
self,
selection: tuple[float, float] = None,
axis: str = "x",
padding: float = 0.0,
fill_color=(0, 0, 0.35, 0.2),
**kwargs,
) -> LinearRegionSelector:
"""
Add a :class:`.LinearRegionSelector`.
Selectors are just ``Graphic`` objects, so you can manage, remove, or delete them
from a plot area just like any other ``Graphic``.
Parameters
----------
selection: (float, float)
initial (min, max) of the selection
axis: "x" | "y"
axis the selector can move along
padding: float, default 100.0
Extends the linear selector along the perpendicular axis to make it easier to interact with.
kwargs
passed to ``LinearRegionSelector``
Returns
-------
LinearRegionSelector
"""
if axis == "x":
size = self._data.value.shape[0]
center = size / 2
limits = (0, self._data.value.shape[1])
elif axis == "y":
size = self._data.value.shape[1]
center = size / 2
limits = (0, self._data.value.shape[0])
else:
raise ValueError("`axis` must be one of 'x' | 'y'")
# default padding is 25% the height or width of the image
if padding is None:
size *= 1.25
else:
size += padding
if selection is None:
selection = limits[0], int(limits[1] * 0.25)
if padding is None:
size *= 1.25
else:
size += padding
selector = LinearRegionSelector(
selection=selection,
limits=limits,
size=size,
center=center,
axis=axis,
fill_color=fill_color,
parent=self,
**kwargs,
)
self._plot_area.add_graphic(selector, center=False)
return selector
def add_rectangle_selector(
self,
selection: tuple[float, float, float, float] = None,
fill_color=(0, 0, 0.35, 0.2),
**kwargs,
) -> RectangleSelector:
"""
Add a :class:`.RectangleSelector`.
Selectors are just ``Graphic`` objects, so you can manage, remove, or delete them
from a plot area just like any other ``Graphic``.
Parameters
----------
selection: (float, float, float, float), optional
initial (xmin, xmax, ymin, ymax) of the selection
"""
# default selection is 25% of the diagonal
if selection is None:
diagonal = math.sqrt(
self._data.value.shape[0] ** 2 + self._data.value.shape[1] ** 2
)
selection = (0, int(diagonal / 4), 0, int(diagonal / 4))
# min/max limits are image shape
# rows are ys, columns are xs
limits = (0, self._data.value.shape[1], 0, self._data.value.shape[0])
selector = RectangleSelector(
selection=selection,
limits=limits,
fill_color=fill_color,
parent=self,
**kwargs,
)
self._plot_area.add_graphic(selector, center=False)
return selector
def add_polygon_selector(
self,
selection: List[tuple[float, float]] = None,
fill_color=(0, 0, 0.35, 0.2),
**kwargs,
) -> PolygonSelector:
"""
Add a :class:`.PolygonSelector`.
Selectors are just ``Graphic`` objects, so you can manage, remove, or delete them
from a plot area just like any other ``Graphic``.
Parameters
----------
selection: list[tuple[float, float]], optional
Initial points for the polygon. If not given or None, you'll start drawing the selection (clicking adds points to the polygon).
"""
# min/max limits are image shape
# rows are ys, columns are xs
limits = (0, self._data.value.shape[1], 0, self._data.value.shape[0])
selector = PolygonSelector(
selection,
limits,
fill_color=fill_color,
parent=self,
**kwargs,
)
self._plot_area.add_graphic(selector, center=False)
return selector
def format_pick_info(self, pick_info: dict) -> str:
col, row = pick_info["index"]
if self.data.value.ndim == 2:
val = self.data[row, col]
info = f"{val:.4g}"
else:
info = "\n".join(
f"{channel}: {val:.4g}"
for channel, val in zip("rgba", self.data[row, col])
)
return info