|
7 | 7 |
|
8 | 8 | import os |
9 | 9 | import os.path as osp |
| 10 | +from collections import OrderedDict |
10 | 11 | from typing import TYPE_CHECKING, Any |
11 | 12 |
|
12 | 13 | import numpy as np |
|
46 | 47 | from plotpy.mathutils.arrayfuncs import get_nan_range |
47 | 48 | from plotpy.mathutils.colormap import FULLRANGE, get_cmap |
48 | 49 | from plotpy.styles.image import RawImageParam |
| 50 | +from plotpy.widgets.colormap.widget import EditableColormap |
49 | 51 |
|
50 | 52 | if TYPE_CHECKING: |
51 | 53 | import guidata.io |
|
57 | 59 | from plotpy.interfaces import IItemType |
58 | 60 | from plotpy.plot import BasePlot |
59 | 61 | from plotpy.styles.base import ItemParameters |
60 | | - from plotpy.widgets.colormap.widget import EditableColormap |
| 62 | + |
| 63 | +#: Lookup tables already computed by :func:`_get_colormap_lut`, most recent last |
| 64 | +_LUT_CACHE: OrderedDict[tuple, tuple[tuple[int, ...], np.ndarray]] = OrderedDict() |
| 65 | +_LUT_CACHE_SIZE = 32 |
| 66 | + |
| 67 | + |
| 68 | +def _lut_cache_key( |
| 69 | + table: EditableColormap, alpha: float, alpha_function: int |
| 70 | +) -> tuple | None: |
| 71 | + """Return the key identifying the lookup table of a colormap, if cacheable. |
| 72 | +
|
| 73 | + The key is built from the colormap content (not its identity), so that an |
| 74 | + edited colormap gets a new lookup table. |
| 75 | +
|
| 76 | + Args: |
| 77 | + table: Colormap |
| 78 | + alpha: Global alpha value |
| 79 | + alpha_function: Alpha function (see :class:`plotpy.constants.LUTAlpha`) |
| 80 | +
|
| 81 | + Returns: |
| 82 | + Cache key, or None if the lookup table must be computed: colormap |
| 83 | + subclasses may compute colors differently, and the colors of translucent |
| 84 | + stops depend on the order in which the stops were inserted. |
| 85 | + """ |
| 86 | + if type(table) is not EditableColormap: |
| 87 | + return None |
| 88 | + stops = tuple((stop.pos, stop.rgb) for stop in table.colorStops()) |
| 89 | + if any(stop.a != 255 for stop in table.colorStops()): |
| 90 | + return None |
| 91 | + return (stops, table.mode(), table.invert, alpha, alpha_function) |
| 92 | + |
| 93 | + |
| 94 | +def _compute_colormap_lut( |
| 95 | + table: EditableColormap, alpha: float, alpha_function: int |
| 96 | +) -> tuple[tuple[int, ...], np.ndarray]: |
| 97 | + """Compute the color table and the lookup table of a colormap. |
| 98 | +
|
| 99 | + Args: |
| 100 | + table: Colormap |
| 101 | + alpha: Global alpha value |
| 102 | + alpha_function: Alpha function (see :class:`plotpy.constants.LUTAlpha`) |
| 103 | +
|
| 104 | + Returns: |
| 105 | + Color table (256 colors) and lookup table (``LUT_SIZE`` ARGB values) |
| 106 | + """ |
| 107 | + colors = tuple(table.colorTable(FULLRANGE)) |
| 108 | + lut = np.zeros((LUT_SIZE,), np.uint32) |
| 109 | + for i in range(LUT_SIZE): |
| 110 | + if alpha_function == LUTAlpha.NONE.value: |
| 111 | + pix_alpha = 1.0 |
| 112 | + elif alpha_function == LUTAlpha.CONSTANT.value: |
| 113 | + pix_alpha = alpha |
| 114 | + else: |
| 115 | + x = i / float(LUT_SIZE - 1) |
| 116 | + if alpha_function == LUTAlpha.LINEAR.value: |
| 117 | + # Linear alpha function |
| 118 | + pix_alpha = alpha * x |
| 119 | + elif alpha_function == LUTAlpha.SIGMOID.value: |
| 120 | + # Sigmoid alpha function |
| 121 | + pix_alpha = alpha / (1 + np.exp(-10 * x)) |
| 122 | + elif alpha_function == LUTAlpha.TANH.value: |
| 123 | + # Hyperbolic tangent alpha function |
| 124 | + pix_alpha = alpha * np.tanh(5 * x) |
| 125 | + elif alpha_function == LUTAlpha.STEP.value: |
| 126 | + # Fully transparent lowest value and `alpha` transparent elsewhere |
| 127 | + pix_alpha = alpha if x > 0 else 0 |
| 128 | + else: |
| 129 | + raise ValueError(f"Invalid alpha function {alpha_function}") |
| 130 | + # pylint: disable=too-many-function-args |
| 131 | + alpha_channel = max(min(np.uint32(255 * pix_alpha + 0.5), 255), 0) << 24 |
| 132 | + lut[i] = ( |
| 133 | + np.uint32((table.rgb(FULLRANGE, i / LUT_MAX)) & 0xFFFFFF) | alpha_channel |
| 134 | + ) |
| 135 | + return colors, lut |
| 136 | + |
| 137 | + |
| 138 | +def _get_colormap_lut( |
| 139 | + table: EditableColormap, alpha: float, alpha_function: int |
| 140 | +) -> tuple[tuple[int, ...], np.ndarray]: |
| 141 | + """Return the color table and the lookup table of a colormap. |
| 142 | +
|
| 143 | + Computing them takes more than 1000 calls to the colormap: they are cached, |
| 144 | + since every new image item usually shows the same colormap. |
| 145 | +
|
| 146 | + Args: |
| 147 | + table: Colormap |
| 148 | + alpha: Global alpha value |
| 149 | + alpha_function: Alpha function (see :class:`plotpy.constants.LUTAlpha`) |
| 150 | +
|
| 151 | + Returns: |
| 152 | + Color table (256 colors) and lookup table (``LUT_SIZE`` ARGB values), |
| 153 | + the latter being read-only |
| 154 | + """ |
| 155 | + key = _lut_cache_key(table, alpha, alpha_function) |
| 156 | + if key is not None and key in _LUT_CACHE: |
| 157 | + _LUT_CACHE.move_to_end(key) |
| 158 | + return _LUT_CACHE[key] |
| 159 | + colors, lut = _compute_colormap_lut(table, alpha, alpha_function) |
| 160 | + lut.flags.writeable = False |
| 161 | + if key is not None: |
| 162 | + _LUT_CACHE[key] = (colors, lut) |
| 163 | + if len(_LUT_CACHE) > _LUT_CACHE_SIZE: |
| 164 | + _LUT_CACHE.popitem(last=False) |
| 165 | + return colors, lut |
61 | 166 |
|
62 | 167 |
|
63 | 168 | class BaseImageItem(QwtPlotItem): |
@@ -466,37 +571,11 @@ def set_color_map( |
466 | 571 | if invert is not None: |
467 | 572 | table.invert = invert |
468 | 573 | self.cmap_table = table |
469 | | - self.cmap = table.colorTable(FULLRANGE) |
470 | | - cmap_a = self.lut[3] |
471 | | - alpha = self.param.alpha |
472 | | - alpha_function = self.param.alpha_function |
473 | | - for i in range(LUT_SIZE): |
474 | | - if alpha_function == LUTAlpha.NONE.value: |
475 | | - pix_alpha = 1.0 |
476 | | - elif alpha_function == LUTAlpha.CONSTANT.value: |
477 | | - pix_alpha = alpha |
478 | | - else: |
479 | | - x = i / float(LUT_SIZE - 1) |
480 | | - if alpha_function == LUTAlpha.LINEAR.value: |
481 | | - # Linear alpha function |
482 | | - pix_alpha = alpha * x |
483 | | - elif alpha_function == LUTAlpha.SIGMOID.value: |
484 | | - # Sigmoid alpha function |
485 | | - pix_alpha = alpha / (1 + np.exp(-10 * x)) |
486 | | - elif alpha_function == LUTAlpha.TANH.value: |
487 | | - # Hyperbolic tangent alpha function |
488 | | - pix_alpha = alpha * np.tanh(5 * x) |
489 | | - elif alpha_function == LUTAlpha.STEP.value: |
490 | | - # Fully transparent lowest value and `alpha` transparent elsewhere |
491 | | - pix_alpha = alpha if x > 0 else 0 |
492 | | - else: |
493 | | - raise ValueError(f"Invalid alpha function {alpha_function}") |
494 | | - # pylint: disable=too-many-function-args |
495 | | - alpha_channel = max(min(np.uint32(255 * pix_alpha + 0.5), 255), 0) << 24 |
496 | | - cmap_a[i] = ( |
497 | | - np.uint32((table.rgb(FULLRANGE, i / LUT_MAX)) & 0xFFFFFF) |
498 | | - | alpha_channel |
499 | | - ) |
| 574 | + colors, lut = _get_colormap_lut( |
| 575 | + table, self.param.alpha, self.param.alpha_function |
| 576 | + ) |
| 577 | + self.cmap = list(colors) |
| 578 | + self.lut[3][:] = lut |
500 | 579 | plot = self.plot() |
501 | 580 | if plot: |
502 | 581 | plot.update_colormap_axis(self) |
|
0 commit comments