Skip to main content

eventcv_core/
io.rs

1use std::{error::Error, fmt, io, path::Path};
2
3use crate::{EventStream, EventStreamBuilder};
4
5mod aedat;
6mod bag;
7#[cfg(feature = "hdf5")]
8mod h5;
9mod npz;
10mod prophesee;
11mod text;
12
13pub use aedat::read_aedat;
14pub use bag::{open_bag_slice, read_bag, write_bag, BagSliceSource};
15#[cfg(feature = "hdf5")]
16pub use h5::{
17    open_hdf5_slice, read_hdf5, read_hdf5_frame, write_hdf5_frame, write_hdf5_stream,
18    Hdf5FrameSink, Hdf5SliceSource,
19};
20pub use npz::{read_npz, read_npz_frame, write_npz_frame, write_npz_stream};
21pub use prophesee::read_dat;
22pub use text::{
23    load_rows, open_text_slice, read_text, write_text_stream, ColumnOrder, RawRow, TextOptions,
24    TextReader, TimeUnit,
25};
26
27use crate::representation::EventFrame;
28use crate::viz::{render_frame, Colormap};
29
30/// A single event as produced by a reader, before it is placed on the sensor grid.
31#[derive(Clone, Copy, Debug, PartialEq, Eq)]
32pub struct RawEvent {
33    pub x: u16,
34    pub y: u16,
35    pub t: i64,
36    pub p: bool,
37}
38
39/// A bounded-memory source of events. Readers parse on demand so multi-gigabyte
40/// files never need to be resident; the consumer decides what to accumulate.
41pub trait EventSource {
42    fn sensor_size(&self) -> (usize, usize);
43    fn timestamp_scale_ms(&self) -> f64;
44    /// Returns the next event, or `None` at end of stream.
45    fn next_event(&mut self) -> Result<Option<RawEvent>, IoError>;
46}
47
48/// Drains a source into an in-memory [`EventStream`], dropping out-of-bounds events.
49pub fn read_all(source: impl EventSource) -> Result<EventStream, IoError> {
50    read_capped(source, None)
51}
52
53/// Like [`read_all`] but stops after `max` kept events (for previewing huge files).
54pub fn read_capped(
55    mut source: impl EventSource,
56    max: Option<usize>,
57) -> Result<EventStream, IoError> {
58    let (width, height) = source.sensor_size();
59    let mut builder = EventStreamBuilder::new(width, height, source.timestamp_scale_ms());
60    while let Some(event) = source.next_event()? {
61        builder.push(event.x, event.y, event.t, event.p);
62        if max.is_some_and(|max| builder.len() >= max) {
63            break;
64        }
65    }
66    Ok(builder.build())
67}
68
69/// Options for the unified [`load`] entry point. Most fields apply to a single
70/// format; readers ignore the ones they do not need.
71#[derive(Clone, Debug, Default)]
72pub struct LoadOptions {
73    /// `(width, height)`. `None` infers it from the data (coordinate range), or from the
74    /// message for rosbag. An explicit value overrides and, for HDF5, skips the scan.
75    pub sensor_size: Option<(usize, usize)>,
76    /// Timestamp unit. `None` infers it (HDF5/text): a fractional text value means
77    /// seconds, otherwise the unit is chosen from the recording span (see
78    /// `TimeUnit::infer_from_span`). Ignored for rosbag (always ROS `sec`+`nsec`).
79    pub time_unit: Option<TimeUnit>,
80    /// Text column order.
81    pub order: ColumnOrder,
82    /// Rosbag topic to read (defaults to `/davis/left/events`).
83    pub topic: Option<String>,
84    /// Cap on the number of events to read.
85    pub max_events: Option<usize>,
86    /// Absolute timestamp (µs, the file's own time base): events before it are skipped.
87    /// `max_events` then caps the events *after* the offset. `None`/`<= 0` reads from the start.
88    pub offset: Option<i64>,
89}
90
91#[derive(Clone, Copy, Debug, PartialEq, Eq)]
92enum Format {
93    Npz,
94    Text,
95    Hdf5,
96    Rosbag,
97    Aedat,
98    Aedat4,
99    PropheseeDat,
100    PropheseeRaw,
101    Png,
102}
103
104fn detect_format(path: &Path) -> Result<Format, IoError> {
105    let extension = path
106        .extension()
107        .and_then(|extension| extension.to_str())
108        .map(str::to_ascii_lowercase);
109    match extension.as_deref() {
110        Some("npz") => Ok(Format::Npz),
111        Some("txt") | Some("csv") => Ok(Format::Text),
112        Some("h5") | Some("hdf5") => Ok(Format::Hdf5),
113        Some("bag") => Ok(Format::Rosbag),
114        Some("aedat") => Ok(Format::Aedat),
115        Some("aedat4") => Ok(Format::Aedat4),
116        Some("dat") => Ok(Format::PropheseeDat),
117        Some("raw") => Ok(Format::PropheseeRaw),
118        Some("png") => Ok(Format::Png),
119        Some(other) => Err(IoError::Unsupported(format!(
120            "unrecognised file extension: .{other}"
121        ))),
122        None => Err(IoError::Unsupported(
123            "file has no extension to detect its format".to_owned(),
124        )),
125    }
126}
127
128/// Loads events from any supported file, detected by extension — the OpenCV-style
129/// single entry point. Supported today: `.npz`, `.txt`/`.csv`, `.bag`, `.h5`/`.hdf5`,
130/// `.aedat` (AEDAT 2.0), and `.dat` (Prophesee CD).
131pub fn load(path: impl AsRef<Path>, options: LoadOptions) -> Result<EventStream, IoError> {
132    let path = path.as_ref();
133    let Some(cutoff) = options.offset.filter(|&offset| offset > 0) else {
134        return load_format(path, &options);
135    };
136    // The offset is an absolute timestamp, so the whole recording must be read before
137    // skipping; the cap then applies to the events at/after the offset.
138    let mut read_options = options.clone();
139    read_options.max_events = None;
140    load_format(path, &read_options).map(|stream| skip_before(stream, cutoff, options.max_events))
141}
142
143/// Drops events before the absolute timestamp `cutoff` (µs), then keeps at most `max` of the rest.
144fn skip_before(stream: EventStream, cutoff: i64, max: Option<usize>) -> EventStream {
145    let (ts, (xs, ys, ps)) = (stream.ts(), (stream.xs(), stream.ys(), stream.ps()));
146    if ts.is_empty() {
147        return stream;
148    }
149    let (width, height) = stream.sensor_size();
150    let mut builder = EventStreamBuilder::new(width, height, stream.timestamp_scale_ms());
151    for index in (0..ts.len()).filter(|&index| ts[index] >= cutoff) {
152        builder.push(xs[index], ys[index], ts[index], ps[index]);
153        if max.is_some_and(|max| builder.len() >= max) {
154            break;
155        }
156    }
157    builder.build()
158}
159
160fn load_format(path: &Path, options: &LoadOptions) -> Result<EventStream, IoError> {
161    match detect_format(path)? {
162        Format::Npz => npz::read_npz(path, options.sensor_size),
163        Format::Text => text::load_text(path, options),
164        Format::Rosbag => bag::read_bag(path, options),
165        Format::Hdf5 => {
166            #[cfg(feature = "hdf5")]
167            {
168                h5::read_hdf5(path, options)
169            }
170            #[cfg(not(feature = "hdf5"))]
171            {
172                Err(IoError::Unsupported(
173                    "HDF5 support is not built in; rebuild with --features hdf5".to_owned(),
174                ))
175            }
176        }
177        Format::Aedat => aedat::read_aedat(path, options),
178        Format::Aedat4 => Err(IoError::Unsupported(
179            "AEDAT4 (.aedat4, iniVation DV FlatBuffer/LZ4) reading is not implemented yet"
180                .to_owned(),
181        )),
182        Format::PropheseeDat => prophesee::read_dat(path, options),
183        Format::PropheseeRaw => Err(IoError::Unsupported(
184            "Prophesee .raw (EVT2/EVT3) reading is not implemented yet".to_owned(),
185        )),
186        Format::Png => Err(IoError::Unsupported(
187            "PNG is a frame export format, not an event stream; use save_frame".to_owned(),
188        )),
189    }
190}
191
192/// Random-access view over a file's events: fetch an arbitrary time or count range
193/// without materialising the whole stream. This backs the lazy [`open`] handle — the
194/// OpenCV `VideoCapture` to [`load`]'s `imread`. Each call returns a new [`EventStream`].
195pub trait SliceSource: Send {
196    fn sensor_size(&self) -> (usize, usize);
197    fn timestamp_scale_ms(&self) -> f64;
198    /// Total events in the file.
199    fn n_events(&self) -> usize;
200    /// `(t_min, t_max)` in microseconds across the whole file; `(0, 0)` when empty.
201    fn time_span(&self) -> (i64, i64);
202    /// Events whose index lies in `[i0, i1)` (clamped to the file).
203    fn slice_index(&self, i0: usize, i1: usize) -> Result<EventStream, IoError>;
204    /// Events whose timestamp (µs) lies in the half-open window `[t0, t1)`.
205    fn slice_time(&self, t0: i64, t1: i64) -> Result<EventStream, IoError>;
206}
207
208/// A [`SliceSource`] backed by an already-loaded stream — the universal fallback for
209/// formats without native random access (npz/txt/bag today). Slicing is entirely
210/// in-RAM, so `open()` returns a working handle for every supported format from day one.
211pub struct MemorySliceSource {
212    stream: EventStream,
213}
214
215impl MemorySliceSource {
216    pub fn new(stream: EventStream) -> Self {
217        Self { stream }
218    }
219
220    fn rebuild(&self, indices: impl Iterator<Item = usize>) -> EventStream {
221        let (width, height) = self.stream.sensor_size();
222        let mut builder = EventStreamBuilder::new(width, height, self.stream.timestamp_scale_ms());
223        let (xs, ys, ts, ps) = (
224            self.stream.xs(),
225            self.stream.ys(),
226            self.stream.ts(),
227            self.stream.ps(),
228        );
229        for index in indices {
230            builder.push(xs[index], ys[index], ts[index], ps[index]);
231        }
232        builder.build()
233    }
234}
235
236impl SliceSource for MemorySliceSource {
237    fn sensor_size(&self) -> (usize, usize) {
238        self.stream.sensor_size()
239    }
240
241    fn timestamp_scale_ms(&self) -> f64 {
242        self.stream.timestamp_scale_ms()
243    }
244
245    fn n_events(&self) -> usize {
246        self.stream.len()
247    }
248
249    fn time_span(&self) -> (i64, i64) {
250        let ts = self.stream.ts();
251        match (ts.iter().min(), ts.iter().max()) {
252            (Some(&min), Some(&max)) => (min, max),
253            _ => (0, 0),
254        }
255    }
256
257    fn slice_index(&self, i0: usize, i1: usize) -> Result<EventStream, IoError> {
258        let i0 = i0.min(self.stream.len());
259        let i1 = i1.clamp(i0, self.stream.len());
260        Ok(self.rebuild(i0..i1))
261    }
262
263    fn slice_time(&self, t0: i64, t1: i64) -> Result<EventStream, IoError> {
264        let ts = self.stream.ts();
265        Ok(self.rebuild((0..ts.len()).filter(|&index| ts[index] >= t0 && ts[index] < t1)))
266    }
267}
268
269/// A boxed [`SliceSource`] — the handle [`open`] returns.
270pub type Reader = Box<dyn SliceSource>;
271
272/// Opens a file for lazy slicing, detected by extension (the `VideoCapture` analogue to
273/// [`load`]'s `imread`). HDF5 is sliced in place by binary-searching its timestamp
274/// dataset; every other format is loaded once and sliced in memory. Same `LoadOptions`
275/// as [`load`] (`max_events` is ignored — slicing supersedes it).
276pub fn open(path: impl AsRef<Path>, options: LoadOptions) -> Result<Reader, IoError> {
277    let path = path.as_ref();
278    match detect_format(path)? {
279        Format::Hdf5 => {
280            #[cfg(feature = "hdf5")]
281            {
282                Ok(Box::new(h5::open_hdf5_slice(path, &options)?))
283            }
284            #[cfg(not(feature = "hdf5"))]
285            {
286                Err(IoError::Unsupported(
287                    "HDF5 support is not built in; rebuild with --features hdf5".to_owned(),
288                ))
289            }
290        }
291        Format::Text => Ok(Box::new(text::open_text_slice(path, &options)?)),
292        Format::Rosbag => Ok(Box::new(bag::open_bag_slice(path, &options)?)),
293        _ => Ok(Box::new(MemorySliceSource::new(load(path, options)?))),
294    }
295}
296
297/// Options for the [`save_stream`] / [`save_frame`] writers — the symmetric mirror of
298/// [`LoadOptions`]. Most formats ignore every field; readers and writers agree on the rest.
299#[derive(Clone, Debug, Default)]
300pub struct SaveOptions {
301    /// Rosbag topic to write the `dvs_msgs/EventArray` messages on (defaults to
302    /// `/davis/left/events`, matching the reader).
303    pub topic: Option<String>,
304    /// PNG frame export only: the colour map (default [`Colormap::Viridis`]).
305    pub colormap: Colormap,
306    /// PNG frame export only: auto-contrast the field to its data range. `None` = `true`.
307    pub normalize: Option<bool>,
308}
309
310/// Persists an [`EventStream`] to `path`, the format chosen by extension — the symmetric
311/// counterpart of [`load`]. npz/HDF5/rosbag round-trip exactly (metadata stored); txt
312/// stores `t x y p` and recovers sensor size / time unit on load via inference or options.
313pub fn save_stream(
314    path: impl AsRef<Path>,
315    stream: &EventStream,
316    options: &SaveOptions,
317) -> Result<(), IoError> {
318    let path = path.as_ref();
319    match detect_format(path)? {
320        Format::Npz => npz::write_npz_stream(path, stream),
321        Format::Text => text::write_text_stream(path, stream),
322        Format::Rosbag => bag::write_bag(path, stream, options.topic.as_deref()),
323        Format::Hdf5 => {
324            #[cfg(feature = "hdf5")]
325            {
326                h5::write_hdf5_stream(path, stream)
327            }
328            #[cfg(not(feature = "hdf5"))]
329            {
330                Err(IoError::Unsupported(
331                    "HDF5 support is not built in; rebuild with --features hdf5".to_owned(),
332                ))
333            }
334        }
335        other => Err(IoError::Unsupported(format!(
336            "saving an event stream as {other:?} is not supported"
337        ))),
338    }
339}
340
341/// Persists an [`EventFrame`] (a computed representation) to `path`, preserving its shape,
342/// dtype, `kind`, and `channel_names`. Supported: npz (default build) and HDF5.
343pub fn save_frame(
344    path: impl AsRef<Path>,
345    frame: &EventFrame,
346    options: &SaveOptions,
347) -> Result<(), IoError> {
348    let path = path.as_ref();
349    match detect_format(path)? {
350        Format::Npz => npz::write_npz_frame(path, frame),
351        Format::Hdf5 => {
352            #[cfg(feature = "hdf5")]
353            {
354                h5::write_hdf5_frame(path, frame)
355            }
356            #[cfg(not(feature = "hdf5"))]
357            {
358                Err(IoError::Unsupported(
359                    "HDF5 support is not built in; rebuild with --features hdf5".to_owned(),
360                ))
361            }
362        }
363        Format::Png => write_png_frame(path, frame, options),
364        other => Err(IoError::Unsupported(format!(
365            "saving an event frame as {other:?} is not supported"
366        ))),
367    }
368}
369
370/// Renders a frame through [`render_frame`] (colormapped 2-D view) and encodes it as an
371/// 8-bit RGB PNG. Unlike npz/HDF5 this is a *view*, not a round-trippable dump.
372fn write_png_frame(path: &Path, frame: &EventFrame, options: &SaveOptions) -> Result<(), IoError> {
373    let image = render_frame(frame, options.colormap, options.normalize.unwrap_or(true));
374    let file = std::fs::File::create(path)?;
375    let writer = std::io::BufWriter::new(file);
376    let mut encoder = png::Encoder::new(writer, image.width as u32, image.height as u32);
377    encoder.set_color(png::ColorType::Rgb);
378    encoder.set_depth(png::BitDepth::Eight);
379    encoder
380        .write_header()
381        .and_then(|mut writer| writer.write_image_data(&image.pixels))
382        .map_err(|error| match error {
383            png::EncodingError::IoError(error) => IoError::Io(error),
384            other => IoError::Format(other.to_string()),
385        })
386}
387
388/// Reads an [`EventFrame`] previously written by [`save_frame`], reconstructing its dtype,
389/// `kind`, and `channel_names`. Supported: npz (default build) and HDF5.
390pub fn load_frame(path: impl AsRef<Path>) -> Result<EventFrame, IoError> {
391    let path = path.as_ref();
392    match detect_format(path)? {
393        Format::Npz => npz::read_npz_frame(path),
394        Format::Hdf5 => {
395            #[cfg(feature = "hdf5")]
396            {
397                h5::read_hdf5_frame(path)
398            }
399            #[cfg(not(feature = "hdf5"))]
400            {
401                Err(IoError::Unsupported(
402                    "HDF5 support is not built in; rebuild with --features hdf5".to_owned(),
403                ))
404            }
405        }
406        other => Err(IoError::Unsupported(format!(
407            "loading an event frame from {other:?} is not supported"
408        ))),
409    }
410}
411
412#[derive(Debug)]
413pub enum IoError {
414    Io(io::Error),
415    Parse { line: usize, message: String },
416    Format(String),
417    InvalidSensorSize,
418    Unsupported(String),
419}
420
421impl fmt::Display for IoError {
422    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
423        match self {
424            Self::Io(error) => error.fmt(formatter),
425            Self::Parse { line, message } => write!(formatter, "line {line}: {message}"),
426            Self::Format(message) => formatter.write_str(message),
427            Self::InvalidSensorSize => {
428                formatter.write_str("sensor width and height must be positive")
429            }
430            Self::Unsupported(message) => formatter.write_str(message),
431        }
432    }
433}
434
435impl Error for IoError {
436    fn source(&self) -> Option<&(dyn Error + 'static)> {
437        match self {
438            Self::Io(error) => Some(error),
439            _ => None,
440        }
441    }
442}
443
444impl From<io::Error> for IoError {
445    fn from(error: io::Error) -> Self {
446        Self::Io(error)
447    }
448}
449
450#[cfg(test)]
451mod tests {
452    use super::{load, open, IoError, LoadOptions, MemorySliceSource, SliceSource};
453    use crate::EventStreamBuilder;
454
455    #[test]
456    fn unknown_extension_is_unsupported() {
457        let error = load("recording.mp4", LoadOptions::default()).unwrap_err();
458        assert!(matches!(error, IoError::Unsupported(_)));
459    }
460
461    #[test]
462    fn hdf5_extension_dispatches_per_feature() {
463        // `.h5` is always recognised; with the feature it reaches the reader (a missing
464        // file is then an IO error), without it the dispatch reports missing support.
465        let error = load("recording.h5", LoadOptions::default()).unwrap_err();
466        #[cfg(feature = "hdf5")]
467        assert!(matches!(error, IoError::Io(_)));
468        #[cfg(not(feature = "hdf5"))]
469        match error {
470            IoError::Unsupported(message) => assert!(message.contains("HDF5")),
471            other => panic!("expected unsupported error, got {other:?}"),
472        }
473    }
474
475    #[test]
476    fn text_missing_file_is_reported() {
477        // Text no longer requires sensor_size (it infers); a missing file is an IO error.
478        let error = load("events.txt", LoadOptions::default()).unwrap_err();
479        assert!(matches!(error, IoError::Io(_)));
480    }
481
482    #[test]
483    fn aedat_dispatches_to_the_reader() {
484        // `.aedat` reaches the AEDAT 2.0 reader, so a missing file is an IO error.
485        let error = load("recording.aedat", LoadOptions::default()).unwrap_err();
486        assert!(matches!(error, IoError::Io(_)));
487    }
488
489    #[test]
490    fn prophesee_dat_dispatches_to_the_reader() {
491        let error = load("recording.dat", LoadOptions::default()).unwrap_err();
492        assert!(matches!(error, IoError::Io(_)));
493    }
494
495    #[test]
496    fn aedat4_and_prophesee_raw_are_unsupported() {
497        // These extensions are recognised but their formats are not implemented yet.
498        for path in ["recording.aedat4", "recording.raw"] {
499            match load(path, LoadOptions::default()) {
500                Err(IoError::Unsupported(_)) => {}
501                other => panic!("expected unsupported for {path}, got {other:?}"),
502            }
503        }
504    }
505
506    fn sample_source() -> MemorySliceSource {
507        let mut builder = EventStreamBuilder::new(4, 4, 0.001);
508        builder.push(0, 0, 0, true);
509        builder.push(1, 1, 10, false);
510        builder.push(2, 2, 20, true);
511        builder.push(0, 1, 30, false);
512        MemorySliceSource::new(builder.build())
513    }
514
515    #[test]
516    fn memory_source_reports_span_and_count() {
517        let source = sample_source();
518        assert_eq!(source.n_events(), 4);
519        assert_eq!(source.time_span(), (0, 30));
520    }
521
522    #[test]
523    fn memory_source_slices_by_time_and_index() {
524        let source = sample_source();
525
526        // Half-open [10, 30) keeps t = 10 and 20.
527        assert_eq!(source.slice_time(10, 30).unwrap().ts(), &[10, 20]);
528        assert_eq!(source.slice_index(1, 3).unwrap().ts(), &[10, 20]);
529        assert_eq!(source.slice_index(2, 100).unwrap().len(), 2); // hi clamped to len
530        assert!(source.slice_time(100, 200).unwrap().is_empty());
531    }
532
533    #[test]
534    fn open_rejects_unknown_extension() {
535        // `Reader` is a trait object (not `Debug`), so match rather than `unwrap_err`.
536        match open("recording.mp4", LoadOptions::default()) {
537            Err(IoError::Unsupported(_)) => {}
538            Err(other) => panic!("expected unsupported error, got {other:?}"),
539            Ok(_) => panic!("expected an error for an unknown extension"),
540        }
541    }
542}