Coverage for src/evutils/transforms/transforms.py: 89%
128 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-18 05:24 +0000
1import math
2from typing import Union
4import numpy as np
6from evutils.types import SoaArray
8from .functional._common import sample_range
10class Transform:
11 """Base class for all evutils transforms.
13 Transforms should implement the `_forward_jit` method to allow zero-overhead
14 composition inside a `Compose` pipeline.
15 """
16 #: Sensor size resolved from the input container for the current call, used
17 #: as a fallback when a transform was constructed without an explicit one.
18 #: See :meth:`bind_context`.
19 _ctx_sensor_size = None
21 def bind_context(self, events):
22 """Capture per-call context (currently ``sensor_size``) from ``events``.
24 Called by :meth:`__call__` and by :class:`Compose` before ``_forward_jit``
25 so a transform can fall back to the container's ``sensor_size`` metadata
26 when it was not given one explicitly. Stored transiently, not persisted.
27 """
28 self._ctx_sensor_size = getattr(events, "sensor_size", None)
30 def _resolve_sensor_size(self):
31 """Return the explicit ``sensor_size`` or the one bound from the events.
33 Raises if neither is available.
34 """
35 ss = self.sensor_size if self.sensor_size is not None else self._ctx_sensor_size
36 if ss is None:
37 raise ValueError(
38 f"{type(self).__name__} needs a sensor_size: pass it to the "
39 f"constructor, or attach it to the events "
40 f"(events.sensor_size = (W, H))."
41 )
42 return ss
44 def __call__(self, events: Union[np.ndarray, SoaArray], target=None):
45 """Applies the transform in a standalone manner."""
46 from evutils.transforms.compose import unwrap_events, repack_events
47 if len(events) == 0:
48 if target is not None:
49 return events, target
50 return events
52 self.bind_context(events)
53 t, x, y, p = unwrap_events(events)
54 t, x, y, p = self._forward_jit(t, x, y, p)
55 events = repack_events(events, t, x, y, p)
57 if target is not None:
58 target = self._transform_target(target)
59 return events, target
60 return events
62 def _forward_jit(self, t: np.ndarray, x: np.ndarray, y: np.ndarray, p: np.ndarray):
63 """The pure array-math forward pass, ideally JIT-compiled."""
64 raise NotImplementedError("Transforms must implement _forward_jit")
66 def _transform_target(self, target):
67 """Pure Python transformation of the target (e.g. bounding boxes).
68 Defaults to doing nothing.
69 """
70 return target
72class DropEvent(Transform):
73 """Randomly drops each event with probability ``p``.
75 Tonic-compatible replacement for ``DropRandomEvents`` (kept as an alias).
76 Uses an independent Bernoulli mask per event, so the number dropped is
77 binomially distributed around ``p * n`` rather than exactly ``p * n``.
79 Parameters
80 ----------
81 p : float or tuple of float, optional
82 Drop probability in ``[0, 1)``. A ``(lo, hi)`` tuple is sampled uniformly
83 per call. Defaults to ``0.1``.
84 """
85 def __init__(self, p: Union[float, tuple] = 0.1, seed: int | None = None):
86 if not isinstance(p, (tuple, list)):
87 if math.isnan(p) or p < 0 or p >= 1:
88 raise ValueError("p must be in [0, 1)")
89 self.p = p
90 self.seed = seed
92 def _forward_jit(self, t, x, y, pol):
93 from evutils.transforms.functional import _drop_random_events_jit
94 prob = sample_range(self.p)
95 if prob <= 0:
96 return t, x, y, pol
97 return _drop_random_events_jit(t, x, y, pol, prob, self.seed if self.seed is not None else -1)
99 def __repr__(self):
100 return f"{self.__class__.__name__}(p={self.p}, seed={self.seed})"
102# Backwards-compatible alias for the pre-rename class name.
103DropRandomEvents = DropEvent
105class DropEventByTime(Transform):
106 """Drops every event inside one randomly-placed time window.
108 Parameters
109 ----------
110 duration_ratio : float or tuple of float, optional
111 Window length as a fraction of the recording span, in ``[0, 1)``. A
112 ``(lo, hi)`` tuple is sampled uniformly per call. Defaults to ``0.2``.
113 """
114 def __init__(self, duration_ratio: Union[float, tuple] = 0.2, seed: int | None = None):
115 self.duration_ratio = duration_ratio
116 self.seed = seed
118 def _forward_jit(self, t, x, y, p):
119 from evutils.transforms.functional import _drop_by_time_jit
120 ratio = sample_range(self.duration_ratio)
121 if ratio <= 0:
122 return t, x, y, p
123 return _drop_by_time_jit(t, x, y, p, ratio, self.seed if self.seed is not None else -1)
125 def __repr__(self):
126 return f"{self.__class__.__name__}(duration_ratio={self.duration_ratio}, seed={self.seed})"
128class RandomFlipLR(Transform):
129 """Flips events horizontally (``x' = width - 1 - x``) with probability ``p``.
131 Parameters
132 ----------
133 sensor_size : tuple, optional
134 ``(W, H)`` (or ``(W, H, P)``) sensor size; only the width ``W`` is used.
135 If omitted, it is taken from the events' ``sensor_size`` metadata.
136 p : float, optional
137 Probability of performing the flip. Defaults to ``0.5``.
138 """
139 def __init__(self, sensor_size: tuple = None, p: float = 0.5):
140 if not 0 <= p <= 1:
141 raise ValueError("p must be in [0, 1]")
142 self.sensor_size = sensor_size
143 self.p = p
145 def _forward_jit(self, t, x, y, pol):
146 from evutils.transforms.functional import _flip_lr_jit
147 if np.random.rand() <= self.p:
148 width = int(self._resolve_sensor_size()[0])
149 return _flip_lr_jit(t, x, y, pol, width)
150 return t, x, y, pol
152 def __repr__(self):
153 return f"{self.__class__.__name__}(sensor_size={self.sensor_size}, p={self.p})"
155class SpatialJitter(Transform):
156 """Adds correlated Gaussian noise to event coordinates.
158 Parameters
159 ----------
160 sensor_size : tuple, optional
161 ``(W, H)`` (or ``(W, H, P)``) sensor size, used for clipping. Only needed
162 when ``clip_outliers`` is True; if omitted, taken from the events'
163 ``sensor_size`` metadata.
164 var_x, var_y : float, optional
165 Variances of the jitter in x and y. Default ``1.0``.
166 sigma_xy : float, optional
167 Off-diagonal covariance. Default ``0.0``.
168 clip_outliers : bool, optional
169 Drop events jittered outside the sensor. Default ``False``.
170 """
171 def __init__(self, sensor_size: tuple = None, var_x: float = 1.0, var_y: float = 1.0,
172 sigma_xy: float = 0.0, clip_outliers: bool = False):
173 self.sensor_size = sensor_size
174 self.var_x = var_x
175 self.var_y = var_y
176 self.sigma_xy = sigma_xy
177 self.clip_outliers = clip_outliers
179 def _forward_jit(self, t, x, y, p):
180 from evutils.transforms.functional import _spatial_jitter_jit
181 # Sensor size is only consulted when clipping; avoid requiring it otherwise.
182 if self.clip_outliers:
183 ss = self._resolve_sensor_size()
184 width, height = int(ss[0]), int(ss[1])
185 else:
186 width, height = 0, 0
187 return _spatial_jitter_jit(t, x, y, p, width, height, float(self.var_x),
188 float(self.var_y), float(self.sigma_xy),
189 bool(self.clip_outliers))
191 def __repr__(self):
192 return (f"{self.__class__.__name__}(sensor_size={self.sensor_size}, "
193 f"var_x={self.var_x}, var_y={self.var_y}, sigma_xy={self.sigma_xy}, "
194 f"clip_outliers={self.clip_outliers})")
196class TimeSkew(Transform):
197 """Rescales timestamps by an affine map ``t' = t * coefficient + offset``.
199 Parameters
200 ----------
201 coefficient : float or tuple of float
202 Multiplier applied to every timestamp. A ``(lo, hi)`` tuple is sampled
203 uniformly per call.
204 offset : float or tuple of float, optional
205 Added after multiplication. Default ``0``.
206 """
207 def __init__(self, coefficient: Union[float, tuple],
208 offset: Union[float, tuple] = 0):
209 self.coefficient = coefficient
210 self.offset = offset
212 def _forward_jit(self, t, x, y, p):
213 from evutils.transforms.functional import _time_skew_jit
214 coef = sample_range(self.coefficient)
215 off = sample_range(self.offset)
216 return _time_skew_jit(t, x, y, p, coef, off)
218 def __repr__(self):
219 return f"{self.__class__.__name__}(coefficient={self.coefficient}, offset={self.offset})"
221class TimeNormalize(Transform):
222 """Shifts timestamps so the earliest event lands at ``start_ts``.
224 Stateless: each call normalizes the batch it is given on its own. For a
225 chunk-by-chunk stream use ``EventReader(normalize_ts=True)``, which latches
226 the offset from the first chunk across the whole stream; dropping this
227 transform into a ``Compose`` over sequential chunks would instead reset each
228 chunk to ``start_ts`` independently.
230 Parameters
231 ----------
232 start_ts : int, optional
233 Timestamp assigned to the earliest event. Default ``0``.
234 """
235 def __init__(self, start_ts: int = 0):
236 self.start_ts = start_ts
238 def _forward_jit(self, t, x, y, p):
239 from evutils.transforms.functional import _normalize_ts_jit
240 return _normalize_ts_jit(t, x, y, p, int(self.start_ts))
242 def __repr__(self):
243 return f"{self.__class__.__name__}(start_ts={self.start_ts})"
245class TimeJitter(Transform):
246 """Adds Gaussian noise to each timestamp.
248 Parameters
249 ----------
250 std : float
251 Standard deviation of the timestamp noise.
252 clip_negative : bool, optional
253 Drop events with negative jittered timestamps. Default ``True``.
254 sort_timestamps : bool, optional
255 Re-sort by timestamp after jittering. Default ``False``.
256 """
257 def __init__(self, std: float, clip_negative: bool = True,
258 sort_timestamps: bool = False):
259 self.std = std
260 self.clip_negative = clip_negative
261 self.sort_timestamps = sort_timestamps
263 def _forward_jit(self, t, x, y, p):
264 from evutils.transforms.functional import _time_jitter_jit
265 return _time_jitter_jit(t, x, y, p, float(self.std),
266 bool(self.clip_negative), bool(self.sort_timestamps))
268 def __repr__(self):
269 return (f"{self.__class__.__name__}(std={self.std}, "
270 f"clip_negative={self.clip_negative}, sort_timestamps={self.sort_timestamps})")
272class RefractoryPeriod(Transform):
273 """Enforces a per-pixel refractory period.
275 Parameters
276 ----------
277 delta : int or tuple of int
278 Refractory period in timestamp units. A ``(lo, hi)`` tuple is sampled
279 uniformly per call. Events must be sorted by timestamp.
280 """
281 def __init__(self, delta: Union[int, tuple]):
282 self.delta = delta
284 def _forward_jit(self, t, x, y, p):
285 from evutils.transforms.functional import _refractory_period_jit
286 delta = int(sample_range(self.delta))
287 return _refractory_period_jit(t, x, y, p, delta)
289 def __repr__(self):
290 return f"{self.__class__.__name__}(delta={self.delta})"