bevy_sensor/
fixtures.rs

1//! Test fixtures for pre-rendered YCB images
2//!
3//! This module provides utilities for loading pre-rendered images from disk,
4//! enabling testing without GPU access.
5//!
6//! # Usage
7//!
8//! ```ignore
9//! use bevy_sensor::fixtures::TestFixtures;
10//!
11//! let fixtures = TestFixtures::load("test_fixtures/renders")?;
12//!
13//! // Get a specific render
14//! let render = fixtures.get_render("003_cracker_box", 0, 5)?;
15//! let rgb_image = render.to_rgb_image();
16//! let depth_image = render.to_depth_image();
17//! ```
18
19use crate::{CameraIntrinsics, RenderOutput};
20use bevy::prelude::*;
21use serde::{Deserialize, Serialize};
22use std::collections::HashMap;
23use std::fs;
24use std::path::{Path, PathBuf};
25
26/// Error type for fixture loading
27#[derive(Debug)]
28pub enum FixtureError {
29    /// Directory not found
30    NotFound(String),
31    /// Metadata file missing or invalid
32    InvalidMetadata(String),
33    /// Render file missing
34    RenderNotFound {
35        object_id: String,
36        rotation: usize,
37        viewpoint: usize,
38    },
39    /// IO error
40    IoError(std::io::Error),
41    /// JSON parsing error
42    JsonError(serde_json::Error),
43}
44
45impl std::fmt::Display for FixtureError {
46    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47        match self {
48            FixtureError::NotFound(path) => write!(f, "Fixture directory not found: {}", path),
49            FixtureError::InvalidMetadata(msg) => write!(f, "Invalid metadata: {}", msg),
50            FixtureError::RenderNotFound {
51                object_id,
52                rotation,
53                viewpoint,
54            } => write!(
55                f,
56                "Render not found: {} r{} v{}",
57                object_id, rotation, viewpoint
58            ),
59            FixtureError::IoError(e) => write!(f, "IO error: {}", e),
60            FixtureError::JsonError(e) => write!(f, "JSON error: {}", e),
61        }
62    }
63}
64
65impl std::error::Error for FixtureError {}
66
67impl From<std::io::Error> for FixtureError {
68    fn from(e: std::io::Error) -> Self {
69        FixtureError::IoError(e)
70    }
71}
72
73impl From<serde_json::Error> for FixtureError {
74    fn from(e: serde_json::Error) -> Self {
75        FixtureError::JsonError(e)
76    }
77}
78
79/// Dataset metadata from pre-rendering
80#[derive(Debug, Clone, Serialize, Deserialize)]
81pub struct DatasetMetadata {
82    pub version: String,
83    pub objects: Vec<String>,
84    pub viewpoints_per_rotation: usize,
85    pub rotations_per_object: usize,
86    pub renders_per_object: usize,
87    pub resolution: [u32; 2],
88    pub intrinsics: IntrinsicsMetadata,
89    pub viewpoint_config: ViewpointConfigMetadata,
90    pub rotations: Vec<[f32; 3]>,
91}
92
93#[derive(Debug, Clone, Serialize, Deserialize)]
94pub struct IntrinsicsMetadata {
95    pub focal_length: [f32; 2],
96    pub principal_point: [f32; 2],
97    pub image_size: [u32; 2],
98}
99
100#[derive(Debug, Clone, Serialize, Deserialize)]
101pub struct ViewpointConfigMetadata {
102    pub radius: f32,
103    pub yaw_count: usize,
104    pub pitch_angles_deg: Vec<f32>,
105}
106
107/// Metadata for a single render
108#[derive(Debug, Clone, Serialize, Deserialize)]
109pub struct RenderMetadata {
110    pub object_id: String,
111    pub rotation_index: usize,
112    pub viewpoint_index: usize,
113    pub rotation_euler: [f32; 3],
114    pub camera_position: [f32; 3],
115    pub rgba_file: String,
116    pub depth_file: String,
117}
118
119/// Pre-rendered test fixtures loaded from disk
120pub struct TestFixtures {
121    /// Root directory containing fixtures
122    root: PathBuf,
123    /// Dataset metadata
124    pub metadata: DatasetMetadata,
125    /// Per-object render indices
126    indices: HashMap<String, Vec<RenderMetadata>>,
127}
128
129impl TestFixtures {
130    /// Load test fixtures from a directory
131    ///
132    /// # Arguments
133    /// * `path` - Path to the fixtures directory (e.g., "test_fixtures/renders")
134    ///
135    /// # Returns
136    /// * `Ok(TestFixtures)` if loaded successfully
137    /// * `Err(FixtureError)` if loading fails
138    pub fn load<P: AsRef<Path>>(path: P) -> Result<Self, FixtureError> {
139        let root = path.as_ref().to_path_buf();
140
141        if !root.exists() {
142            return Err(FixtureError::NotFound(root.display().to_string()));
143        }
144
145        // Load metadata
146        let metadata_path = root.join("metadata.json");
147        if !metadata_path.exists() {
148            return Err(FixtureError::InvalidMetadata(
149                "metadata.json not found".to_string(),
150            ));
151        }
152
153        let metadata_json = fs::read_to_string(&metadata_path)?;
154        let metadata: DatasetMetadata = serde_json::from_str(&metadata_json)?;
155
156        // Load per-object indices
157        let mut indices = HashMap::new();
158        for object_id in &metadata.objects {
159            let index_path = root.join(object_id).join("index.json");
160            if index_path.exists() {
161                let index_json = fs::read_to_string(&index_path)?;
162                let renders: Vec<RenderMetadata> = serde_json::from_str(&index_json)?;
163                indices.insert(object_id.clone(), renders);
164            }
165        }
166
167        Ok(Self {
168            root,
169            metadata,
170            indices,
171        })
172    }
173
174    /// Check if fixtures exist at the given path
175    pub fn exists<P: AsRef<Path>>(path: P) -> bool {
176        let root = path.as_ref();
177        root.exists() && root.join("metadata.json").exists()
178    }
179
180    /// Get list of available objects
181    pub fn objects(&self) -> &[String] {
182        &self.metadata.objects
183    }
184
185    /// Get number of viewpoints per rotation
186    pub fn viewpoints_per_rotation(&self) -> usize {
187        self.metadata.viewpoints_per_rotation
188    }
189
190    /// Get number of rotations per object
191    pub fn rotations_per_object(&self) -> usize {
192        self.metadata.rotations_per_object
193    }
194
195    /// Get total renders available for an object
196    pub fn renders_for_object(&self, object_id: &str) -> usize {
197        self.indices.get(object_id).map(|v| v.len()).unwrap_or(0)
198    }
199
200    /// Get camera intrinsics (converts from f32 metadata to f64 for TBP precision)
201    pub fn intrinsics(&self) -> CameraIntrinsics {
202        CameraIntrinsics {
203            focal_length: [
204                self.metadata.intrinsics.focal_length[0] as f64,
205                self.metadata.intrinsics.focal_length[1] as f64,
206            ],
207            principal_point: [
208                self.metadata.intrinsics.principal_point[0] as f64,
209                self.metadata.intrinsics.principal_point[1] as f64,
210            ],
211            image_size: self.metadata.intrinsics.image_size,
212        }
213    }
214
215    /// Load a specific render by object, rotation index, and viewpoint index
216    ///
217    /// # Arguments
218    /// * `object_id` - YCB object ID (e.g., "003_cracker_box")
219    /// * `rotation_idx` - Rotation index (0-2 for benchmark rotations)
220    /// * `viewpoint_idx` - Viewpoint index (0-23 for default config)
221    pub fn get_render(
222        &self,
223        object_id: &str,
224        rotation_idx: usize,
225        viewpoint_idx: usize,
226    ) -> Result<RenderOutput, FixtureError> {
227        // Find the render metadata
228        let renders = self
229            .indices
230            .get(object_id)
231            .ok_or_else(|| FixtureError::RenderNotFound {
232                object_id: object_id.to_string(),
233                rotation: rotation_idx,
234                viewpoint: viewpoint_idx,
235            })?;
236
237        let render_meta = renders
238            .iter()
239            .find(|r| r.rotation_index == rotation_idx && r.viewpoint_index == viewpoint_idx)
240            .ok_or_else(|| FixtureError::RenderNotFound {
241                object_id: object_id.to_string(),
242                rotation: rotation_idx,
243                viewpoint: viewpoint_idx,
244            })?;
245
246        // Load RGBA from PNG
247        let rgba_path = self.root.join(object_id).join(&render_meta.rgba_file);
248        let rgba = load_rgba_png(&rgba_path)?;
249
250        // Load depth from binary
251        let depth_path = self.root.join(object_id).join(&render_meta.depth_file);
252        let depth = load_depth_binary(&depth_path)?;
253
254        // Build camera transform from position (looking at origin)
255        let pos = render_meta.camera_position;
256        let camera_transform =
257            Transform::from_xyz(pos[0], pos[1], pos[2]).looking_at(Vec3::ZERO, Vec3::Y);
258
259        // Build object rotation (convert from f32 metadata to f64)
260        let rot = render_meta.rotation_euler;
261        let object_rotation =
262            crate::ObjectRotation::new(rot[0] as f64, rot[1] as f64, rot[2] as f64);
263
264        Ok(RenderOutput {
265            rgba,
266            depth,
267            width: self.metadata.resolution[0],
268            height: self.metadata.resolution[1],
269            intrinsics: self.intrinsics(),
270            camera_transform,
271            object_rotation,
272        })
273    }
274
275    /// Load all renders for an object
276    pub fn get_all_renders(&self, object_id: &str) -> Result<Vec<RenderOutput>, FixtureError> {
277        let renders = self
278            .indices
279            .get(object_id)
280            .ok_or_else(|| FixtureError::RenderNotFound {
281                object_id: object_id.to_string(),
282                rotation: 0,
283                viewpoint: 0,
284            })?;
285
286        let mut outputs = Vec::with_capacity(renders.len());
287        for meta in renders {
288            let output = self.get_render(object_id, meta.rotation_index, meta.viewpoint_index)?;
289            outputs.push(output);
290        }
291
292        Ok(outputs)
293    }
294
295    /// Iterate over all renders for an object
296    pub fn iter_renders<'a>(
297        &'a self,
298        object_id: &'a str,
299    ) -> impl Iterator<Item = Result<(usize, usize, RenderOutput), FixtureError>> + 'a {
300        let renders = self.indices.get(object_id);
301
302        renders.into_iter().flat_map(|v| v.iter()).map(move |meta| {
303            let output = self.get_render(object_id, meta.rotation_index, meta.viewpoint_index)?;
304            Ok((meta.rotation_index, meta.viewpoint_index, output))
305        })
306    }
307}
308
309/// Load RGBA data from a PNG file
310fn load_rgba_png(path: &Path) -> Result<Vec<u8>, FixtureError> {
311    let img = image::open(path).map_err(|e| FixtureError::IoError(std::io::Error::other(e)))?;
312
313    let rgba = img.to_rgba8();
314    Ok(rgba.into_raw())
315}
316
317/// Load depth data from binary f32 file and convert to f64 for TBP precision
318fn load_depth_binary(path: &Path) -> Result<Vec<f64>, FixtureError> {
319    let bytes = fs::read(path)?;
320
321    // Convert from little-endian bytes to f32, then to f64 for TBP precision
322    let depth: Vec<f64> = bytes
323        .chunks_exact(4)
324        .map(|chunk| {
325            let arr: [u8; 4] = chunk.try_into().unwrap();
326            f32::from_le_bytes(arr) as f64
327        })
328        .collect();
329
330    Ok(depth)
331}
332
333#[cfg(test)]
334mod tests {
335    use super::*;
336    use tempfile::TempDir;
337
338    #[test]
339    fn test_fixture_not_found() {
340        let result = TestFixtures::load("/nonexistent/path");
341        assert!(matches!(result, Err(FixtureError::NotFound(_))));
342    }
343
344    #[test]
345    fn test_fixtures_exists() {
346        assert!(!TestFixtures::exists("/nonexistent/path"));
347    }
348
349    #[test]
350    fn test_fixture_error_display() {
351        let errors = vec![
352            FixtureError::NotFound("/path".to_string()),
353            FixtureError::InvalidMetadata("bad json".to_string()),
354            FixtureError::RenderNotFound {
355                object_id: "obj".to_string(),
356                rotation: 0,
357                viewpoint: 5,
358            },
359            FixtureError::IoError(std::io::Error::new(
360                std::io::ErrorKind::NotFound,
361                "file not found",
362            )),
363            FixtureError::JsonError(serde_json::from_str::<String>("invalid").unwrap_err()),
364        ];
365
366        for err in errors {
367            let msg = err.to_string();
368            assert!(!msg.is_empty());
369        }
370    }
371
372    #[test]
373    fn test_fixture_missing_metadata() {
374        let temp_dir = TempDir::new().unwrap();
375        let result = TestFixtures::load(temp_dir.path());
376        assert!(matches!(result, Err(FixtureError::InvalidMetadata(_))));
377    }
378
379    #[test]
380    fn test_fixture_load_metadata() {
381        let temp_dir = TempDir::new().unwrap();
382
383        // Create minimal metadata
384        let metadata = DatasetMetadata {
385            version: "1.0".to_string(),
386            objects: vec!["test_object".to_string()],
387            viewpoints_per_rotation: 24,
388            rotations_per_object: 3,
389            renders_per_object: 72,
390            resolution: [64, 64],
391            intrinsics: IntrinsicsMetadata {
392                focal_length: [55.4, 55.4],
393                principal_point: [32.0, 32.0],
394                image_size: [64, 64],
395            },
396            viewpoint_config: ViewpointConfigMetadata {
397                radius: 0.5,
398                yaw_count: 8,
399                pitch_angles_deg: vec![-30.0, 0.0, 30.0],
400            },
401            rotations: vec![[0.0, 0.0, 0.0], [0.0, 90.0, 0.0], [0.0, 180.0, 0.0]],
402        };
403
404        let metadata_json = serde_json::to_string_pretty(&metadata).unwrap();
405        let metadata_path = temp_dir.path().join("metadata.json");
406        fs::write(&metadata_path, &metadata_json).unwrap();
407
408        // Create object directory with empty index
409        let obj_dir = temp_dir.path().join("test_object");
410        fs::create_dir_all(&obj_dir).unwrap();
411        fs::write(obj_dir.join("index.json"), "[]").unwrap();
412
413        // Load fixtures
414        let fixtures = TestFixtures::load(temp_dir.path()).unwrap();
415
416        assert_eq!(fixtures.objects(), &["test_object"]);
417        assert_eq!(fixtures.viewpoints_per_rotation(), 24);
418        assert_eq!(fixtures.rotations_per_object(), 3);
419        assert_eq!(fixtures.renders_for_object("test_object"), 0);
420        assert_eq!(fixtures.renders_for_object("nonexistent"), 0);
421
422        let intrinsics = fixtures.intrinsics();
423        assert_eq!(intrinsics.image_size, [64, 64]);
424    }
425
426    #[test]
427    fn test_load_depth_binary() {
428        let temp_dir = TempDir::new().unwrap();
429        let depth_path = temp_dir.path().join("test.depth");
430
431        // Write test depth values
432        let depths: Vec<f32> = vec![0.5, 1.0, 2.0, 10.0];
433        let bytes: Vec<u8> = depths.iter().flat_map(|f| f.to_le_bytes()).collect();
434        fs::write(&depth_path, &bytes).unwrap();
435
436        // Load and verify
437        let loaded = load_depth_binary(&depth_path).unwrap();
438        assert_eq!(loaded.len(), 4);
439        assert!((loaded[0] - 0.5).abs() < 0.001);
440        assert!((loaded[1] - 1.0).abs() < 0.001);
441        assert!((loaded[2] - 2.0).abs() < 0.001);
442        assert!((loaded[3] - 10.0).abs() < 0.001);
443    }
444
445    #[test]
446    fn test_metadata_serialization_roundtrip() {
447        let metadata = DatasetMetadata {
448            version: "1.0".to_string(),
449            objects: vec!["obj1".to_string(), "obj2".to_string()],
450            viewpoints_per_rotation: 24,
451            rotations_per_object: 3,
452            renders_per_object: 72,
453            resolution: [64, 64],
454            intrinsics: IntrinsicsMetadata {
455                focal_length: [55.4, 55.4],
456                principal_point: [32.0, 32.0],
457                image_size: [64, 64],
458            },
459            viewpoint_config: ViewpointConfigMetadata {
460                radius: 0.5,
461                yaw_count: 8,
462                pitch_angles_deg: vec![-30.0, 0.0, 30.0],
463            },
464            rotations: vec![[0.0, 0.0, 0.0]],
465        };
466
467        let json = serde_json::to_string(&metadata).unwrap();
468        let loaded: DatasetMetadata = serde_json::from_str(&json).unwrap();
469
470        assert_eq!(loaded.version, metadata.version);
471        assert_eq!(loaded.objects, metadata.objects);
472        assert_eq!(loaded.resolution, metadata.resolution);
473    }
474
475    #[test]
476    fn test_render_metadata_serialization() {
477        let meta = RenderMetadata {
478            object_id: "003_cracker_box".to_string(),
479            rotation_index: 1,
480            viewpoint_index: 5,
481            rotation_euler: [0.0, 90.0, 0.0],
482            camera_position: [0.5, 0.0, 0.0],
483            rgba_file: "r1_v05.png".to_string(),
484            depth_file: "r1_v05.depth".to_string(),
485        };
486
487        let json = serde_json::to_string(&meta).unwrap();
488        let loaded: RenderMetadata = serde_json::from_str(&json).unwrap();
489
490        assert_eq!(loaded.object_id, meta.object_id);
491        assert_eq!(loaded.rotation_index, meta.rotation_index);
492        assert_eq!(loaded.viewpoint_index, meta.viewpoint_index);
493    }
494}
bevy_sensor/fixtures.rs

bevy_sensor/
fixtures.rs
