Coverage for src/evutils/io/_hdf5.py: 92%
178 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-15 09:33 +0000
1"""HDF5 file decoder and encoder.
3Layout (DSEC-compatible): the four event columns are stored under
4``events/{t,x,y,p}``, with ``width`` / ``height`` file attributes and a
5top-level ``ms_to_idx`` index: ``ms_to_idx[ms]`` is the index of the first
6event with ``t >= ms * 1000`` (µs), which makes millisecond-range reads O(1)
7lookups. A ``t_offset`` attribute (DSEC) is honoured on read when present.
8"""
9from __future__ import annotations
11import io
12from datetime import datetime
13from typing import Any
15import h5py
16import hdf5plugin
17import numpy as np
19from .._jit import lazy_njit
20from ..types import EventArray, TriggerArray
21from .common import EventDecoder, EventEncoder
22from ._source import ByteSource
24_EMPTY_EVENTS = EventArray.empty()
27@lazy_njit
28def _fill_ms_to_idx(t: np.ndarray, ms_to_idx: np.ndarray, start_ms: int,
29 end_ms: int, base_idx: int, base_ms: int) -> None:
30 """Fill ``ms_to_idx[start_ms:end_ms+1]`` from the chunk timestamps ``t``.
32 ``ms_to_idx[ms]`` is the global index of the first event with
33 ``t >= (ms + base_ms) * 1000``; ``base_idx`` is the global index of
34 ``t[0]``. ``base_ms`` anchors the index at the recording's first
35 millisecond so absolute (e.g. epoch) timestamps do not blow it up.
37 Parameters
38 ----------
39 t : np.ndarray
40 Timestamps (µs) of the chunk, monotonically non-decreasing.
41 ms_to_idx : np.ndarray
42 The full index array to fill.
43 start_ms : int
44 First (relative) millisecond entry to fill.
45 end_ms : int
46 Last (relative) millisecond entry to fill (inclusive).
47 base_idx : int
48 Global event index of the first element of ``t``.
49 base_ms : int
50 Millisecond of the first event in the recording (index anchor).
52 """
53 idx = 0
54 for ms in range(start_ms, end_ms + 1):
55 while idx < len(t) and t[idx] < (ms + base_ms) * 1000:
56 idx += 1
57 ms_to_idx[ms] = base_idx + idx
60class EventDecoder_HDF5(EventDecoder):
61 """Decode events from an HDF5 file.
63 Two on-disk layouts are detected automatically:
65 * **DSEC / RVT layout** (what :class:`EventEncoder_HDF5` writes): the four
66 columns under ``events/{t,x,y,p}``, optional ``ms_to_idx`` index and
67 DSEC ``t_offset``.
68 * **Prophesee layout** (Metavision ``.hdf5``): a compound ``CD/events``
69 dataset with ``x``/``y``/``p``/``t`` fields. Prophesee files are usually
70 compressed with the ECF codec, a separate HDF5 plugin
71 (https://github.com/prophesee-ai/hdf5_ecf) -- a clear error points there
72 when it is missing. A compound ``events`` dataset at the root is read
73 the same way.
75 Supports both the streaming :meth:`read_chunk` interface used by
76 :class:`~evutils.io.EventReader` and random-access millisecond-range reads
77 via :meth:`read` (backed by the ``ms_to_idx`` index when present).
79 Parameters
80 ----------
81 source
82 Byte source to read from (must be seekable).
83 chunk_size
84 Number of events returned per :meth:`read_chunk` call.
86 """
88 def __init__(self, source: ByteSource, chunk_size: int = 1_000_000):
89 super().__init__(source, chunk_size)
90 self._h5: h5py.File | None = None
91 self._aos: h5py.Dataset | None = None # compound dataset (Prophesee layout)
92 self._ms_to_idx: np.ndarray | None = None
93 self._ms_idx_offset: int = 0 # ms of index entry 0 (absolute-t recordings)
94 self._t_offset: int = 0
95 self._n: int = 0
96 self._pos = 0
98 def init(self) -> None:
99 """Open the HDF5 file and locate the event datasets."""
100 if self._is_initialized:
101 return
103 self._h5 = h5py.File(self._source, "r")
104 node = None
105 if "events" in self._h5:
106 node = self._h5["events"]
107 elif "CD" in self._h5 and "events" in self._h5["CD"]:
108 node = self._h5["CD"]["events"] # Prophesee Metavision layout
109 if node is None:
110 raise ValueError(
111 "HDF5 file contains neither an 'events' group/dataset nor "
112 "'CD/events' (Prophesee layout)"
113 )
115 if isinstance(node, h5py.Dataset):
116 if node.dtype.names is None or not {"t", "x", "y", "p"} <= set(node.dtype.names):
117 raise ValueError(
118 "HDF5 events dataset must be a compound type with t/x/y/p fields"
119 )
120 self._aos = node
121 self._n = node.shape[0]
122 else:
123 self._n = node["t"].shape[0]
125 if "ms_to_idx" in self._h5:
126 self._ms_to_idx = np.asarray(self._h5["ms_to_idx"], dtype=np.int64)
127 if "ms_to_idx_offset" in self._h5:
128 self._ms_idx_offset = int(np.asarray(self._h5["ms_to_idx_offset"]).item())
129 if "t_offset" in self._h5:
130 self._t_offset = int(np.asarray(self._h5["t_offset"]).item())
131 if "width" in self._h5.attrs:
132 self._width = int(self._h5.attrs["width"])
133 if "height" in self._h5.attrs:
134 self._height = int(self._h5.attrs["height"])
136 self._pos = 0
137 self._is_initialized = True
139 def _slice(self, start: int, end: int) -> EventArray:
140 """Materialise ``events[start:end]`` as an :class:`EventArray`."""
141 assert self._h5 is not None
142 try:
143 if self._aos is not None:
144 rec = self._aos[start:end]
145 t = rec["t"].astype(np.int64)
146 if self._t_offset:
147 t += self._t_offset
148 return EventArray(t, rec["x"], rec["y"],
149 np.clip(rec["p"], 0, 1).astype(np.uint8))
150 ev = self._h5["events"]
151 t = ev["t"][start:end].astype(np.int64)
152 if self._t_offset:
153 t += self._t_offset
154 return EventArray(t, ev["x"][start:end], ev["y"][start:end], ev["p"][start:end])
155 except OSError as exc:
156 # h5py raises OSError when a dataset's filter/codec is not loaded.
157 raise OSError(
158 f"Failed to read the HDF5 events dataset ({exc}). If this is a "
159 "Prophesee Metavision file, its events are compressed with the "
160 "ECF codec: install the plugin from "
161 "https://github.com/prophesee-ai/hdf5_ecf (or put it on "
162 "HDF5_PLUGIN_PATH) and retry."
163 ) from exc
165 def read_chunk(self, delta_t_hint: int | None = None,
166 n_events_hint: int | None = None) -> EventArray:
167 if not self._is_initialized:
168 self.init()
170 if self._pos >= self._n:
171 self._eof = True
172 return _EMPTY_EVENTS
174 chunk = self._slice(self._pos, min(self._pos + self._chunk_size, self._n))
175 self._pos += len(chunk)
176 if self._pos >= self._n:
177 self._eof = True
178 return chunk
180 def read_all(self) -> EventArray:
181 """Return every remaining event at once."""
182 if not self._is_initialized:
183 self.init()
184 out = self._slice(self._pos, self._n)
185 self._pos = self._n
186 self._eof = True
187 return out
189 def read(self, start_ms: int = 0, end_ms: int = -1) -> EventArray:
190 """Random-access read of a millisecond time range via ``ms_to_idx``.
192 Parameters
193 ----------
194 start_ms : int, optional
195 Start time in milliseconds, by default 0.
196 end_ms : int, optional
197 End time in milliseconds (exclusive), by default -1 (until the end).
199 Returns
200 -------
201 EventArray
202 Events with ``start_ms * 1000 <= t < end_ms * 1000``.
204 """
205 if not self._is_initialized:
206 self.init()
207 if self._ms_to_idx is None:
208 raise ValueError("HDF5 file has no 'ms_to_idx' index; use read_chunk/read_all")
209 if start_ms < 0:
210 raise ValueError("start_ms must be greater or equal to 0")
211 if 0 <= end_ms < start_ms:
212 raise ValueError("start_ms must be smaller than end_ms")
214 # The index may be anchored at the recording's first millisecond
215 # (absolute-timestamp files); shift the requested range accordingly.
216 last_ms = len(self._ms_to_idx) - 1
217 rel_start = max(start_ms - self._ms_idx_offset, 0)
218 if rel_start > last_ms:
219 return _EMPTY_EVENTS
220 rel_end = end_ms - self._ms_idx_offset if end_ms >= 0 else last_ms
221 if rel_end < 0:
222 return _EMPTY_EVENTS
223 rel_end = min(rel_end, last_ms)
225 return self._slice(int(self._ms_to_idx[rel_start]), int(self._ms_to_idx[rel_end]))
227 def reset(self) -> None:
228 """Reset the reader to the beginning of the file."""
229 self._pos = 0
230 self._eof = False
232 def tell(self) -> int:
233 """Current position, in events (HDF5 has no meaningful byte offset)."""
234 return self._pos
236 def close(self) -> None:
237 """Close the HDF5 handle (the byte source is closed by the reader)."""
238 if self._h5 is not None:
239 self._h5.close()
240 self._h5 = None
243class EventEncoder_HDF5(EventEncoder):
244 """Encode events into an HDF5 file (``events/{t,x,y,p}`` + ``ms_to_idx``).
246 Events must be written in timestamp order (chunks are appended and the
247 millisecond index is built incrementally). The index and final flush
248 happen on :meth:`close`.
250 Parameters
251 ----------
252 writable : io.BufferedIOBase
253 The file-like object to write to (must be readable and seekable,
254 as required by HDF5).
255 width : int, optional
256 The width of the frame.
257 height : int, optional
258 The height of the frame.
259 dt : datetime, optional
260 Unused; HDF5 stores no recording timestamp.
261 chunksize : int, optional
262 HDF5 dataset chunk size, default 10000.
264 """
266 def __init__(self, writable: io.BufferedIOBase, width: int = 1280, height: int = 720,
267 dt: datetime | None = None, chunksize: int = 10000):
268 super().__init__(writable, width=width, height=height, dt=dt)
270 self._chunksize = chunksize
271 self._h5: h5py.File | None = None
272 self._ms_to_idx = np.zeros(0, dtype=np.int64)
273 self._next_ms = 0 # first (relative) ms entry not yet filled
274 self._ms_base = -1 # ms of the first written event (index anchor)
275 self._closed = False
277 def init(self) -> None:
278 """Create the HDF5 structure (groups, datasets, attributes)."""
279 if self._is_initialized:
280 return
282 self._h5 = h5py.File(self._fd, "w")
283 self._compressor = hdf5plugin.Blosc(cname="zstd", clevel=5, shuffle=hdf5plugin.Blosc.SHUFFLE)
285 self._h5.attrs["width"] = self._width
286 self._h5.attrs["height"] = self._height
288 group = self._h5.create_group("events")
289 for name, dtype in (("t", "int64"), ("x", "uint16"), ("y", "uint16"), ("p", "uint8")):
290 group.create_dataset(name, shape=(0,), chunks=(self._chunksize,),
291 maxshape=(None,), dtype=dtype, **self._compressor)
293 self._is_initialized = True
295 def write(self, events: 'np.ndarray | EventArray', triggers: 'np.ndarray | TriggerArray | None' = None) -> int:
296 """Append a chunk of events and extend the millisecond index.
298 Parameters
299 ----------
300 events : np.ndarray or EventArray
301 Array of events to write (timestamps must not go backwards
302 between chunks).
304 Returns
305 -------
306 int
307 Number of events written.
309 """
310 if not self._is_initialized:
311 self.init()
313 n = len(events)
314 if n == 0:
315 return 0
316 assert self._h5 is not None
318 t = np.ascontiguousarray(events["t"], dtype=np.int64)
320 # Extend ms_to_idx up to the last full millisecond of this chunk. The
321 # index is anchored at the first event's millisecond so absolute
322 # (epoch-style) timestamps don't inflate it.
323 if self._ms_base < 0:
324 self._ms_base = int(t[0] // 1000)
325 max_ms = int(t[-1] // 1000) - self._ms_base
326 if max_ms + 1 > len(self._ms_to_idx):
327 self._ms_to_idx = np.resize(self._ms_to_idx, max_ms + 1)
328 _fill_ms_to_idx(t, self._ms_to_idx, self._next_ms, max_ms,
329 self._n_written_events, self._ms_base)
330 self._next_ms = max_ms + 1
332 group = self._h5["events"]
333 total = self._n_written_events + n
334 for name, col in (("t", t), ("x", events["x"]), ("y", events["y"]), ("p", events["p"])):
335 ds = group[name]
336 ds.resize((total,))
337 ds[-n:] = col
339 self._n_written_events += n
340 return n
342 def flush(self) -> None:
343 """Flush the HDF5 buffers to the underlying stream."""
344 if self._h5 is not None:
345 self._h5.flush()
347 def close(self) -> None:
348 """Write the ``ms_to_idx`` index and close the HDF5 handle."""
349 if self._closed:
350 return
351 self._closed = True
352 if not self._is_initialized:
353 self.init() # produce a valid (empty) file even with no writes
354 assert self._h5 is not None
356 # Terminate the index: one entry past the last ms points at the end.
357 idx = np.append(self._ms_to_idx, self._n_written_events)
358 self._h5.create_dataset("ms_to_idx", data=idx.astype(np.uint64), **self._compressor)
359 if self._ms_base > 0:
360 # Anchor for absolute-timestamp recordings; the decoder shifts
361 # requested millisecond ranges by this.
362 self._h5.create_dataset("ms_to_idx_offset", data=np.int64(self._ms_base))
363 self._h5.close()
364 self._h5 = None