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
201    pub fn intrinsics(&self) -> CameraIntrinsics {
202        CameraIntrinsics {
203            focal_length: self.metadata.intrinsics.focal_length,
204            principal_point: self.metadata.intrinsics.principal_point,
205            image_size: self.metadata.intrinsics.image_size,
206        }
207    }
208
209    /// Load a specific render by object, rotation index, and viewpoint index
210    ///
211    /// # Arguments
212    /// * `object_id` - YCB object ID (e.g., "003_cracker_box")
213    /// * `rotation_idx` - Rotation index (0-2 for benchmark rotations)
214    /// * `viewpoint_idx` - Viewpoint index (0-23 for default config)
215    pub fn get_render(
216        &self,
217        object_id: &str,
218        rotation_idx: usize,
219        viewpoint_idx: usize,
220    ) -> Result<RenderOutput, FixtureError> {
221        // Find the render metadata
222        let renders = self
223            .indices
224            .get(object_id)
225            .ok_or_else(|| FixtureError::RenderNotFound {
226                object_id: object_id.to_string(),
227                rotation: rotation_idx,
228                viewpoint: viewpoint_idx,
229            })?;
230
231        let render_meta = renders
232            .iter()
233            .find(|r| r.rotation_index == rotation_idx && r.viewpoint_index == viewpoint_idx)
234            .ok_or_else(|| FixtureError::RenderNotFound {
235                object_id: object_id.to_string(),
236                rotation: rotation_idx,
237                viewpoint: viewpoint_idx,
238            })?;
239
240        // Load RGBA from PNG
241        let rgba_path = self.root.join(object_id).join(&render_meta.rgba_file);
242        let rgba = load_rgba_png(&rgba_path)?;
243
244        // Load depth from binary
245        let depth_path = self.root.join(object_id).join(&render_meta.depth_file);
246        let depth = load_depth_binary(&depth_path)?;
247
248        // Build camera transform from position (looking at origin)
249        let pos = render_meta.camera_position;
250        let camera_transform =
251            Transform::from_xyz(pos[0], pos[1], pos[2]).looking_at(Vec3::ZERO, Vec3::Y);
252
253        // Build object rotation
254        let rot = render_meta.rotation_euler;
255        let object_rotation = crate::ObjectRotation::new(rot[0], rot[1], rot[2]);
256
257        Ok(RenderOutput {
258            rgba,
259            depth,
260            width: self.metadata.resolution[0],
261            height: self.metadata.resolution[1],
262            intrinsics: self.intrinsics(),
263            camera_transform,
264            object_rotation,
265        })
266    }
267
268    /// Load all renders for an object
269    pub fn get_all_renders(&self, object_id: &str) -> Result<Vec<RenderOutput>, FixtureError> {
270        let renders = self
271            .indices
272            .get(object_id)
273            .ok_or_else(|| FixtureError::RenderNotFound {
274                object_id: object_id.to_string(),
275                rotation: 0,
276                viewpoint: 0,
277            })?;
278
279        let mut outputs = Vec::with_capacity(renders.len());
280        for meta in renders {
281            let output = self.get_render(object_id, meta.rotation_index, meta.viewpoint_index)?;
282            outputs.push(output);
283        }
284
285        Ok(outputs)
286    }
287
288    /// Iterate over all renders for an object
289    pub fn iter_renders<'a>(
290        &'a self,
291        object_id: &'a str,
292    ) -> impl Iterator<Item = Result<(usize, usize, RenderOutput), FixtureError>> + 'a {
293        let renders = self.indices.get(object_id);
294
295        renders.into_iter().flat_map(|v| v.iter()).map(move |meta| {
296            let output = self.get_render(object_id, meta.rotation_index, meta.viewpoint_index)?;
297            Ok((meta.rotation_index, meta.viewpoint_index, output))
298        })
299    }
300}
301
302/// Load RGBA data from a PNG file
303fn load_rgba_png(path: &Path) -> Result<Vec<u8>, FixtureError> {
304    let img = image::open(path).map_err(|e| FixtureError::IoError(std::io::Error::other(e)))?;
305
306    let rgba = img.to_rgba8();
307    Ok(rgba.into_raw())
308}
309
310/// Load depth data from binary f32 file
311fn load_depth_binary(path: &Path) -> Result<Vec<f32>, FixtureError> {
312    let bytes = fs::read(path)?;
313
314    // Convert from little-endian bytes to f32
315    let depth: Vec<f32> = bytes
316        .chunks_exact(4)
317        .map(|chunk| {
318            let arr: [u8; 4] = chunk.try_into().unwrap();
319            f32::from_le_bytes(arr)
320        })
321        .collect();
322
323    Ok(depth)
324}
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329    use tempfile::TempDir;
330
331    #[test]
332    fn test_fixture_not_found() {
333        let result = TestFixtures::load("/nonexistent/path");
334        assert!(matches!(result, Err(FixtureError::NotFound(_))));
335    }
336
337    #[test]
338    fn test_fixtures_exists() {
339        assert!(!TestFixtures::exists("/nonexistent/path"));
340    }
341
342    #[test]
343    fn test_fixture_error_display() {
344        let errors = vec![
345            FixtureError::NotFound("/path".to_string()),
346            FixtureError::InvalidMetadata("bad json".to_string()),
347            FixtureError::RenderNotFound {
348                object_id: "obj".to_string(),
349                rotation: 0,
350                viewpoint: 5,
351            },
352            FixtureError::IoError(std::io::Error::new(
353                std::io::ErrorKind::NotFound,
354                "file not found",
355            )),
356            FixtureError::JsonError(serde_json::from_str::<String>("invalid").unwrap_err()),
357        ];
358
359        for err in errors {
360            let msg = err.to_string();
361            assert!(!msg.is_empty());
362        }
363    }
364
365    #[test]
366    fn test_fixture_missing_metadata() {
367        let temp_dir = TempDir::new().unwrap();
368        let result = TestFixtures::load(temp_dir.path());
369        assert!(matches!(result, Err(FixtureError::InvalidMetadata(_))));
370    }
371
372    #[test]
373    fn test_fixture_load_metadata() {
374        let temp_dir = TempDir::new().unwrap();
375
376        // Create minimal metadata
377        let metadata = DatasetMetadata {
378            version: "1.0".to_string(),
379            objects: vec!["test_object".to_string()],
380            viewpoints_per_rotation: 24,
381            rotations_per_object: 3,
382            renders_per_object: 72,
383            resolution: [64, 64],
384            intrinsics: IntrinsicsMetadata {
385                focal_length: [55.4, 55.4],
386                principal_point: [32.0, 32.0],
387                image_size: [64, 64],
388            },
389            viewpoint_config: ViewpointConfigMetadata {
390                radius: 0.5,
391                yaw_count: 8,
392                pitch_angles_deg: vec![-30.0, 0.0, 30.0],
393            },
394            rotations: vec![[0.0, 0.0, 0.0], [0.0, 90.0, 0.0], [0.0, 180.0, 0.0]],
395        };
396
397        let metadata_json = serde_json::to_string_pretty(&metadata).unwrap();
398        let metadata_path = temp_dir.path().join("metadata.json");
399        fs::write(&metadata_path, &metadata_json).unwrap();
400
401        // Create object directory with empty index
402        let obj_dir = temp_dir.path().join("test_object");
403        fs::create_dir_all(&obj_dir).unwrap();
404        fs::write(obj_dir.join("index.json"), "[]").unwrap();
405
406        // Load fixtures
407        let fixtures = TestFixtures::load(temp_dir.path()).unwrap();
408
409        assert_eq!(fixtures.objects(), &["test_object"]);
410        assert_eq!(fixtures.viewpoints_per_rotation(), 24);
411        assert_eq!(fixtures.rotations_per_object(), 3);
412        assert_eq!(fixtures.renders_for_object("test_object"), 0);
413        assert_eq!(fixtures.renders_for_object("nonexistent"), 0);
414
415        let intrinsics = fixtures.intrinsics();
416        assert_eq!(intrinsics.image_size, [64, 64]);
417    }
418
419    #[test]
420    fn test_load_depth_binary() {
421        let temp_dir = TempDir::new().unwrap();
422        let depth_path = temp_dir.path().join("test.depth");
423
424        // Write test depth values
425        let depths: Vec<f32> = vec![0.5, 1.0, 2.0, 10.0];
426        let bytes: Vec<u8> = depths.iter().flat_map(|f| f.to_le_bytes()).collect();
427        fs::write(&depth_path, &bytes).unwrap();
428
429        // Load and verify
430        let loaded = load_depth_binary(&depth_path).unwrap();
431        assert_eq!(loaded.len(), 4);
432        assert!((loaded[0] - 0.5).abs() < 0.001);
433        assert!((loaded[1] - 1.0).abs() < 0.001);
434        assert!((loaded[2] - 2.0).abs() < 0.001);
435        assert!((loaded[3] - 10.0).abs() < 0.001);
436    }
437
438    #[test]
439    fn test_metadata_serialization_roundtrip() {
440        let metadata = DatasetMetadata {
441            version: "1.0".to_string(),
442            objects: vec!["obj1".to_string(), "obj2".to_string()],
443            viewpoints_per_rotation: 24,
444            rotations_per_object: 3,
445            renders_per_object: 72,
446            resolution: [64, 64],
447            intrinsics: IntrinsicsMetadata {
448                focal_length: [55.4, 55.4],
449                principal_point: [32.0, 32.0],
450                image_size: [64, 64],
451            },
452            viewpoint_config: ViewpointConfigMetadata {
453                radius: 0.5,
454                yaw_count: 8,
455                pitch_angles_deg: vec![-30.0, 0.0, 30.0],
456            },
457            rotations: vec![[0.0, 0.0, 0.0]],
458        };
459
460        let json = serde_json::to_string(&metadata).unwrap();
461        let loaded: DatasetMetadata = serde_json::from_str(&json).unwrap();
462
463        assert_eq!(loaded.version, metadata.version);
464        assert_eq!(loaded.objects, metadata.objects);
465        assert_eq!(loaded.resolution, metadata.resolution);
466    }
467
468    #[test]
469    fn test_render_metadata_serialization() {
470        let meta = RenderMetadata {
471            object_id: "003_cracker_box".to_string(),
472            rotation_index: 1,
473            viewpoint_index: 5,
474            rotation_euler: [0.0, 90.0, 0.0],
475            camera_position: [0.5, 0.0, 0.0],
476            rgba_file: "r1_v05.png".to_string(),
477            depth_file: "r1_v05.depth".to_string(),
478        };
479
480        let json = serde_json::to_string(&meta).unwrap();
481        let loaded: RenderMetadata = serde_json::from_str(&json).unwrap();
482
483        assert_eq!(loaded.object_id, meta.object_id);
484        assert_eq!(loaded.rotation_index, meta.rotation_index);
485        assert_eq!(loaded.viewpoint_index, meta.viewpoint_index);
486    }
487}
bevy_sensor/fixtures.rs

bevy_sensor/
fixtures.rs
