File: _3d_overlay.py

package info (click to toggle)
python-mne 1.9.0-2
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 131,492 kB
  • sloc: python: 213,302; javascript: 12,910; sh: 447; makefile: 144
file content (184 lines) | stat: -rw-r--r-- 5,549 bytes parent folder | download
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
"""Classes to handle overlapping surfaces."""

# Authors: The MNE-Python contributors.
# License: BSD-3-Clause
# Copyright the MNE-Python contributors.

from collections import OrderedDict

import numpy as np

from ..utils import logger


class _Overlay:
    def __init__(self, scalars, colormap, rng, opacity, name):
        self._scalars = scalars
        self._colormap = colormap
        assert rng is not None
        self._rng = rng
        self._opacity = opacity
        self._name = name

    def to_colors(self):
        from matplotlib.colors import Colormap, ListedColormap

        from ._3d import _get_cmap

        if isinstance(self._colormap, str):
            cmap = _get_cmap(self._colormap)
        elif isinstance(self._colormap, Colormap):
            cmap = self._colormap
        else:
            cmap = ListedColormap(
                self._colormap / 255.0, name=str(type(self._colormap))
            )
        logger.debug(
            f"Color mapping {repr(self._name)} with {cmap.name} "
            f"colormap and range {self._rng}"
        )

        rng = self._rng
        assert rng is not None
        scalars = self._norm(rng)

        colors = cmap(scalars)
        if self._opacity is not None:
            colors[:, 3] *= self._opacity
        return colors

    def _norm(self, rng):
        if rng[0] == rng[1]:
            factor = 1 if rng[0] == 0 else 1e-6 * rng[0]
        else:
            factor = rng[1] - rng[0]
        return (self._scalars - rng[0]) / factor


class _LayeredMesh:
    def __init__(self, renderer, vertices, triangles, normals):
        self._renderer = renderer
        self._vertices = vertices
        self._triangles = triangles
        self._normals = normals

        self._polydata = None
        self._actor = None
        self._is_mapped = False

        self._current_colors = None
        self._cached_colors = None
        self._overlays = OrderedDict()

        self._default_scalars = np.ones(vertices.shape)
        self._default_scalars_name = "Data"

    def map(self):
        kwargs = {
            "color": None,
            "pickable": True,
            "rgba": True,
        }
        mesh_data = self._renderer.mesh(
            x=self._vertices[:, 0],
            y=self._vertices[:, 1],
            z=self._vertices[:, 2],
            triangles=self._triangles,
            normals=self._normals,
            scalars=self._default_scalars,
            **kwargs,
        )
        self._actor, self._polydata = mesh_data
        self._is_mapped = True

    def _compute_over(self, B, A):
        assert A.ndim == B.ndim == 2
        assert A.shape[1] == B.shape[1] == 4
        A_w = A[:, 3:]  # * 1
        B_w = B[:, 3:] * (1 - A_w)
        C = A.copy()
        C[:, :3] *= A_w
        C[:, :3] += B[:, :3] * B_w
        C[:, 3:] += B_w
        C[:, :3] /= C[:, 3:]
        return np.clip(C, 0, 1, out=C)

    def _compose_overlays(self):
        B = cache = None
        for overlay in self._overlays.values():
            A = overlay.to_colors()
            if B is None:
                B = A
            else:
                cache = B
                B = self._compute_over(cache, A)
        return B, cache

    def add_overlay(self, scalars, colormap, rng, opacity, name):
        overlay = _Overlay(
            scalars=scalars,
            colormap=colormap,
            rng=rng,
            opacity=opacity,
            name=name,
        )
        self._overlays[name] = overlay
        colors = overlay.to_colors()
        if self._current_colors is None:
            self._current_colors = colors
        else:
            # save previous colors to cache
            self._cached_colors = self._current_colors
            self._current_colors = self._compute_over(self._cached_colors, colors)

        # apply the texture
        self._apply()

    def remove_overlay(self, names):
        to_update = False
        if not isinstance(names, list):
            names = [names]
        for name in names:
            if name in self._overlays:
                del self._overlays[name]
                to_update = True
        if to_update:
            self.update()

    def _apply(self):
        if self._current_colors is None or self._renderer is None:
            return
        self._polydata[self._default_scalars_name] = self._current_colors

    def update(self, colors=None):
        if colors is not None and self._cached_colors is not None:
            self._current_colors = self._compute_over(self._cached_colors, colors)
        else:
            self._current_colors, self._cached_colors = self._compose_overlays()
        self._apply()

    def _clean(self):
        mapper = self._actor.GetMapper()
        mapper.SetLookupTable(None)
        self._actor.SetMapper(None)
        self._actor = None
        self._polydata = None
        self._renderer = None

    def update_overlay(self, name, scalars=None, colormap=None, opacity=None, rng=None):
        overlay = self._overlays.get(name, None)
        if overlay is None:
            return
        if scalars is not None:
            overlay._scalars = scalars
        if colormap is not None:
            overlay._colormap = colormap
        if opacity is not None:
            overlay._opacity = opacity
        if rng is not None:
            overlay._rng = rng
        # partial update: use cache if possible
        if name == list(self._overlays.keys())[-1]:
            self.update(colors=overlay.to_colors())
        else:  # full update
            self.update()