Coverage for src/evutils/io/_hdf5.py: 91%
205 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-18 05:24 +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
14import h5py
15import hdf5plugin
16import numpy as np
18from ..jit import lazy_njit
19from ..types import EventArray, TriggerArray
20from .common import EventDecoder, EventEncoder
21from ._source import ByteSource
23_EMPTY_EVENTS = EventArray.empty()
25@lazy_njit
26def _fill_ms_to_idx(t: np.ndarray, ms_to_idx: np.ndarray, start_ms: int,
27 end_ms: int, base_idx: int, base_ms: int) -> None:
28 """Fill ``ms_to_idx[start_ms:end_ms+1]`` from the chunk timestamps ``t``.
30 ``ms_to_idx[ms]`` is the global index of the first event with
31 ``t >= (ms + base_ms) * 1000``; ``base_idx`` is the global index of
32 ``t[0]``. ``base_ms`` anchors the index at the recording's first
33 millisecond so absolute (e.g. epoch) timestamps do not blow it up.
35 Parameters
36 ----------
37 t : np.ndarray
38 Timestamps (µs) of the chunk, monotonically non-decreasing.
39 ms_to_idx : np.ndarray
40 The full index array to fill.
41 start_ms : int
42 First (relative) millisecond entry to fill.
43 end_ms : int
44 Last (relative) millisecond entry to fill (inclusive).
45 base_idx : int
46 Global event index of the first element of ``t``.
47 base_ms : int
48 Millisecond of the first event in the recording (index anchor).
50 """
51 idx = 0
52 for ms in range(start_ms, end_ms + 1):
53 while idx < len(t) and t[idx] < (ms + base_ms) * 1000:
54 idx += 1
55 ms_to_idx[ms] = base_idx + idx
57class EventDecoder_HDF5(EventDecoder):
58 """Decode events from an HDF5 file.
60 Two on-disk layouts are detected automatically:
62 * **DSEC / RVT layout** (what :class:`EventEncoder_HDF5` writes): the four
63 columns under ``events/{t,x,y,p}``, optional ``ms_to_idx`` index and
64 DSEC ``t_offset``.
65 * **Prophesee layout** (Metavision ``.hdf5``): a compound ``CD/events``
66 dataset with ``x``/``y``/``p``/``t`` fields. Prophesee files are usually
67 compressed with the ECF codec, a separate HDF5 plugin
68 (https://github.com/prophesee-ai/hdf5_ecf) -- a clear error points there
69 when it is missing. A compound ``events`` dataset at the root is read
70 the same way.
72 Supports both the streaming :meth:`read_chunk` interface used by
73 :class:`~evutils.io.EventReader` and random-access millisecond-range reads
74 via :meth:`read` (backed by the ``ms_to_idx`` index when present).
76 Parameters
77 ----------
78 source
79 Byte source to read from (must be seekable).
80 chunk_size
81 Number of events returned per :meth:`read_chunk` call.
83 """
85 #: Columns are index-addressable datasets, so seeking sets the event
86 #: position directly (index) or via searchsorted on the timestamp column
87 #: (time). The separate :meth:`read` remains for ms-range random access.
88 SUPPORTS_SEEK = True
90 def __init__(self, source: ByteSource, chunk_size: int = 1_000_000):
91 super().__init__(source, chunk_size)
92 self._h5: h5py.File | None = None
93 self._aos: h5py.Dataset | None = None # compound dataset (Prophesee layout)
94 self._ms_to_idx: np.ndarray | None = None
95 self._ms_idx_offset: int = 0 # ms of index entry 0 (absolute-t recordings)
96 self._t_offset: int = 0
97 self._n: int = 0
98 self._pos = 0
99 self._t_cache: np.ndarray | None = None # lazily-loaded t column for time seek
101 def init(self) -> None:
102 """Open the HDF5 file and locate the event datasets."""
103 if self._is_initialized:
104 return
106 # h5py needs a seekable file-like. A non-seekable source (pipe,
107 # compressed stream) is slurped into an in-memory buffer (mirrors the
108 # NPZ/CSV decoders).
109 h5_source: "ByteSource | io.BytesIO" = self._source
110 if not self._source.seekable():
111 h5_source = io.BytesIO(self._source.read(-1))
112 self._h5 = h5py.File(h5_source, "r")
113 node = None
114 if "events" in self._h5:
115 node = self._h5["events"]
116 elif "CD" in self._h5 and "events" in self._h5["CD"]:
117 node = self._h5["CD"]["events"] # Prophesee Metavision layout
118 if node is None:
119 raise ValueError(
120 "HDF5 file contains neither an 'events' group/dataset nor "
121 "'CD/events' (Prophesee layout)"
122 )
124 if isinstance(node, h5py.Dataset):
125 if node.dtype.names is None or not {"t", "x", "y", "p"} <= set(node.dtype.names):
126 raise ValueError(
127 "HDF5 events dataset must be a compound type with t/x/y/p fields"
128 )
129 self._aos = node
130 self._n = node.shape[0]
131 else:
132 self._n = node["t"].shape[0]
134 if "ms_to_idx" in self._h5:
135 self._ms_to_idx = np.asarray(self._h5["ms_to_idx"], dtype=np.int64)
136 if "ms_to_idx_offset" in self._h5:
137 self._ms_idx_offset = int(np.asarray(self._h5["ms_to_idx_offset"]).item())
138 if "t_offset" in self._h5:
139 self._t_offset = int(np.asarray(self._h5["t_offset"]).item())
140 if "width" in self._h5.attrs:
141 self._width = int(self._h5.attrs["width"])
142 if "height" in self._h5.attrs:
143 self._height = int(self._h5.attrs["height"])
145 self._pos = 0
146 self._is_initialized = True
148 def _slice(self, start: int, end: int) -> EventArray:
149 """Materialise ``events[start:end]`` as an :class:`EventArray`."""
150 assert self._h5 is not None
151 try:
152 if self._aos is not None:
153 rec = self._aos[start:end]
154 t = rec["t"].astype(np.int64)
155 if self._t_offset:
156 t += self._t_offset
157 return EventArray(t, rec["x"], rec["y"],
158 np.clip(rec["p"], 0, 1).astype(np.uint8))
159 ev = self._h5["events"]
160 t = ev["t"][start:end].astype(np.int64)
161 if self._t_offset:
162 t += self._t_offset
163 return EventArray(t, ev["x"][start:end], ev["y"][start:end], ev["p"][start:end])
164 except OSError as exc:
165 # h5py raises OSError when a dataset's filter/codec is not loaded.
166 raise OSError(
167 f"Failed to read the HDF5 events dataset ({exc}). If this is a "
168 "Prophesee Metavision file, its events are compressed with the "
169 "ECF codec: install the plugin from "
170 "https://github.com/prophesee-ai/hdf5_ecf (or put it on "
171 "HDF5_PLUGIN_PATH) and retry."
172 ) from exc
174 def read_chunk(self, delta_t_hint: int | None = None,
175 n_events_hint: int | None = None) -> EventArray:
176 if not self._is_initialized:
177 self.init()
179 if self._pos >= self._n:
180 self._eof = True
181 return _EMPTY_EVENTS
183 chunk = self._slice(self._pos, min(self._pos + self._chunk_size, self._n))
184 self._pos += len(chunk)
185 if self._pos >= self._n:
186 self._eof = True
187 return chunk
189 def read_all(self) -> EventArray:
190 """Return every remaining event at once."""
191 if not self._is_initialized:
192 self.init()
193 out = self._slice(self._pos, self._n)
194 self._pos = self._n
195 self._eof = True
196 return out
198 def read(self, start_ms: int = 0, end_ms: int = -1) -> EventArray:
199 """Random-access read of a millisecond time range via ``ms_to_idx``.
201 Parameters
202 ----------
203 start_ms : int, optional
204 Start time in milliseconds, by default 0.
205 end_ms : int, optional
206 End time in milliseconds (exclusive), by default -1 (until the end).
208 Returns
209 -------
210 EventArray
211 Events with ``start_ms * 1000 <= t < end_ms * 1000``.
213 """
214 if not self._is_initialized:
215 self.init()
216 if self._ms_to_idx is None:
217 raise ValueError("HDF5 file has no 'ms_to_idx' index; use read_chunk/read_all")
218 if start_ms < 0:
219 raise ValueError("start_ms must be greater or equal to 0")
220 if 0 <= end_ms < start_ms:
221 raise ValueError("start_ms must be smaller than end_ms")
223 # The index may be anchored at the recording's first millisecond
224 # (absolute-timestamp files); shift the requested range accordingly.
225 last_ms = len(self._ms_to_idx) - 1
226 rel_start = max(start_ms - self._ms_idx_offset, 0)
227 if rel_start > last_ms:
228 return _EMPTY_EVENTS
229 rel_end = end_ms - self._ms_idx_offset if end_ms >= 0 else last_ms
230 if rel_end < 0:
231 return _EMPTY_EVENTS
232 rel_end = min(rel_end, last_ms)
234 return self._slice(int(self._ms_to_idx[rel_start]), int(self._ms_to_idx[rel_end]))
236 def _all_t(self) -> np.ndarray:
237 """Lazily load the whole timestamp column (user timeline, +t_offset)."""
238 if self._t_cache is None:
239 assert self._h5 is not None
240 if self._aos is not None:
241 t = self._aos["t"][:].astype(np.int64)
242 else:
243 t = self._h5["events"]["t"][:].astype(np.int64)
244 if self._t_offset:
245 t = t + self._t_offset
246 self._t_cache = t
247 return self._t_cache
249 def seek(self, t: int | None = None, n: int | None = None) -> tuple["SeekResult", "EventArray", "TriggerArray | None"]:
250 """Seek to an absolute timestamp (µs) or event index. See base class.
252 Index seek sets the event position directly; time seek does a
253 ``searchsorted`` over the (lazily cached) timestamp column, so it is
254 exact regardless of ``t_offset`` / ``ms_to_idx`` anchoring.
255 """
256 from .common import SeekResult
257 if not self._is_initialized:
258 self.init()
259 axis, val = self._seek_axis(t, n)
261 if axis == "t":
262 idx = int(np.searchsorted(self._all_t(), val, side="left"))
263 else:
264 idx = val
266 idx = max(0, min(idx, self._n))
267 self._pos = idx
268 self._eof = idx >= self._n
270 landed_ts = int(self._all_t()[idx]) if idx < self._n else val
271 return SeekResult(ts=landed_ts, index=idx, eof=self._eof), _EMPTY_EVENTS, None
273 def reset(self) -> None:
274 """Reset the reader to the beginning of the file."""
275 self._pos = 0
276 self._eof = False
278 def tell(self) -> int:
279 """Current position, in events (HDF5 has no meaningful byte offset)."""
280 return self._pos
282 def close(self) -> None:
283 """Close the HDF5 handle (the byte source is closed by the reader)."""
284 if self._h5 is not None:
285 self._h5.close()
286 self._h5 = None
288class EventEncoder_HDF5(EventEncoder):
289 """Encode events into an HDF5 file (``events/{t,x,y,p}`` + ``ms_to_idx``).
291 Events must be written in timestamp order (chunks are appended and the
292 millisecond index is built incrementally). The index and final flush
293 happen on :meth:`close`.
295 Parameters
296 ----------
297 writable : io.BufferedIOBase
298 The file-like object to write to (must be readable and seekable,
299 as required by HDF5).
300 width : int, optional
301 The width of the frame.
302 height : int, optional
303 The height of the frame.
304 dt : datetime, optional
305 Unused; HDF5 stores no recording timestamp.
306 chunksize : int, optional
307 HDF5 dataset chunk size, default 10000.
309 """
311 def __init__(self, writable: io.BufferedIOBase, width: int = 1280, height: int = 720,
312 dt: datetime | None = None, chunksize: int = 10000):
313 super().__init__(writable, width=width, height=height, dt=dt)
315 self._chunksize = chunksize
316 self._h5: h5py.File | None = None
317 self._ms_to_idx = np.zeros(0, dtype=np.int64)
318 self._next_ms = 0 # first (relative) ms entry not yet filled
319 self._ms_base = -1 # ms of the first written event (index anchor)
320 self._closed = False
322 def init(self) -> None:
323 """Create the HDF5 structure (groups, datasets, attributes)."""
324 if self._is_initialized:
325 return
327 self._h5 = h5py.File(self._fd, "w")
328 self._compressor = hdf5plugin.Blosc(cname="zstd", clevel=5, shuffle=hdf5plugin.Blosc.SHUFFLE)
330 self._h5.attrs["width"] = self._width
331 self._h5.attrs["height"] = self._height
333 group = self._h5.create_group("events")
334 for name, dtype in (("t", "int64"), ("x", "uint16"), ("y", "uint16"), ("p", "uint8")):
335 group.create_dataset(name, shape=(0,), chunks=(self._chunksize,),
336 maxshape=(None,), dtype=dtype, **self._compressor)
338 self._is_initialized = True
340 def write(self, events: 'np.ndarray | EventArray', triggers: 'np.ndarray | TriggerArray | None' = None) -> int:
341 """Append a chunk of events and extend the millisecond index.
343 Parameters
344 ----------
345 events : np.ndarray or EventArray
346 Array of events to write (timestamps must not go backwards
347 between chunks).
349 Returns
350 -------
351 int
352 Number of events written.
354 """
355 if not self._is_initialized:
356 self.init()
358 n = len(events)
359 if n == 0:
360 return 0
361 assert self._h5 is not None
363 t = np.ascontiguousarray(events["t"], dtype=np.int64)
365 # Extend ms_to_idx up to the last full millisecond of this chunk. The
366 # index is anchored at the first event's millisecond so absolute
367 # (epoch-style) timestamps don't inflate it.
368 if self._ms_base < 0:
369 self._ms_base = int(t[0] // 1000)
370 max_ms = int(t[-1] // 1000) - self._ms_base
371 if max_ms + 1 > len(self._ms_to_idx):
372 self._ms_to_idx = np.resize(self._ms_to_idx, max_ms + 1)
373 _fill_ms_to_idx(t, self._ms_to_idx, self._next_ms, max_ms,
374 self._n_written_events, self._ms_base)
375 self._next_ms = max_ms + 1
377 group = self._h5["events"]
378 total = self._n_written_events + n
379 for name, col in (("t", t), ("x", events["x"]), ("y", events["y"]), ("p", events["p"])):
380 ds = group[name]
381 ds.resize((total,))
382 ds[-n:] = col
384 self._n_written_events += n
385 return n
387 def flush(self) -> None:
388 """Flush the HDF5 buffers to the underlying stream."""
389 if self._h5 is not None:
390 self._h5.flush()
392 def close(self) -> None:
393 """Write the ``ms_to_idx`` index and close the HDF5 handle."""
394 if self._closed:
395 return
396 self._closed = True
397 if not self._is_initialized:
398 self.init() # produce a valid (empty) file even with no writes
399 assert self._h5 is not None
401 # Terminate the index: one entry past the last ms points at the end.
402 idx = np.append(self._ms_to_idx, self._n_written_events)
403 self._h5.create_dataset("ms_to_idx", data=idx.astype(np.uint64), **self._compressor)
404 if self._ms_base > 0:
405 # Anchor for absolute-timestamp recordings; the decoder shifts
406 # requested millisecond ranges by this.
407 self._h5.create_dataset("ms_to_idx_offset", data=np.int64(self._ms_base))
408 self._h5.close()
409 self._h5 = None