1use bevy::prelude::*;
54use std::f32::consts::PI;
55use std::path::Path;
56
57mod render;
60
61pub mod batch;
63
64pub mod backend;
66
67pub mod cache;
69
70pub mod fixtures;
72
73pub use ycbust::{
75 self, DownloadOptions, Subset as YcbSubset, REPRESENTATIVE_OBJECTS, TBP_SIMILAR_OBJECTS,
76 TBP_STANDARD_OBJECTS,
77};
78
79pub mod ycb {
81 pub use ycbust::{
82 download_ycb, DownloadOptions, Subset, REPRESENTATIVE_OBJECTS, TBP_SIMILAR_OBJECTS,
83 TBP_STANDARD_OBJECTS,
84 };
85
86 use std::path::Path;
87
88 pub async fn download_models<P: AsRef<Path>>(
101 output_dir: P,
102 subset: Subset,
103 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
104 download_ycb(subset, output_dir.as_ref(), DownloadOptions::default()).await?;
105 Ok(())
106 }
107
108 pub async fn download_models_with_options<P: AsRef<Path>>(
110 output_dir: P,
111 subset: Subset,
112 options: DownloadOptions,
113 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
114 download_ycb(subset, output_dir.as_ref(), options).await?;
115 Ok(())
116 }
117
118 pub async fn download_objects<P: AsRef<Path>>(
124 output_dir: P,
125 object_ids: &[&str],
126 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
127 ycbust::download_objects(object_ids, output_dir.as_ref(), DownloadOptions::default())
128 .await?;
129 Ok(())
130 }
131
132 pub fn models_exist<P: AsRef<Path>>(output_dir: P) -> bool {
134 ycbust::object_mesh_path(output_dir.as_ref(), "003_cracker_box").exists()
135 }
136
137 pub fn object_mesh_path<P: AsRef<Path>>(output_dir: P, object_id: &str) -> std::path::PathBuf {
139 ycbust::object_mesh_path(output_dir.as_ref(), object_id)
140 }
141
142 pub fn object_texture_path<P: AsRef<Path>>(
144 output_dir: P,
145 object_id: &str,
146 ) -> std::path::PathBuf {
147 ycbust::object_texture_path(output_dir.as_ref(), object_id)
148 }
149}
150
151pub fn initialize() {
185 use std::sync::atomic::{AtomicBool, Ordering};
187 static INITIALIZED: AtomicBool = AtomicBool::new(false);
188
189 if !INITIALIZED.swap(true, Ordering::SeqCst) {
190 let config = backend::BackendConfig::new();
192 config.apply_env();
193 }
194}
195
196#[derive(Clone, Debug, PartialEq)]
199pub struct ObjectRotation {
200 pub pitch: f64,
202 pub yaw: f64,
204 pub roll: f64,
206}
207
208impl ObjectRotation {
209 pub fn new(pitch: f64, yaw: f64, roll: f64) -> Self {
211 Self { pitch, yaw, roll }
212 }
213
214 pub fn from_array(arr: [f64; 3]) -> Self {
216 Self {
217 pitch: arr[0],
218 yaw: arr[1],
219 roll: arr[2],
220 }
221 }
222
223 pub fn identity() -> Self {
225 Self::new(0.0, 0.0, 0.0)
226 }
227
228 pub fn tbp_benchmark_rotations() -> Vec<Self> {
231 vec![
232 Self::from_array([0.0, 0.0, 0.0]),
233 Self::from_array([0.0, 90.0, 0.0]),
234 Self::from_array([0.0, 180.0, 0.0]),
235 ]
236 }
237
238 pub fn tbp_known_orientations() -> Vec<Self> {
241 vec![
242 Self::from_array([0.0, 0.0, 0.0]), Self::from_array([0.0, 90.0, 0.0]), Self::from_array([0.0, 180.0, 0.0]), Self::from_array([0.0, 270.0, 0.0]), Self::from_array([90.0, 0.0, 0.0]), Self::from_array([-90.0, 0.0, 0.0]), Self::from_array([45.0, 45.0, 0.0]),
251 Self::from_array([45.0, 135.0, 0.0]),
252 Self::from_array([45.0, 225.0, 0.0]),
253 Self::from_array([45.0, 315.0, 0.0]),
254 Self::from_array([-45.0, 45.0, 0.0]),
255 Self::from_array([-45.0, 135.0, 0.0]),
256 Self::from_array([-45.0, 225.0, 0.0]),
257 Self::from_array([-45.0, 315.0, 0.0]),
258 ]
259 }
260
261 pub fn to_quat(&self) -> Quat {
263 Quat::from_euler(
264 EulerRot::XYZ,
265 (self.pitch as f32).to_radians(),
266 (self.yaw as f32).to_radians(),
267 (self.roll as f32).to_radians(),
268 )
269 }
270
271 pub fn to_transform(&self) -> Transform {
273 Transform::from_rotation(self.to_quat())
274 }
275}
276
277impl Default for ObjectRotation {
278 fn default() -> Self {
279 Self::identity()
280 }
281}
282
283#[derive(Clone, Debug)]
286pub struct ViewpointConfig {
287 pub radius: f32,
289 pub yaw_count: usize,
291 pub pitch_angles_deg: Vec<f32>,
293}
294
295impl Default for ViewpointConfig {
296 fn default() -> Self {
297 Self {
298 radius: 0.5,
299 yaw_count: 8,
300 pitch_angles_deg: vec![-30.0, 0.0, 30.0],
303 }
304 }
305}
306
307impl ViewpointConfig {
308 pub fn viewpoint_count(&self) -> usize {
310 self.yaw_count * self.pitch_angles_deg.len()
311 }
312}
313
314#[derive(Clone, Debug, Resource)]
316pub struct SensorConfig {
317 pub viewpoints: ViewpointConfig,
319 pub object_rotations: Vec<ObjectRotation>,
321 pub output_dir: String,
323 pub filename_pattern: String,
325}
326
327impl Default for SensorConfig {
328 fn default() -> Self {
329 Self {
330 viewpoints: ViewpointConfig::default(),
331 object_rotations: vec![ObjectRotation::identity()],
332 output_dir: ".".to_string(),
333 filename_pattern: "capture_{rot}_{view}.png".to_string(),
334 }
335 }
336}
337
338impl SensorConfig {
339 pub fn tbp_benchmark() -> Self {
341 Self {
342 viewpoints: ViewpointConfig::default(),
343 object_rotations: ObjectRotation::tbp_benchmark_rotations(),
344 output_dir: ".".to_string(),
345 filename_pattern: "capture_{rot}_{view}.png".to_string(),
346 }
347 }
348
349 pub fn tbp_full_training() -> Self {
351 Self {
352 viewpoints: ViewpointConfig::default(),
353 object_rotations: ObjectRotation::tbp_known_orientations(),
354 output_dir: ".".to_string(),
355 filename_pattern: "capture_{rot}_{view}.png".to_string(),
356 }
357 }
358
359 pub fn total_captures(&self) -> usize {
361 self.viewpoints.viewpoint_count() * self.object_rotations.len()
362 }
363}
364
365pub fn generate_viewpoints(config: &ViewpointConfig) -> Vec<Transform> {
372 let mut views = Vec::with_capacity(config.viewpoint_count());
373
374 for pitch_deg in &config.pitch_angles_deg {
375 let pitch = pitch_deg.to_radians();
376
377 for i in 0..config.yaw_count {
378 let yaw = (i as f32) * 2.0 * PI / (config.yaw_count as f32);
379
380 let x = config.radius * pitch.cos() * yaw.sin();
385 let y = config.radius * pitch.sin();
386 let z = config.radius * pitch.cos() * yaw.cos();
387
388 let transform = Transform::from_xyz(x, y, z).looking_at(Vec3::ZERO, Vec3::Y);
389 views.push(transform);
390 }
391 }
392 views
393}
394
395#[derive(Component)]
397pub struct CaptureTarget;
398
399#[derive(Component)]
401pub struct CaptureCamera;
402
403#[derive(Clone, Debug, PartialEq)]
411pub struct RenderConfig {
412 pub width: u32,
414 pub height: u32,
416 pub zoom: f32,
419 pub near_plane: f32,
421 pub far_plane: f32,
423 pub lighting: LightingConfig,
425}
426
427#[derive(Clone, Debug, PartialEq)]
431pub struct LightingConfig {
432 pub ambient_brightness: f32,
434 pub key_light_intensity: f32,
436 pub key_light_position: [f32; 3],
438 pub fill_light_intensity: f32,
440 pub fill_light_position: [f32; 3],
442 pub shadows_enabled: bool,
444}
445
446impl Default for LightingConfig {
447 fn default() -> Self {
448 Self {
449 ambient_brightness: 0.3,
450 key_light_intensity: 1500.0,
451 key_light_position: [4.0, 8.0, 4.0],
452 fill_light_intensity: 500.0,
453 fill_light_position: [-4.0, 2.0, -4.0],
454 shadows_enabled: false,
455 }
456 }
457}
458
459impl LightingConfig {
460 pub fn bright() -> Self {
462 Self {
463 ambient_brightness: 0.5,
464 key_light_intensity: 2000.0,
465 key_light_position: [4.0, 8.0, 4.0],
466 fill_light_intensity: 800.0,
467 fill_light_position: [-4.0, 2.0, -4.0],
468 shadows_enabled: false,
469 }
470 }
471
472 pub fn soft() -> Self {
474 Self {
475 ambient_brightness: 0.4,
476 key_light_intensity: 1000.0,
477 key_light_position: [3.0, 6.0, 3.0],
478 fill_light_intensity: 600.0,
479 fill_light_position: [-3.0, 3.0, -3.0],
480 shadows_enabled: false,
481 }
482 }
483
484 pub fn unlit() -> Self {
486 Self {
487 ambient_brightness: 1.0,
488 key_light_intensity: 0.0,
489 key_light_position: [0.0, 0.0, 0.0],
490 fill_light_intensity: 0.0,
491 fill_light_position: [0.0, 0.0, 0.0],
492 shadows_enabled: false,
493 }
494 }
495}
496
497impl Default for RenderConfig {
498 fn default() -> Self {
499 Self::tbp_default()
500 }
501}
502
503impl RenderConfig {
504 pub fn tbp_default() -> Self {
512 Self {
513 width: 64,
514 height: 64,
515 zoom: 4.0,
516 near_plane: 0.01,
517 far_plane: 10.0,
518 lighting: LightingConfig::default(),
519 }
520 }
521
522 pub fn preview() -> Self {
524 Self {
525 width: 256,
526 height: 256,
527 zoom: 1.0,
528 near_plane: 0.01,
529 far_plane: 10.0,
530 lighting: LightingConfig::default(),
531 }
532 }
533
534 pub fn high_res() -> Self {
536 Self {
537 width: 512,
538 height: 512,
539 zoom: 1.0,
540 near_plane: 0.01,
541 far_plane: 10.0,
542 lighting: LightingConfig::default(),
543 }
544 }
545
546 pub fn fov_radians(&self) -> f32 {
553 let base_hfov_rad = 90.0_f32.to_radians();
554 let half_tan = (base_hfov_rad / 2.0).tan() / self.zoom;
555 2.0 * half_tan.atan()
556 }
557
558 pub fn intrinsics(&self) -> CameraIntrinsics {
566 self.intrinsics_for_size(self.width, self.height)
567 }
568
569 pub fn intrinsics_for_size(&self, width: u32, height: u32) -> CameraIntrinsics {
574 let base_hfov_rad = 90.0_f64.to_radians();
575 let fx_norm = (base_hfov_rad / 2.0).tan() / self.zoom as f64;
577 let fx = (width as f64 / 2.0) / fx_norm;
579 let fy = fx; CameraIntrinsics {
582 focal_length: [fx, fy],
583 principal_point: [width as f64 / 2.0, height as f64 / 2.0],
584 image_size: [width, height],
585 }
586 }
587}
588
589#[derive(Clone, Debug, PartialEq)]
594pub struct CameraIntrinsics {
595 pub focal_length: [f64; 2],
597 pub principal_point: [f64; 2],
599 pub image_size: [u32; 2],
601}
602
603impl CameraIntrinsics {
604 pub fn project(&self, point: Vec3) -> Option<[f64; 2]> {
606 if point.z <= 0.0 {
607 return None;
608 }
609 let x = (point.x as f64 / point.z as f64) * self.focal_length[0] + self.principal_point[0];
610 let y = (point.y as f64 / point.z as f64) * self.focal_length[1] + self.principal_point[1];
611 Some([x, y])
612 }
613
614 pub fn unproject(&self, pixel: [f64; 2], depth: f64) -> [f64; 3] {
616 let x = (pixel[0] - self.principal_point[0]) / self.focal_length[0] * depth;
617 let y = (pixel[1] - self.principal_point[1]) / self.focal_length[1] * depth;
618 [x, y, depth]
619 }
620}
621
622#[derive(Clone, Debug)]
624pub struct RenderOutput {
625 pub rgba: Vec<u8>,
627 pub depth: Vec<f64>,
631 pub width: u32,
633 pub height: u32,
635 pub intrinsics: CameraIntrinsics,
637 pub camera_transform: Transform,
639 pub object_rotation: ObjectRotation,
641}
642
643impl RenderOutput {
644 pub fn get_rgba(&self, x: u32, y: u32) -> Option<[u8; 4]> {
646 if x >= self.width || y >= self.height {
647 return None;
648 }
649 let idx = ((y * self.width + x) * 4) as usize;
650 Some([
651 self.rgba[idx],
652 self.rgba[idx + 1],
653 self.rgba[idx + 2],
654 self.rgba[idx + 3],
655 ])
656 }
657
658 pub fn get_depth(&self, x: u32, y: u32) -> Option<f64> {
660 if x >= self.width || y >= self.height {
661 return None;
662 }
663 let idx = (y * self.width + x) as usize;
664 Some(self.depth[idx])
665 }
666
667 pub fn get_rgb(&self, x: u32, y: u32) -> Option<[u8; 3]> {
669 self.get_rgba(x, y).map(|rgba| [rgba[0], rgba[1], rgba[2]])
670 }
671
672 pub fn to_rgb_image(&self) -> Vec<Vec<[u8; 3]>> {
674 let mut image = Vec::with_capacity(self.height as usize);
675 for y in 0..self.height {
676 let mut row = Vec::with_capacity(self.width as usize);
677 for x in 0..self.width {
678 row.push(self.get_rgb(x, y).unwrap_or([0, 0, 0]));
679 }
680 image.push(row);
681 }
682 image
683 }
684
685 pub fn to_depth_image(&self) -> Vec<Vec<f64>> {
687 let mut image = Vec::with_capacity(self.height as usize);
688 for y in 0..self.height {
689 let mut row = Vec::with_capacity(self.width as usize);
690 for x in 0..self.width {
691 row.push(self.get_depth(x, y).unwrap_or(0.0));
692 }
693 image.push(row);
694 }
695 image
696 }
697}
698
699#[derive(Debug, Clone)]
701pub enum RenderError {
702 MeshNotFound(String),
704 TextureNotFound(String),
706 FileNotFound { path: String, reason: String },
708 FileWriteFailed { path: String, reason: String },
710 DirectoryCreationFailed { path: String, reason: String },
712 RenderFailed(String),
714 InvalidConfig(String),
716 InvalidInput(String),
718 SerializationError(String),
720 DataParsingError(String),
722 RenderTimeout { duration_secs: u64 },
724}
725
726impl std::fmt::Display for RenderError {
727 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
728 match self {
729 RenderError::MeshNotFound(path) => write!(f, "Mesh not found: {}", path),
730 RenderError::TextureNotFound(path) => write!(f, "Texture not found: {}", path),
731 RenderError::FileNotFound { path, reason } => {
732 write!(f, "File not found at {}: {}", path, reason)
733 }
734 RenderError::FileWriteFailed { path, reason } => {
735 write!(f, "Failed to write file {}: {}", path, reason)
736 }
737 RenderError::DirectoryCreationFailed { path, reason } => {
738 write!(f, "Failed to create directory {}: {}", path, reason)
739 }
740 RenderError::RenderFailed(msg) => write!(f, "Render failed: {}", msg),
741 RenderError::InvalidConfig(msg) => write!(f, "Invalid config: {}", msg),
742 RenderError::InvalidInput(msg) => write!(f, "Invalid input: {}", msg),
743 RenderError::SerializationError(msg) => write!(f, "Serialization error: {}", msg),
744 RenderError::DataParsingError(msg) => write!(f, "Data parsing error: {}", msg),
745 RenderError::RenderTimeout { duration_secs } => {
746 write!(f, "Render timeout after {} seconds", duration_secs)
747 }
748 }
749 }
750}
751
752impl std::error::Error for RenderError {}
753
754pub fn render_to_buffer(
779 object_dir: &Path,
780 camera_transform: &Transform,
781 object_rotation: &ObjectRotation,
782 config: &RenderConfig,
783) -> Result<RenderOutput, RenderError> {
784 render::render_headless(object_dir, camera_transform, object_rotation, config)
786}
787
788pub fn render_all_viewpoints(
801 object_dir: &Path,
802 viewpoint_config: &ViewpointConfig,
803 rotations: &[ObjectRotation],
804 render_config: &RenderConfig,
805) -> Result<Vec<RenderOutput>, RenderError> {
806 let viewpoints = generate_viewpoints(viewpoint_config);
807 let mut outputs = Vec::with_capacity(viewpoints.len() * rotations.len());
808
809 for rotation in rotations {
810 for viewpoint in &viewpoints {
811 let output = render_to_buffer(object_dir, viewpoint, rotation, render_config)?;
812 outputs.push(output);
813 }
814 }
815
816 Ok(outputs)
817}
818
819pub fn render_to_buffer_cached(
889 object_dir: &Path,
890 camera_transform: &Transform,
891 object_rotation: &ObjectRotation,
892 config: &RenderConfig,
893 cache: &mut cache::ModelCache,
894) -> Result<RenderOutput, RenderError> {
895 let mesh_path = object_dir.join("google_16k/textured.obj");
896 let texture_path = object_dir.join("google_16k/texture_map.png");
897
898 cache.cache_scene(mesh_path.clone());
900 cache.cache_texture(texture_path.clone());
901
902 render::render_headless(object_dir, camera_transform, object_rotation, config)
904}
905
906pub fn render_to_files(
923 object_dir: &Path,
924 camera_transform: &Transform,
925 object_rotation: &ObjectRotation,
926 config: &RenderConfig,
927 rgba_path: &Path,
928 depth_path: &Path,
929) -> Result<(), RenderError> {
930 render::render_to_files(
931 object_dir,
932 camera_transform,
933 object_rotation,
934 config,
935 rgba_path,
936 depth_path,
937 )
938}
939
940pub use batch::{
942 BatchRenderConfig, BatchRenderError, BatchRenderOutput, BatchRenderRequest, BatchRenderer,
943 BatchState, RenderStatus,
944};
945
946pub use render::RenderSession;
949
950pub use render::PersistentRenderer;
956
957pub fn create_batch_renderer(config: &BatchRenderConfig) -> Result<BatchRenderer, RenderError> {
975 Ok(BatchRenderer::new(config.clone()))
976}
977
978pub fn queue_render_request(
1003 renderer: &mut BatchRenderer,
1004 request: BatchRenderRequest,
1005) -> Result<(), RenderError> {
1006 renderer
1007 .queue_request(request)
1008 .map_err(|e| RenderError::RenderFailed(e.to_string()))
1009}
1010
1011pub fn render_next_in_batch(
1033 renderer: &mut BatchRenderer,
1034 _timeout_ms: u32,
1035) -> Result<Option<BatchRenderOutput>, RenderError> {
1036 if let Some(request) = renderer.pending_requests.pop_front() {
1037 let output = render_to_buffer(
1038 &request.object_dir,
1039 &request.viewpoint,
1040 &request.object_rotation,
1041 &request.render_config,
1042 )?;
1043 let batch_output = BatchRenderOutput::from_render_output(request, output);
1044 renderer.completed_results.push(batch_output.clone());
1045 renderer.renders_processed += 1;
1046 Ok(Some(batch_output))
1047 } else {
1048 Ok(None)
1049 }
1050}
1051
1052pub fn render_batch(
1071 requests: Vec<BatchRenderRequest>,
1072 config: &BatchRenderConfig,
1073) -> Result<Vec<BatchRenderOutput>, RenderError> {
1074 if requests.is_empty() {
1075 return Ok(Vec::new());
1076 }
1077
1078 if requests.len() > 1 && requests_share_batch_context(&requests) {
1079 let first_request = requests[0].clone();
1080 let viewpoints: Vec<Transform> = requests.iter().map(|request| request.viewpoint).collect();
1081 let outputs = render::render_headless_sequence(
1082 &first_request.object_dir,
1083 &viewpoints,
1084 &first_request.object_rotation,
1085 &first_request.render_config,
1086 )?;
1087
1088 return Ok(requests
1089 .into_iter()
1090 .zip(outputs)
1091 .map(|(request, output)| BatchRenderOutput::from_render_output(request, output))
1092 .collect());
1093 }
1094
1095 let mut renderer = create_batch_renderer(config)?;
1096
1097 for request in requests {
1099 queue_render_request(&mut renderer, request)?;
1100 }
1101
1102 let mut results = Vec::new();
1104 while let Some(output) = render_next_in_batch(&mut renderer, config.frame_timeout_ms)? {
1105 results.push(output);
1106 }
1107
1108 Ok(results)
1109}
1110
1111fn requests_share_batch_context(requests: &[BatchRenderRequest]) -> bool {
1112 let Some(first) = requests.first() else {
1113 return true;
1114 };
1115
1116 requests.iter().all(|request| {
1117 request.object_dir == first.object_dir
1118 && request.object_rotation == first.object_rotation
1119 && request.render_config == first.render_config
1120 })
1121}
1122
1123pub use bevy::prelude::{Quat, Transform, Vec3};
1125
1126#[cfg(test)]
1127mod tests {
1128 use super::*;
1129
1130 #[test]
1131 fn test_object_rotation_identity() {
1132 let rot = ObjectRotation::identity();
1133 assert_eq!(rot.pitch, 0.0);
1134 assert_eq!(rot.yaw, 0.0);
1135 assert_eq!(rot.roll, 0.0);
1136 }
1137
1138 #[test]
1139 fn test_object_rotation_from_array() {
1140 let rot = ObjectRotation::from_array([10.0, 20.0, 30.0]);
1141 assert_eq!(rot.pitch, 10.0);
1142 assert_eq!(rot.yaw, 20.0);
1143 assert_eq!(rot.roll, 30.0);
1144 }
1145
1146 #[test]
1147 fn test_requests_share_batch_context_for_homogeneous_batch() {
1148 let config = RenderConfig::tbp_default();
1149 let request = BatchRenderRequest {
1150 object_dir: "/tmp/ycb/003_cracker_box".into(),
1151 viewpoint: Transform::IDENTITY,
1152 object_rotation: ObjectRotation::identity(),
1153 render_config: config.clone(),
1154 };
1155
1156 assert!(requests_share_batch_context(&[
1157 request.clone(),
1158 BatchRenderRequest {
1159 viewpoint: Transform::from_xyz(1.0, 0.0, 0.0),
1160 ..request
1161 },
1162 ]));
1163 }
1164
1165 #[test]
1166 fn test_requests_share_batch_context_rejects_mixed_objects() {
1167 let config = RenderConfig::tbp_default();
1168 let request = BatchRenderRequest {
1169 object_dir: "/tmp/ycb/003_cracker_box".into(),
1170 viewpoint: Transform::IDENTITY,
1171 object_rotation: ObjectRotation::identity(),
1172 render_config: config.clone(),
1173 };
1174
1175 assert!(!requests_share_batch_context(&[
1176 request.clone(),
1177 BatchRenderRequest {
1178 object_dir: "/tmp/ycb/005_tomato_soup_can".into(),
1179 ..request
1180 },
1181 ]));
1182 }
1183
1184 #[test]
1185 fn test_tbp_benchmark_rotations() {
1186 let rotations = ObjectRotation::tbp_benchmark_rotations();
1187 assert_eq!(rotations.len(), 3);
1188 assert_eq!(rotations[0], ObjectRotation::from_array([0.0, 0.0, 0.0]));
1189 assert_eq!(rotations[1], ObjectRotation::from_array([0.0, 90.0, 0.0]));
1190 assert_eq!(rotations[2], ObjectRotation::from_array([0.0, 180.0, 0.0]));
1191 }
1192
1193 #[test]
1194 fn test_tbp_known_orientations_count() {
1195 let orientations = ObjectRotation::tbp_known_orientations();
1196 assert_eq!(orientations.len(), 14);
1197 }
1198
1199 #[test]
1200 fn test_rotation_to_quat() {
1201 let rot = ObjectRotation::identity();
1202 let quat = rot.to_quat();
1203 assert!((quat.w - 1.0).abs() < 0.001);
1205 assert!(quat.x.abs() < 0.001);
1206 assert!(quat.y.abs() < 0.001);
1207 assert!(quat.z.abs() < 0.001);
1208 }
1209
1210 #[test]
1211 fn test_rotation_90_yaw() {
1212 let rot = ObjectRotation::new(0.0, 90.0, 0.0);
1213 let quat = rot.to_quat();
1214 assert!((quat.w - 0.707).abs() < 0.01);
1216 assert!((quat.y - 0.707).abs() < 0.01);
1217 }
1218
1219 #[test]
1220 fn test_viewpoint_config_default() {
1221 let config = ViewpointConfig::default();
1222 assert_eq!(config.radius, 0.5);
1223 assert_eq!(config.yaw_count, 8);
1224 assert_eq!(config.pitch_angles_deg.len(), 3);
1225 }
1226
1227 #[test]
1228 fn test_viewpoint_count() {
1229 let config = ViewpointConfig::default();
1230 assert_eq!(config.viewpoint_count(), 24); }
1232
1233 #[test]
1234 fn test_generate_viewpoints_count() {
1235 let config = ViewpointConfig::default();
1236 let viewpoints = generate_viewpoints(&config);
1237 assert_eq!(viewpoints.len(), 24);
1238 }
1239
1240 #[test]
1241 fn test_viewpoints_spherical_radius() {
1242 let config = ViewpointConfig::default();
1243 let viewpoints = generate_viewpoints(&config);
1244
1245 for (i, transform) in viewpoints.iter().enumerate() {
1246 let actual_radius = transform.translation.length();
1247 assert!(
1248 (actual_radius - config.radius).abs() < 0.001,
1249 "Viewpoint {} has incorrect radius: {} (expected {})",
1250 i,
1251 actual_radius,
1252 config.radius
1253 );
1254 }
1255 }
1256
1257 #[test]
1258 fn test_viewpoints_looking_at_origin() {
1259 let config = ViewpointConfig::default();
1260 let viewpoints = generate_viewpoints(&config);
1261
1262 for (i, transform) in viewpoints.iter().enumerate() {
1263 let forward = transform.forward();
1264 let to_origin = (Vec3::ZERO - transform.translation).normalize();
1265 let dot = forward.dot(to_origin);
1266 assert!(
1267 dot > 0.99,
1268 "Viewpoint {} not looking at origin, dot product: {}",
1269 i,
1270 dot
1271 );
1272 }
1273 }
1274
1275 #[test]
1276 fn test_sensor_config_default() {
1277 let config = SensorConfig::default();
1278 assert_eq!(config.object_rotations.len(), 1);
1279 assert_eq!(config.total_captures(), 24);
1280 }
1281
1282 #[test]
1283 fn test_sensor_config_tbp_benchmark() {
1284 let config = SensorConfig::tbp_benchmark();
1285 assert_eq!(config.object_rotations.len(), 3);
1286 assert_eq!(config.total_captures(), 72); }
1288
1289 #[test]
1290 fn test_sensor_config_tbp_full() {
1291 let config = SensorConfig::tbp_full_training();
1292 assert_eq!(config.object_rotations.len(), 14);
1293 assert_eq!(config.total_captures(), 336); }
1295
1296 #[test]
1297 fn test_ycb_representative_objects() {
1298 assert_eq!(crate::ycb::REPRESENTATIVE_OBJECTS.len(), 3);
1300 assert!(crate::ycb::REPRESENTATIVE_OBJECTS.contains(&"003_cracker_box"));
1301 }
1302
1303 #[test]
1304 fn test_ycb_tbp_standard_objects() {
1305 assert_eq!(crate::ycb::TBP_STANDARD_OBJECTS.len(), 10);
1306 assert!(crate::ycb::TBP_STANDARD_OBJECTS.contains(&"025_mug"));
1307 }
1308
1309 #[test]
1310 fn test_ycb_tbp_similar_objects() {
1311 assert_eq!(crate::ycb::TBP_SIMILAR_OBJECTS.len(), 10);
1312 assert!(crate::ycb::TBP_SIMILAR_OBJECTS.contains(&"003_cracker_box"));
1313 }
1314
1315 #[test]
1316 fn test_ycb_object_mesh_path() {
1317 let path = crate::ycb::object_mesh_path("/tmp/ycb", "003_cracker_box");
1318 assert_eq!(
1319 path,
1320 std::path::Path::new("/tmp/ycb")
1321 .join("003_cracker_box")
1322 .join("google_16k")
1323 .join("textured.obj")
1324 );
1325 }
1326
1327 #[test]
1328 fn test_ycb_object_texture_path() {
1329 let path = crate::ycb::object_texture_path("/tmp/ycb", "003_cracker_box");
1330 assert_eq!(
1331 path,
1332 std::path::Path::new("/tmp/ycb")
1333 .join("003_cracker_box")
1334 .join("google_16k")
1335 .join("texture_map.png")
1336 );
1337 }
1338
1339 #[test]
1344 fn test_render_config_tbp_default() {
1345 let config = RenderConfig::tbp_default();
1346 assert_eq!(config.width, 64);
1348 assert_eq!(config.height, 64);
1349 assert!(config.zoom > 0.0);
1351 assert!(config.near_plane > 0.0);
1353 assert!(config.far_plane > config.near_plane);
1354 }
1355
1356 #[test]
1357 fn test_render_config_preview() {
1358 let config = RenderConfig::preview();
1359 assert_eq!(config.width, 256);
1360 assert_eq!(config.height, 256);
1361 }
1362
1363 #[test]
1364 fn test_render_config_default_is_tbp() {
1365 let default = RenderConfig::default();
1366 let tbp = RenderConfig::tbp_default();
1367 assert_eq!(default.width, tbp.width);
1368 assert_eq!(default.height, tbp.height);
1369 }
1370
1371 #[test]
1372 fn test_render_config_fov() {
1373 let config = RenderConfig::tbp_default();
1374 let fov = config.fov_radians();
1375 assert!(fov > 0.0);
1378 assert!(fov < PI);
1379
1380 let zoomed = RenderConfig {
1382 zoom: config.zoom * 2.0,
1383 ..config
1384 };
1385 assert!(zoomed.fov_radians() < fov);
1386 }
1387
1388 #[test]
1389 fn test_render_config_intrinsics() {
1390 let config = RenderConfig::tbp_default();
1391 let intrinsics = config.intrinsics();
1392
1393 assert_eq!(intrinsics.image_size, [config.width, config.height]);
1395 assert_eq!(
1396 intrinsics.principal_point,
1397 [config.width as f64 / 2.0, config.height as f64 / 2.0]
1398 );
1399 assert_eq!(intrinsics.focal_length[0], intrinsics.focal_length[1]);
1401 assert!(intrinsics.focal_length[0] > 0.0);
1402 }
1403
1404 #[test]
1405 fn test_render_config_intrinsics_for_size_uses_tbp_zoom_formula() {
1406 let config = RenderConfig {
1407 width: 64,
1408 height: 64,
1409 zoom: 4.0,
1410 ..RenderConfig::tbp_default()
1411 };
1412
1413 let intrinsics = config.intrinsics_for_size(64, 64);
1414
1415 assert!((intrinsics.focal_length[0] - 128.0).abs() < 1e-9);
1418 assert!((intrinsics.focal_length[1] - 128.0).abs() < 1e-9);
1419 assert_ne!(intrinsics.focal_length[0], 64.0 * config.zoom as f64);
1420 assert_eq!(intrinsics.principal_point, [32.0, 32.0]);
1421 assert_eq!(intrinsics.image_size, [64, 64]);
1422 }
1423
1424 #[test]
1425 fn test_render_config_intrinsics_for_size_tracks_actual_readback_size() {
1426 let config = RenderConfig {
1427 width: 64,
1428 height: 64,
1429 zoom: 4.0,
1430 ..RenderConfig::tbp_default()
1431 };
1432
1433 let intrinsics = config.intrinsics_for_size(128, 96);
1434
1435 assert!((intrinsics.focal_length[0] - 256.0).abs() < 1e-9);
1436 assert!((intrinsics.focal_length[1] - 256.0).abs() < 1e-9);
1437 assert_eq!(intrinsics.principal_point, [64.0, 48.0]);
1438 assert_eq!(intrinsics.image_size, [128, 96]);
1439 }
1440
1441 #[test]
1442 fn test_camera_intrinsics_project() {
1443 let intrinsics = CameraIntrinsics {
1444 focal_length: [100.0, 100.0],
1445 principal_point: [32.0, 32.0],
1446 image_size: [64, 64],
1447 };
1448
1449 let center = intrinsics.project(Vec3::new(0.0, 0.0, 1.0));
1451 assert!(center.is_some());
1452 let [x, y] = center.unwrap();
1453 assert!((x - 32.0).abs() < 0.001);
1454 assert!((y - 32.0).abs() < 0.001);
1455
1456 let behind = intrinsics.project(Vec3::new(0.0, 0.0, -1.0));
1458 assert!(behind.is_none());
1459 }
1460
1461 #[test]
1462 fn test_camera_intrinsics_unproject() {
1463 let intrinsics = CameraIntrinsics {
1464 focal_length: [100.0, 100.0],
1465 principal_point: [32.0, 32.0],
1466 image_size: [64, 64],
1467 };
1468
1469 let point = intrinsics.unproject([32.0, 32.0], 1.0);
1471 assert!((point[0]).abs() < 0.001); assert!((point[1]).abs() < 0.001); assert!((point[2] - 1.0).abs() < 0.001); }
1475
1476 #[test]
1477 fn test_render_output_get_rgba() {
1478 let output = RenderOutput {
1479 rgba: vec![
1480 255, 0, 0, 255, 0, 255, 0, 255, 0, 0, 255, 255, 255, 255, 255, 255,
1481 ],
1482 depth: vec![1.0, 2.0, 3.0, 4.0],
1483 width: 2,
1484 height: 2,
1485 intrinsics: RenderConfig::tbp_default().intrinsics(),
1486 camera_transform: Transform::IDENTITY,
1487 object_rotation: ObjectRotation::identity(),
1488 };
1489
1490 assert_eq!(output.get_rgba(0, 0), Some([255, 0, 0, 255]));
1492 assert_eq!(output.get_rgba(1, 0), Some([0, 255, 0, 255]));
1494 assert_eq!(output.get_rgba(0, 1), Some([0, 0, 255, 255]));
1496 assert_eq!(output.get_rgba(1, 1), Some([255, 255, 255, 255]));
1498 assert_eq!(output.get_rgba(2, 0), None);
1500 }
1501
1502 #[test]
1503 fn test_render_output_get_depth() {
1504 let output = RenderOutput {
1505 rgba: vec![0u8; 16],
1506 depth: vec![1.0, 2.0, 3.0, 4.0],
1507 width: 2,
1508 height: 2,
1509 intrinsics: RenderConfig::tbp_default().intrinsics(),
1510 camera_transform: Transform::IDENTITY,
1511 object_rotation: ObjectRotation::identity(),
1512 };
1513
1514 assert_eq!(output.get_depth(0, 0), Some(1.0));
1515 assert_eq!(output.get_depth(1, 0), Some(2.0));
1516 assert_eq!(output.get_depth(0, 1), Some(3.0));
1517 assert_eq!(output.get_depth(1, 1), Some(4.0));
1518 assert_eq!(output.get_depth(2, 0), None);
1519 }
1520
1521 #[test]
1522 fn test_render_output_to_rgb_image() {
1523 let output = RenderOutput {
1524 rgba: vec![
1525 255, 0, 0, 255, 0, 255, 0, 255, 0, 0, 255, 255, 255, 255, 255, 255,
1526 ],
1527 depth: vec![1.0, 2.0, 3.0, 4.0],
1528 width: 2,
1529 height: 2,
1530 intrinsics: RenderConfig::tbp_default().intrinsics(),
1531 camera_transform: Transform::IDENTITY,
1532 object_rotation: ObjectRotation::identity(),
1533 };
1534
1535 let image = output.to_rgb_image();
1536 assert_eq!(image.len(), 2); assert_eq!(image[0].len(), 2); assert_eq!(image[0][0], [255, 0, 0]); assert_eq!(image[0][1], [0, 255, 0]); assert_eq!(image[1][0], [0, 0, 255]); assert_eq!(image[1][1], [255, 255, 255]); }
1543
1544 #[test]
1545 fn test_render_output_to_depth_image() {
1546 let output = RenderOutput {
1547 rgba: vec![0u8; 16],
1548 depth: vec![1.0, 2.0, 3.0, 4.0],
1549 width: 2,
1550 height: 2,
1551 intrinsics: RenderConfig::tbp_default().intrinsics(),
1552 camera_transform: Transform::IDENTITY,
1553 object_rotation: ObjectRotation::identity(),
1554 };
1555
1556 let depth_image = output.to_depth_image();
1557 assert_eq!(depth_image.len(), 2);
1558 assert_eq!(depth_image[0], vec![1.0, 2.0]);
1559 assert_eq!(depth_image[1], vec![3.0, 4.0]);
1560 }
1561
1562 #[test]
1563 fn test_render_error_display() {
1564 let err = RenderError::MeshNotFound("/path/to/mesh.obj".to_string());
1565 assert!(err.to_string().contains("Mesh not found"));
1566 assert!(err.to_string().contains("/path/to/mesh.obj"));
1567 }
1568
1569 #[test]
1574 fn test_object_rotation_extreme_angles() {
1575 let rot = ObjectRotation::new(450.0, -720.0, 1080.0);
1577 let quat = rot.to_quat();
1578 assert!((quat.length() - 1.0).abs() < 0.001);
1580 }
1581
1582 #[test]
1583 fn test_object_rotation_to_transform() {
1584 let rot = ObjectRotation::new(45.0, 90.0, 0.0);
1585 let transform = rot.to_transform();
1586 assert_eq!(transform.translation, Vec3::ZERO);
1588 assert!(transform.rotation != Quat::IDENTITY);
1590 }
1591
1592 #[test]
1593 fn test_viewpoint_config_single_viewpoint() {
1594 let config = ViewpointConfig {
1595 radius: 1.0,
1596 yaw_count: 1,
1597 pitch_angles_deg: vec![0.0],
1598 };
1599 assert_eq!(config.viewpoint_count(), 1);
1600 let viewpoints = generate_viewpoints(&config);
1601 assert_eq!(viewpoints.len(), 1);
1602 let pos = viewpoints[0].translation;
1604 assert!((pos.x).abs() < 0.001);
1605 assert!((pos.y).abs() < 0.001);
1606 assert!((pos.z - 1.0).abs() < 0.001);
1607 }
1608
1609 #[test]
1610 fn test_viewpoint_radius_scaling() {
1611 let config1 = ViewpointConfig {
1612 radius: 0.5,
1613 yaw_count: 4,
1614 pitch_angles_deg: vec![0.0],
1615 };
1616 let config2 = ViewpointConfig {
1617 radius: 2.0,
1618 yaw_count: 4,
1619 pitch_angles_deg: vec![0.0],
1620 };
1621
1622 let v1 = generate_viewpoints(&config1);
1623 let v2 = generate_viewpoints(&config2);
1624
1625 for (vp1, vp2) in v1.iter().zip(v2.iter()) {
1627 let ratio = vp2.translation.length() / vp1.translation.length();
1628 assert!((ratio - 4.0).abs() < 0.01); }
1630 }
1631
1632 #[test]
1633 fn test_camera_intrinsics_project_at_z_zero() {
1634 let intrinsics = CameraIntrinsics {
1635 focal_length: [100.0, 100.0],
1636 principal_point: [32.0, 32.0],
1637 image_size: [64, 64],
1638 };
1639
1640 let result = intrinsics.project(Vec3::new(1.0, 1.0, 0.0));
1642 assert!(result.is_none());
1643 }
1644
1645 #[test]
1646 fn test_camera_intrinsics_roundtrip() {
1647 let intrinsics = CameraIntrinsics {
1648 focal_length: [100.0, 100.0],
1649 principal_point: [32.0, 32.0],
1650 image_size: [64, 64],
1651 };
1652
1653 let original = Vec3::new(0.5, -0.3, 2.0);
1655 let projected = intrinsics.project(original).unwrap();
1656
1657 let unprojected = intrinsics.unproject(projected, original.z as f64);
1659
1660 assert!((unprojected[0] - original.x as f64).abs() < 0.001); assert!((unprojected[1] - original.y as f64).abs() < 0.001); assert!((unprojected[2] - original.z as f64).abs() < 0.001); }
1665
1666 #[test]
1667 fn test_render_output_empty() {
1668 let output = RenderOutput {
1669 rgba: vec![],
1670 depth: vec![],
1671 width: 0,
1672 height: 0,
1673 intrinsics: RenderConfig::tbp_default().intrinsics(),
1674 camera_transform: Transform::IDENTITY,
1675 object_rotation: ObjectRotation::identity(),
1676 };
1677
1678 assert_eq!(output.get_rgba(0, 0), None);
1680 assert_eq!(output.get_depth(0, 0), None);
1681 assert!(output.to_rgb_image().is_empty());
1682 assert!(output.to_depth_image().is_empty());
1683 }
1684
1685 #[test]
1686 fn test_render_output_1x1() {
1687 let output = RenderOutput {
1688 rgba: vec![128, 64, 32, 255],
1689 depth: vec![0.5],
1690 width: 1,
1691 height: 1,
1692 intrinsics: RenderConfig::tbp_default().intrinsics(),
1693 camera_transform: Transform::IDENTITY,
1694 object_rotation: ObjectRotation::identity(),
1695 };
1696
1697 assert_eq!(output.get_rgba(0, 0), Some([128, 64, 32, 255]));
1698 assert_eq!(output.get_depth(0, 0), Some(0.5));
1699 assert_eq!(output.get_rgb(0, 0), Some([128, 64, 32]));
1700
1701 let rgb_img = output.to_rgb_image();
1702 assert_eq!(rgb_img.len(), 1);
1703 assert_eq!(rgb_img[0].len(), 1);
1704 assert_eq!(rgb_img[0][0], [128, 64, 32]);
1705 }
1706
1707 #[test]
1708 fn test_render_config_high_res() {
1709 let config = RenderConfig::high_res();
1710 assert_eq!(config.width, 512);
1711 assert_eq!(config.height, 512);
1712
1713 let intrinsics = config.intrinsics();
1714 assert_eq!(intrinsics.image_size, [512, 512]);
1715 assert_eq!(intrinsics.principal_point, [256.0, 256.0]);
1716 }
1717
1718 #[test]
1719 fn test_render_config_zoom_affects_fov() {
1720 let base = RenderConfig {
1725 zoom: 2.0,
1726 ..RenderConfig::tbp_default()
1727 };
1728 let doubled = RenderConfig {
1729 zoom: 4.0,
1730 ..RenderConfig::tbp_default()
1731 };
1732
1733 assert!(doubled.fov_radians() < base.fov_radians());
1735
1736 let base_half_tan = (base.fov_radians() / 2.0).tan();
1738 let doubled_half_tan = (doubled.fov_radians() / 2.0).tan();
1739 assert!((base_half_tan / doubled_half_tan - 2.0).abs() < 1e-4);
1740 }
1741
1742 #[test]
1743 fn test_render_config_zoom_affects_intrinsics() {
1744 let a = RenderConfig {
1747 zoom: 2.0,
1748 ..RenderConfig::tbp_default()
1749 };
1750 let b = RenderConfig {
1751 zoom: 4.0,
1752 ..RenderConfig::tbp_default()
1753 };
1754
1755 let fx_a = a.intrinsics().focal_length[0];
1756 let fx_b = b.intrinsics().focal_length[0];
1757
1758 assert!(fx_b > fx_a);
1760
1761 assert!((fx_a / a.zoom as f64 - fx_b / b.zoom as f64).abs() < 1e-9);
1763 }
1764
1765 #[test]
1766 fn test_lighting_config_variants() {
1767 let default = LightingConfig::default();
1768 let bright = LightingConfig::bright();
1769 let soft = LightingConfig::soft();
1770 let unlit = LightingConfig::unlit();
1771
1772 assert!(bright.key_light_intensity > default.key_light_intensity);
1774
1775 assert_eq!(unlit.key_light_intensity, 0.0);
1777 assert_eq!(unlit.fill_light_intensity, 0.0);
1778 assert_eq!(unlit.ambient_brightness, 1.0);
1779
1780 assert!(soft.key_light_intensity < default.key_light_intensity);
1782 }
1783
1784 #[test]
1785 fn test_all_render_error_variants() {
1786 let errors = vec![
1787 RenderError::MeshNotFound("mesh.obj".to_string()),
1788 RenderError::TextureNotFound("texture.png".to_string()),
1789 RenderError::RenderFailed("GPU error".to_string()),
1790 RenderError::InvalidConfig("bad config".to_string()),
1791 ];
1792
1793 for err in errors {
1794 let msg = err.to_string();
1796 assert!(!msg.is_empty());
1797 }
1798 }
1799
1800 #[test]
1801 fn test_tbp_known_orientations_unique() {
1802 let orientations = ObjectRotation::tbp_known_orientations();
1803
1804 let quats: Vec<Quat> = orientations.iter().map(|r| r.to_quat()).collect();
1806
1807 for (i, q1) in quats.iter().enumerate() {
1808 for (j, q2) in quats.iter().enumerate() {
1809 if i != j {
1810 let dot = q1.dot(*q2).abs();
1812 assert!(
1813 dot < 0.999,
1814 "Orientations {} and {} produce same quaternion",
1815 i,
1816 j
1817 );
1818 }
1819 }
1820 }
1821 }
1822}