spatialrust_episode/
augment.rs1use spatialrust_sync::MemoryEpisode;
4
5use crate::{Episode, EpisodeResult};
6
7#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
9pub enum AugmentationOp {
10 TemporalSubsample {
12 stride: usize,
14 },
15 Reverse,
17}
18
19#[derive(Clone, Copy, Debug, Default)]
21pub struct EpisodeAugmentor;
22
23impl EpisodeAugmentor {
24 pub fn apply(&self, episode: &Episode, op: AugmentationOp) -> EpisodeResult<Episode> {
26 let records = episode.memory.records();
27 let next = match op {
28 AugmentationOp::TemporalSubsample { stride } => {
29 if stride == 0 {
30 return Err(crate::EpisodeError::InvalidConfiguration(
31 "stride must be positive".into(),
32 ));
33 }
34 records.iter().step_by(stride).cloned().collect::<Vec<_>>()
35 }
36 AugmentationOp::Reverse => records.iter().rev().cloned().collect::<Vec<_>>(),
37 };
38 let memory = MemoryEpisode::from_records(next);
39 let mut out = Episode::try_new(format!("{}::aug", episode.id.0), memory)?;
40 out.annotations = episode.annotations.clone();
41 out.provenance = episode.provenance.clone();
42 Ok(out)
43 }
44}
45
46#[cfg(test)]
47mod tests {
48 use super::{AugmentationOp, EpisodeAugmentor};
49 use crate::Episode;
50 use spatialrust_core::{
51 PointBuffer, PointBufferSet, PointCloud, SpatialMetadata, StandardSchemas, Timestamp,
52 };
53 use spatialrust_records::{SchemaVersion, SpatialRecord};
54 use spatialrust_sync::{ClockDomain, MemoryEpisode, StampedRecord, StampedTime};
55
56 fn sample(nanos: u64) -> StampedRecord {
57 let mut buffers = PointBufferSet::new();
58 buffers.insert("x", PointBuffer::from_f32(vec![0.0]));
59 buffers.insert("y", PointBuffer::from_f32(vec![0.0]));
60 buffers.insert("z", PointBuffer::from_f32(vec![0.0]));
61 let cloud = PointCloud::try_from_parts(
62 StandardSchemas::point_xyz(),
63 buffers,
64 SpatialMetadata::default(),
65 )
66 .unwrap();
67 let record =
68 SpatialRecord::try_from_cloud("point", SchemaVersion::new(1, 0), cloud).unwrap();
69 StampedRecord::new(
70 "lidar",
71 StampedTime::exact("host", ClockDomain::HostSteady, Timestamp::from_nanos(nanos)),
72 record,
73 )
74 }
75
76 #[test]
77 fn subsamples_episode() {
78 let memory = MemoryEpisode::from_records(vec![sample(1), sample(2), sample(3), sample(4)]);
79 let episode = Episode::try_new("ep0", memory).unwrap();
80 let out = EpisodeAugmentor
81 .apply(&episode, AugmentationOp::TemporalSubsample { stride: 2 })
82 .unwrap();
83 assert_eq!(out.memory.records().len(), 2);
84 }
85}