use crate::object_classifier::ObjectClass;
use crate::sdf_compress::{svo_diff_hash, CompressedSdf, SvoChunkData};
use libasp::scene::{SdfSceneDelta, SdfSceneDescriptor};
pub struct EdgeStreamEncoder {
scene_version: u32,
last_svo_hash: u64,
frame_count: u64,
keyframe_interval: u32,
}
#[derive(Debug)]
pub enum AspEdgePacket {
Keyframe {
scene: SdfSceneDescriptor,
frame_id: u64,
},
Delta { delta: SdfSceneDelta, frame_id: u64 },
Skip { frame_id: u64 },
}
impl EdgeStreamEncoder {
pub fn new(keyframe_interval: u32) -> Self {
Self {
scene_version: 0,
last_svo_hash: 0,
frame_count: 0,
keyframe_interval,
}
}
pub fn encode_frame(
&mut self,
compressed: &CompressedSdf,
classifications: &[(u8, ObjectClass)],
) -> AspEdgePacket {
self.frame_count += 1;
let is_keyframe =
self.frame_count == 1 || (self.frame_count % self.keyframe_interval as u64) == 0;
if is_keyframe {
self.encode_keyframe(compressed, classifications)
} else {
self.encode_delta(compressed)
}
}
fn encode_keyframe(
&mut self,
compressed: &CompressedSdf,
classifications: &[(u8, ObjectClass)],
) -> AspEdgePacket {
self.scene_version += 1;
let asdf_data = extract_asdf_data(compressed);
let (bounds_min, bounds_max) = extract_bounds(compressed);
let labels: Option<Vec<(u8, String)>> = if !classifications.is_empty() {
Some(
classifications
.iter()
.map(|(id, class)| (*id, class.label().to_string()))
.collect(),
)
} else {
None
};
let mut scene = SdfSceneDescriptor::new(asdf_data)
.with_bounds(bounds_min, bounds_max)
.with_render_resolution(0);
scene.scene_version = self.scene_version;
scene.classification_labels = labels;
self.last_svo_hash = compute_compressed_hash(compressed);
AspEdgePacket::Keyframe {
scene,
frame_id: self.frame_count,
}
}
fn encode_delta(&mut self, compressed: &CompressedSdf) -> AspEdgePacket {
let new_hash = compute_compressed_hash(compressed);
if new_hash == self.last_svo_hash {
return AspEdgePacket::Skip {
frame_id: self.frame_count,
};
}
self.last_svo_hash = new_hash;
let ref_version = self.scene_version;
self.scene_version += 1;
let delta_data = extract_delta_data(compressed);
let delta = SdfSceneDelta::svo_chunk_delta(ref_version, delta_data);
AspEdgePacket::Delta {
delta,
frame_id: self.frame_count,
}
}
pub fn scene_version(&self) -> u32 {
self.scene_version
}
pub fn frame_count(&self) -> u64 {
self.frame_count
}
}
fn extract_asdf_data(compressed: &CompressedSdf) -> Vec<u8> {
match compressed {
CompressedSdf::Primitives { asdf_data, .. } => asdf_data.clone(),
CompressedSdf::SvoChunks { chunks, .. } => {
let total_len: usize = chunks.iter().map(|c| c.data.len()).sum();
let mut data = Vec::with_capacity(total_len);
for chunk in chunks {
data.extend_from_slice(&chunk.data);
}
data
}
CompressedSdf::Hybrid {
asdf_data,
svo_chunks,
..
} => {
let total_len =
asdf_data.len() + svo_chunks.iter().map(|c| c.data.len()).sum::<usize>();
let mut data = Vec::with_capacity(total_len);
data.extend_from_slice(asdf_data);
for chunk in svo_chunks {
data.extend_from_slice(&chunk.data);
}
data
}
}
}
fn extract_bounds(compressed: &CompressedSdf) -> ([f32; 3], [f32; 3]) {
match compressed {
CompressedSdf::SvoChunks { chunks, .. }
| CompressedSdf::Hybrid {
svo_chunks: chunks, ..
} => {
if let Some(chunk) = chunks.first() {
(chunk.bounds_min, chunk.bounds_max)
} else {
([-10.0; 3], [10.0; 3])
}
}
_ => ([-10.0; 3], [10.0; 3]),
}
}
fn compute_compressed_hash(compressed: &CompressedSdf) -> u64 {
match compressed {
CompressedSdf::Primitives { asdf_data, .. } => svo_diff_hash(asdf_data),
CompressedSdf::SvoChunks { chunks, .. } => {
let mut combined_hash = 0u64;
for chunk in chunks {
combined_hash ^= svo_diff_hash(&chunk.data);
}
combined_hash
}
CompressedSdf::Hybrid {
asdf_data,
svo_chunks,
..
} => {
let mut combined = svo_diff_hash(asdf_data);
for chunk in svo_chunks {
combined ^= svo_diff_hash(&chunk.data);
}
combined
}
}
}
fn extract_delta_data(compressed: &CompressedSdf) -> Vec<u8> {
extract_asdf_data(compressed)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_encoder_first_frame_is_keyframe() {
let mut encoder = EdgeStreamEncoder::new(30);
let compressed = CompressedSdf::SvoChunks {
chunks: vec![SvoChunkData {
chunk_id: 0,
data: vec![1, 2, 3, 4],
node_count: 1,
bounds_min: [-1.0, -1.0, -1.0],
bounds_max: [1.0, 1.0, 1.0],
}],
total_nodes: 1,
};
let packet = encoder.encode_frame(&compressed, &[]);
assert!(matches!(packet, AspEdgePacket::Keyframe { .. }));
}
#[test]
fn test_encoder_skip_on_no_change() {
let mut encoder = EdgeStreamEncoder::new(100);
let compressed = CompressedSdf::SvoChunks {
chunks: vec![SvoChunkData {
chunk_id: 0,
data: vec![1, 2, 3, 4],
node_count: 1,
bounds_min: [-1.0; 3],
bounds_max: [1.0; 3],
}],
total_nodes: 1,
};
let _ = encoder.encode_frame(&compressed, &[]);
let packet = encoder.encode_frame(&compressed, &[]);
assert!(matches!(packet, AspEdgePacket::Skip { .. }));
}
#[test]
fn test_encoder_delta_on_change() {
let mut encoder = EdgeStreamEncoder::new(100);
let compressed1 = CompressedSdf::SvoChunks {
chunks: vec![SvoChunkData {
chunk_id: 0,
data: vec![1, 2, 3, 4],
node_count: 1,
bounds_min: [-1.0; 3],
bounds_max: [1.0; 3],
}],
total_nodes: 1,
};
let compressed2 = CompressedSdf::SvoChunks {
chunks: vec![SvoChunkData {
chunk_id: 0,
data: vec![5, 6, 7, 8], node_count: 1,
bounds_min: [-1.0; 3],
bounds_max: [1.0; 3],
}],
total_nodes: 1,
};
let _ = encoder.encode_frame(&compressed1, &[]);
let packet = encoder.encode_frame(&compressed2, &[]);
assert!(matches!(packet, AspEdgePacket::Delta { .. }));
}
#[test]
fn test_encoder_periodic_keyframe() {
let mut encoder = EdgeStreamEncoder::new(3);
let compressed = CompressedSdf::Primitives {
primitives: vec![],
asdf_data: vec![1, 2, 3],
};
let _ = encoder.encode_frame(&compressed, &[]); let _ = encoder.encode_frame(&compressed, &[]); let packet = encoder.encode_frame(&compressed, &[]); assert!(matches!(packet, AspEdgePacket::Keyframe { .. }));
}
#[test]
fn test_scene_version_increments() {
let mut encoder = EdgeStreamEncoder::new(100);
assert_eq!(encoder.scene_version(), 0);
let compressed = CompressedSdf::Primitives {
primitives: vec![],
asdf_data: vec![1],
};
let _ = encoder.encode_frame(&compressed, &[]);
assert_eq!(encoder.scene_version(), 1);
}
#[test]
fn test_frame_count_increments() {
let mut encoder = EdgeStreamEncoder::new(100);
assert_eq!(encoder.frame_count(), 0);
let compressed = CompressedSdf::Primitives {
primitives: vec![],
asdf_data: vec![],
};
let _ = encoder.encode_frame(&compressed, &[]);
let _ = encoder.encode_frame(&compressed, &[]);
assert_eq!(encoder.frame_count(), 2);
}
#[test]
fn test_extract_bounds_from_svo() {
let bounds = extract_bounds(&CompressedSdf::SvoChunks {
chunks: vec![SvoChunkData {
chunk_id: 0,
data: vec![],
node_count: 1,
bounds_min: [-5.0, -3.0, -1.0],
bounds_max: [5.0, 3.0, 1.0],
}],
total_nodes: 1,
});
assert_eq!(bounds.0, [-5.0, -3.0, -1.0]);
assert_eq!(bounds.1, [5.0, 3.0, 1.0]);
}
#[test]
fn test_extract_bounds_default_for_primitives() {
let bounds = extract_bounds(&CompressedSdf::Primitives {
primitives: vec![],
asdf_data: vec![],
});
assert_eq!(bounds.0, [-10.0; 3]);
assert_eq!(bounds.1, [10.0; 3]);
}
}