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#[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
39pub trait EventSource {
42 fn sensor_size(&self) -> (usize, usize);
43 fn timestamp_scale_ms(&self) -> f64;
44 fn next_event(&mut self) -> Result<Option<RawEvent>, IoError>;
46}
47
48pub fn read_all(source: impl EventSource) -> Result<EventStream, IoError> {
50 read_capped(source, None)
51}
52
53pub 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#[derive(Clone, Debug, Default)]
72pub struct LoadOptions {
73 pub sensor_size: Option<(usize, usize)>,
76 pub time_unit: Option<TimeUnit>,
80 pub order: ColumnOrder,
82 pub topic: Option<String>,
84 pub max_events: Option<usize>,
86 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
128pub 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 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
143fn 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
192pub trait SliceSource: Send {
196 fn sensor_size(&self) -> (usize, usize);
197 fn timestamp_scale_ms(&self) -> f64;
198 fn n_events(&self) -> usize;
200 fn time_span(&self) -> (i64, i64);
202 fn slice_index(&self, i0: usize, i1: usize) -> Result<EventStream, IoError>;
204 fn slice_time(&self, t0: i64, t1: i64) -> Result<EventStream, IoError>;
206}
207
208pub 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
269pub type Reader = Box<dyn SliceSource>;
271
272pub 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#[derive(Clone, Debug, Default)]
300pub struct SaveOptions {
301 pub topic: Option<String>,
304 pub colormap: Colormap,
306 pub normalize: Option<bool>,
308}
309
310pub 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
341pub 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
370fn 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
388pub 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 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 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 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 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 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); assert!(source.slice_time(100, 200).unwrap().is_empty());
531 }
532
533 #[test]
534 fn open_rejects_unknown_extension() {
535 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}