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
|
# -*- coding: utf-8 -*-
# Copyright (c) Vispy Development Team. All Rights Reserved.
# Distributed under the (new) BSD License. See LICENSE.txt for more info.
from abc import ABCMeta, abstractmethod
import numpy as np
from vispy.gloo import VertexBuffer
from ..shaders import Function
class BaseFilter(object):
"""Superclass for all filters."""
def _attach(self, visual):
"""Called when a filter should be attached to a visual.
Parameters
----------
visual : instance of BaseVisual
The visual that called this.
"""
raise NotImplementedError(self)
def _detach(self, visual):
"""Called when a filter should be detached from a visual.
Parameters
----------
visual : instance of BaseVisual
The visual that called this.
"""
raise NotImplementedError(self)
class Filter(BaseFilter):
"""Base class for all filters that use fragment and/or vertex shaders.
Parameters
----------
vcode : str | Function | None
Vertex shader code. If None, ``vhook`` and ``vpos`` will
be ignored.
vhook : {'pre', 'post'}
Hook name to attach the vertex shader to.
vpos : int
Position in the hook to attach the vertex shader.
fcode : str | Function | None
Fragment shader code. If None, ``fhook`` and ``fpos`` will
be ignored.
fhook : {'pre', 'post'}
Hook name to attach the fragment shader to.
fpos : int
Position in the hook to attach the fragment shader.
Attributes
----------
vshader : Function | None
Vertex shader.
fshader : Function | None
Fragment shader.
"""
def __init__(self, vcode=None, vhook='post', vpos=5,
fcode=None, fhook='post', fpos=5):
super(Filter, self).__init__()
if vcode is not None:
self.vshader = Function(vcode) if isinstance(vcode, str) else vcode
self._vexpr = self.vshader()
self._vhook = vhook
self._vpos = vpos
else:
self.vshader = None
if fcode is not None:
self.fshader = Function(fcode) if isinstance(fcode, str) else fcode
self._fexpr = self.fshader()
self._fhook = fhook
self._fpos = fpos
else:
self.fshader = None
self._attached = False
@property
def attached(self):
return self._attached
def _attach(self, visual):
"""Called when a filter should be attached to a visual.
Parameters
----------
visual : instance of Visual
The visual that called this.
"""
if self.vshader:
hook = visual._get_hook('vert', self._vhook)
hook.add(self._vexpr, position=self._vpos)
if self.fshader:
hook = visual._get_hook('frag', self._fhook)
hook.add(self._fexpr, position=self._fpos)
self._attached = True
self._visual = visual
def _detach(self, visual):
"""Called when a filter should be detached from a visual.
Parameters
----------
visual : instance of Visual
The visual that called this.
"""
if self.vshader:
hook = visual._get_hook('vert', self._vhook)
hook.remove(self._vexpr)
if self.fshader:
hook = visual._get_hook('frag', self._fhook)
hook.remove(self._fexpr)
self._attached = False
self._visual = None
class PrimitivePickingFilter(Filter, metaclass=ABCMeta):
"""Abstract base class for Visual-specific filters to implement a
primitive-picking mode.
Subclasses must (and usually only need to) implement
:py:meth:`_get_picking_ids`.
"""
def __init__(self, fpos=9, *, discard_transparent=False):
# fpos is set to 9 by default to put it near the end, but before the
# default PickingFilter
vfunc = Function("""\
varying vec4 v_marker_picking_color;
void prepare_marker_picking() {
v_marker_picking_color = $ids;
}
""")
ffunc = Function("""\
varying vec4 v_marker_picking_color;
void marker_picking_filter() {
if ( $enabled != 1 ) {
return;
}
if ( $discard_transparent == 1 && gl_FragColor.a == 0.0 ) {
discard;
}
gl_FragColor = v_marker_picking_color;
}
""")
self._id_colors = VertexBuffer(np.zeros((0, 4), dtype=np.float32))
vfunc['ids'] = self._id_colors
self._n_primitives = 0
super().__init__(vcode=vfunc, fcode=ffunc, fpos=fpos)
self.enabled = False
self.discard_transparent = discard_transparent
@abstractmethod
def _get_picking_ids(self):
"""Return a 1D array of picking IDs for the vertices in the visual.
Generally, this method should be implemented to:
1. Calculate the number of primitives in the visual (may be
persisted in `self._n_primitives`).
2. Calculate a range of picking ids for each primitive in the
visual. IDs should start from 1, reserving 0 for the background. If
primitives comprise multiple vertices (triangles), ids may need to
be repeated.
The return value should be an array of uint32 with shape
(num_vertices,).
If no change to the picking IDs is needed (for example, the number of
primitives has not changed), this method should return `None`.
"""
raise NotImplementedError(self)
def _update_id_colors(self):
"""Calculate the colors encoding the picking IDs for the visual.
For performance, this method will not update the id colors VertexBuffer
if :py:meth:`_get_picking_ids` returns `None`.
"""
# this should remain untouched
ids = self._get_picking_ids()
if ids is not None:
id_colors = self._pack_ids_into_rgba(ids)
self._id_colors.set_data(id_colors)
@staticmethod
def _pack_ids_into_rgba(ids):
"""Pack an array of uint32 primitive ids into float32 RGBA colors."""
if ids.dtype != np.uint32:
raise ValueError(f"ids must be uint32, got {ids.dtype}")
return np.divide(
ids.view(np.uint8).reshape(-1, 4),
255,
dtype=np.float32
)
def _on_data_updated(self, event=None):
if not self.attached:
return
self._update_id_colors()
@property
def enabled(self):
return self._enabled
@enabled.setter
def enabled(self, e):
self._enabled = bool(e)
self.fshader['enabled'] = int(self._enabled)
self._on_data_updated()
@property
def discard_transparent(self):
return self._discard_transparent
@discard_transparent.setter
def discard_transparent(self, d):
self._discard_transparent = bool(d)
self.fshader['discard_transparent'] = int(self._discard_transparent)
def _attach(self, visual):
super()._attach(visual)
visual.events.data_updated.connect(self._on_data_updated)
self._on_data_updated()
def _detach(self, visual):
visual.events.data_updated.disconnect(self._on_data_updated)
super()._detach(visual)
|