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::{self, DownloadOptions, Subset as YcbSubset, REPRESENTATIVE_OBJECTS, TEN_OBJECTS};
75
76pub mod ycb {
78 pub use ycbust::{download_ycb, DownloadOptions, Subset, REPRESENTATIVE_OBJECTS, TEN_OBJECTS};
79
80 use std::path::Path;
81
82 pub async fn download_models<P: AsRef<Path>>(
95 output_dir: P,
96 subset: Subset,
97 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
98 let options = DownloadOptions {
99 overwrite: false,
100 full: false,
101 show_progress: true,
102 delete_archives: true,
103 };
104 download_ycb(subset, output_dir.as_ref(), options).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 fn models_exist<P: AsRef<Path>>(output_dir: P) -> bool {
120 let path = output_dir.as_ref();
121 path.join("003_cracker_box/google_16k/textured.obj")
123 .exists()
124 }
125
126 pub fn object_mesh_path<P: AsRef<Path>>(output_dir: P, object_id: &str) -> std::path::PathBuf {
128 output_dir
129 .as_ref()
130 .join(object_id)
131 .join("google_16k")
132 .join("textured.obj")
133 }
134
135 pub fn object_texture_path<P: AsRef<Path>>(
137 output_dir: P,
138 object_id: &str,
139 ) -> std::path::PathBuf {
140 output_dir
141 .as_ref()
142 .join(object_id)
143 .join("google_16k")
144 .join("texture_map.png")
145 }
146}
147
148pub fn initialize() {
182 use std::sync::atomic::{AtomicBool, Ordering};
184 static INITIALIZED: AtomicBool = AtomicBool::new(false);
185
186 if !INITIALIZED.swap(true, Ordering::SeqCst) {
187 let config = backend::BackendConfig::new();
189 config.apply_env();
190 }
191}
192
193#[derive(Clone, Debug, PartialEq)]
196pub struct ObjectRotation {
197 pub pitch: f64,
199 pub yaw: f64,
201 pub roll: f64,
203}
204
205impl ObjectRotation {
206 pub fn new(pitch: f64, yaw: f64, roll: f64) -> Self {
208 Self { pitch, yaw, roll }
209 }
210
211 pub fn from_array(arr: [f64; 3]) -> Self {
213 Self {
214 pitch: arr[0],
215 yaw: arr[1],
216 roll: arr[2],
217 }
218 }
219
220 pub fn identity() -> Self {
222 Self::new(0.0, 0.0, 0.0)
223 }
224
225 pub fn tbp_benchmark_rotations() -> Vec<Self> {
228 vec![
229 Self::from_array([0.0, 0.0, 0.0]),
230 Self::from_array([0.0, 90.0, 0.0]),
231 Self::from_array([0.0, 180.0, 0.0]),
232 ]
233 }
234
235 pub fn tbp_known_orientations() -> Vec<Self> {
238 vec![
239 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]),
248 Self::from_array([45.0, 135.0, 0.0]),
249 Self::from_array([45.0, 225.0, 0.0]),
250 Self::from_array([45.0, 315.0, 0.0]),
251 Self::from_array([-45.0, 45.0, 0.0]),
252 Self::from_array([-45.0, 135.0, 0.0]),
253 Self::from_array([-45.0, 225.0, 0.0]),
254 Self::from_array([-45.0, 315.0, 0.0]),
255 ]
256 }
257
258 pub fn to_quat(&self) -> Quat {
260 Quat::from_euler(
261 EulerRot::XYZ,
262 (self.pitch as f32).to_radians(),
263 (self.yaw as f32).to_radians(),
264 (self.roll as f32).to_radians(),
265 )
266 }
267
268 pub fn to_transform(&self) -> Transform {
270 Transform::from_rotation(self.to_quat())
271 }
272}
273
274impl Default for ObjectRotation {
275 fn default() -> Self {
276 Self::identity()
277 }
278}
279
280#[derive(Clone, Debug)]
283pub struct ViewpointConfig {
284 pub radius: f32,
286 pub yaw_count: usize,
288 pub pitch_angles_deg: Vec<f32>,
290}
291
292impl Default for ViewpointConfig {
293 fn default() -> Self {
294 Self {
295 radius: 0.5,
296 yaw_count: 8,
297 pitch_angles_deg: vec![-30.0, 0.0, 30.0],
300 }
301 }
302}
303
304impl ViewpointConfig {
305 pub fn viewpoint_count(&self) -> usize {
307 self.yaw_count * self.pitch_angles_deg.len()
308 }
309}
310
311#[derive(Clone, Debug, Resource)]
313pub struct SensorConfig {
314 pub viewpoints: ViewpointConfig,
316 pub object_rotations: Vec<ObjectRotation>,
318 pub output_dir: String,
320 pub filename_pattern: String,
322}
323
324impl Default for SensorConfig {
325 fn default() -> Self {
326 Self {
327 viewpoints: ViewpointConfig::default(),
328 object_rotations: vec![ObjectRotation::identity()],
329 output_dir: ".".to_string(),
330 filename_pattern: "capture_{rot}_{view}.png".to_string(),
331 }
332 }
333}
334
335impl SensorConfig {
336 pub fn tbp_benchmark() -> Self {
338 Self {
339 viewpoints: ViewpointConfig::default(),
340 object_rotations: ObjectRotation::tbp_benchmark_rotations(),
341 output_dir: ".".to_string(),
342 filename_pattern: "capture_{rot}_{view}.png".to_string(),
343 }
344 }
345
346 pub fn tbp_full_training() -> Self {
348 Self {
349 viewpoints: ViewpointConfig::default(),
350 object_rotations: ObjectRotation::tbp_known_orientations(),
351 output_dir: ".".to_string(),
352 filename_pattern: "capture_{rot}_{view}.png".to_string(),
353 }
354 }
355
356 pub fn total_captures(&self) -> usize {
358 self.viewpoints.viewpoint_count() * self.object_rotations.len()
359 }
360}
361
362pub fn generate_viewpoints(config: &ViewpointConfig) -> Vec<Transform> {
369 let mut views = Vec::with_capacity(config.viewpoint_count());
370
371 for pitch_deg in &config.pitch_angles_deg {
372 let pitch = pitch_deg.to_radians();
373
374 for i in 0..config.yaw_count {
375 let yaw = (i as f32) * 2.0 * PI / (config.yaw_count as f32);
376
377 let x = config.radius * pitch.cos() * yaw.sin();
382 let y = config.radius * pitch.sin();
383 let z = config.radius * pitch.cos() * yaw.cos();
384
385 let transform = Transform::from_xyz(x, y, z).looking_at(Vec3::ZERO, Vec3::Y);
386 views.push(transform);
387 }
388 }
389 views
390}
391
392#[derive(Component)]
394pub struct CaptureTarget;
395
396#[derive(Component)]
398pub struct CaptureCamera;
399
400#[derive(Clone, Debug)]
408pub struct RenderConfig {
409 pub width: u32,
411 pub height: u32,
413 pub zoom: f32,
416 pub near_plane: f32,
418 pub far_plane: f32,
420 pub lighting: LightingConfig,
422}
423
424#[derive(Clone, Debug)]
428pub struct LightingConfig {
429 pub ambient_brightness: f32,
431 pub key_light_intensity: f32,
433 pub key_light_position: [f32; 3],
435 pub fill_light_intensity: f32,
437 pub fill_light_position: [f32; 3],
439 pub shadows_enabled: bool,
441}
442
443impl Default for LightingConfig {
444 fn default() -> Self {
445 Self {
446 ambient_brightness: 0.3,
447 key_light_intensity: 1500.0,
448 key_light_position: [4.0, 8.0, 4.0],
449 fill_light_intensity: 500.0,
450 fill_light_position: [-4.0, 2.0, -4.0],
451 shadows_enabled: false,
452 }
453 }
454}
455
456impl LightingConfig {
457 pub fn bright() -> Self {
459 Self {
460 ambient_brightness: 0.5,
461 key_light_intensity: 2000.0,
462 key_light_position: [4.0, 8.0, 4.0],
463 fill_light_intensity: 800.0,
464 fill_light_position: [-4.0, 2.0, -4.0],
465 shadows_enabled: false,
466 }
467 }
468
469 pub fn soft() -> Self {
471 Self {
472 ambient_brightness: 0.4,
473 key_light_intensity: 1000.0,
474 key_light_position: [3.0, 6.0, 3.0],
475 fill_light_intensity: 600.0,
476 fill_light_position: [-3.0, 3.0, -3.0],
477 shadows_enabled: false,
478 }
479 }
480
481 pub fn unlit() -> Self {
483 Self {
484 ambient_brightness: 1.0,
485 key_light_intensity: 0.0,
486 key_light_position: [0.0, 0.0, 0.0],
487 fill_light_intensity: 0.0,
488 fill_light_position: [0.0, 0.0, 0.0],
489 shadows_enabled: false,
490 }
491 }
492}
493
494impl Default for RenderConfig {
495 fn default() -> Self {
496 Self::tbp_default()
497 }
498}
499
500impl RenderConfig {
501 pub fn tbp_default() -> Self {
505 Self {
506 width: 64,
507 height: 64,
508 zoom: 1.0,
509 near_plane: 0.01,
510 far_plane: 10.0,
511 lighting: LightingConfig::default(),
512 }
513 }
514
515 pub fn preview() -> Self {
517 Self {
518 width: 256,
519 height: 256,
520 zoom: 1.0,
521 near_plane: 0.01,
522 far_plane: 10.0,
523 lighting: LightingConfig::default(),
524 }
525 }
526
527 pub fn high_res() -> Self {
529 Self {
530 width: 512,
531 height: 512,
532 zoom: 1.0,
533 near_plane: 0.01,
534 far_plane: 10.0,
535 lighting: LightingConfig::default(),
536 }
537 }
538
539 pub fn fov_radians(&self) -> f32 {
543 let base_fov_deg = 60.0_f32;
544 (base_fov_deg / self.zoom).to_radians()
545 }
546
547 pub fn intrinsics(&self) -> CameraIntrinsics {
552 let fov = self.fov_radians() as f64;
553 let fy = (self.height as f64 / 2.0) / (fov / 2.0).tan();
555 let fx = fy; CameraIntrinsics {
558 focal_length: [fx, fy],
559 principal_point: [self.width as f64 / 2.0, self.height as f64 / 2.0],
560 image_size: [self.width, self.height],
561 }
562 }
563}
564
565#[derive(Clone, Debug, PartialEq)]
570pub struct CameraIntrinsics {
571 pub focal_length: [f64; 2],
573 pub principal_point: [f64; 2],
575 pub image_size: [u32; 2],
577}
578
579impl CameraIntrinsics {
580 pub fn project(&self, point: Vec3) -> Option<[f64; 2]> {
582 if point.z <= 0.0 {
583 return None;
584 }
585 let x = (point.x as f64 / point.z as f64) * self.focal_length[0] + self.principal_point[0];
586 let y = (point.y as f64 / point.z as f64) * self.focal_length[1] + self.principal_point[1];
587 Some([x, y])
588 }
589
590 pub fn unproject(&self, pixel: [f64; 2], depth: f64) -> [f64; 3] {
592 let x = (pixel[0] - self.principal_point[0]) / self.focal_length[0] * depth;
593 let y = (pixel[1] - self.principal_point[1]) / self.focal_length[1] * depth;
594 [x, y, depth]
595 }
596}
597
598#[derive(Clone, Debug)]
600pub struct RenderOutput {
601 pub rgba: Vec<u8>,
603 pub depth: Vec<f64>,
607 pub width: u32,
609 pub height: u32,
611 pub intrinsics: CameraIntrinsics,
613 pub camera_transform: Transform,
615 pub object_rotation: ObjectRotation,
617}
618
619impl RenderOutput {
620 pub fn get_rgba(&self, x: u32, y: u32) -> Option<[u8; 4]> {
622 if x >= self.width || y >= self.height {
623 return None;
624 }
625 let idx = ((y * self.width + x) * 4) as usize;
626 Some([
627 self.rgba[idx],
628 self.rgba[idx + 1],
629 self.rgba[idx + 2],
630 self.rgba[idx + 3],
631 ])
632 }
633
634 pub fn get_depth(&self, x: u32, y: u32) -> Option<f64> {
636 if x >= self.width || y >= self.height {
637 return None;
638 }
639 let idx = (y * self.width + x) as usize;
640 Some(self.depth[idx])
641 }
642
643 pub fn get_rgb(&self, x: u32, y: u32) -> Option<[u8; 3]> {
645 self.get_rgba(x, y).map(|rgba| [rgba[0], rgba[1], rgba[2]])
646 }
647
648 pub fn to_rgb_image(&self) -> Vec<Vec<[u8; 3]>> {
650 let mut image = Vec::with_capacity(self.height as usize);
651 for y in 0..self.height {
652 let mut row = Vec::with_capacity(self.width as usize);
653 for x in 0..self.width {
654 row.push(self.get_rgb(x, y).unwrap_or([0, 0, 0]));
655 }
656 image.push(row);
657 }
658 image
659 }
660
661 pub fn to_depth_image(&self) -> Vec<Vec<f64>> {
663 let mut image = Vec::with_capacity(self.height as usize);
664 for y in 0..self.height {
665 let mut row = Vec::with_capacity(self.width as usize);
666 for x in 0..self.width {
667 row.push(self.get_depth(x, y).unwrap_or(0.0));
668 }
669 image.push(row);
670 }
671 image
672 }
673}
674
675#[derive(Debug, Clone)]
677pub enum RenderError {
678 MeshNotFound(String),
680 TextureNotFound(String),
682 FileNotFound { path: String, reason: String },
684 FileWriteFailed { path: String, reason: String },
686 DirectoryCreationFailed { path: String, reason: String },
688 RenderFailed(String),
690 InvalidConfig(String),
692 InvalidInput(String),
694 SerializationError(String),
696 DataParsingError(String),
698 RenderTimeout { duration_secs: u64 },
700}
701
702impl std::fmt::Display for RenderError {
703 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
704 match self {
705 RenderError::MeshNotFound(path) => write!(f, "Mesh not found: {}", path),
706 RenderError::TextureNotFound(path) => write!(f, "Texture not found: {}", path),
707 RenderError::FileNotFound { path, reason } => {
708 write!(f, "File not found at {}: {}", path, reason)
709 }
710 RenderError::FileWriteFailed { path, reason } => {
711 write!(f, "Failed to write file {}: {}", path, reason)
712 }
713 RenderError::DirectoryCreationFailed { path, reason } => {
714 write!(f, "Failed to create directory {}: {}", path, reason)
715 }
716 RenderError::RenderFailed(msg) => write!(f, "Render failed: {}", msg),
717 RenderError::InvalidConfig(msg) => write!(f, "Invalid config: {}", msg),
718 RenderError::InvalidInput(msg) => write!(f, "Invalid input: {}", msg),
719 RenderError::SerializationError(msg) => write!(f, "Serialization error: {}", msg),
720 RenderError::DataParsingError(msg) => write!(f, "Data parsing error: {}", msg),
721 RenderError::RenderTimeout { duration_secs } => {
722 write!(f, "Render timeout after {} seconds", duration_secs)
723 }
724 }
725 }
726}
727
728impl std::error::Error for RenderError {}
729
730pub fn render_to_buffer(
755 object_dir: &Path,
756 camera_transform: &Transform,
757 object_rotation: &ObjectRotation,
758 config: &RenderConfig,
759) -> Result<RenderOutput, RenderError> {
760 render::render_headless(object_dir, camera_transform, object_rotation, config)
762}
763
764pub fn render_all_viewpoints(
777 object_dir: &Path,
778 viewpoint_config: &ViewpointConfig,
779 rotations: &[ObjectRotation],
780 render_config: &RenderConfig,
781) -> Result<Vec<RenderOutput>, RenderError> {
782 let viewpoints = generate_viewpoints(viewpoint_config);
783 let mut outputs = Vec::with_capacity(viewpoints.len() * rotations.len());
784
785 for rotation in rotations {
786 for viewpoint in &viewpoints {
787 let output = render_to_buffer(object_dir, viewpoint, rotation, render_config)?;
788 outputs.push(output);
789 }
790 }
791
792 Ok(outputs)
793}
794
795pub fn render_to_buffer_cached(
861 object_dir: &Path,
862 camera_transform: &Transform,
863 object_rotation: &ObjectRotation,
864 config: &RenderConfig,
865 cache: &mut cache::ModelCache,
866) -> Result<RenderOutput, RenderError> {
867 let mesh_path = object_dir.join("google_16k/textured.obj");
868 let texture_path = object_dir.join("google_16k/texture_map.png");
869
870 cache.cache_scene(mesh_path.clone());
872 cache.cache_texture(texture_path.clone());
873
874 render::render_headless(object_dir, camera_transform, object_rotation, config)
876}
877
878pub fn render_to_files(
895 object_dir: &Path,
896 camera_transform: &Transform,
897 object_rotation: &ObjectRotation,
898 config: &RenderConfig,
899 rgba_path: &Path,
900 depth_path: &Path,
901) -> Result<(), RenderError> {
902 render::render_to_files(
903 object_dir,
904 camera_transform,
905 object_rotation,
906 config,
907 rgba_path,
908 depth_path,
909 )
910}
911
912pub use batch::{
914 BatchRenderConfig, BatchRenderError, BatchRenderOutput, BatchRenderRequest, BatchRenderer,
915 BatchState, RenderStatus,
916};
917
918pub fn create_batch_renderer(config: &BatchRenderConfig) -> Result<BatchRenderer, RenderError> {
936 Ok(BatchRenderer::new(config.clone()))
939}
940
941pub fn queue_render_request(
966 renderer: &mut BatchRenderer,
967 request: BatchRenderRequest,
968) -> Result<(), RenderError> {
969 renderer
970 .queue_request(request)
971 .map_err(|e| RenderError::RenderFailed(e.to_string()))
972}
973
974pub fn render_next_in_batch(
996 renderer: &mut BatchRenderer,
997 _timeout_ms: u32,
998) -> Result<Option<BatchRenderOutput>, RenderError> {
999 if let Some(request) = renderer.pending_requests.pop_front() {
1002 let output = render_to_buffer(
1003 &request.object_dir,
1004 &request.viewpoint,
1005 &request.object_rotation,
1006 &request.render_config,
1007 )?;
1008 let batch_output = BatchRenderOutput::from_render_output(request, output);
1009 renderer.completed_results.push(batch_output.clone());
1010 renderer.renders_processed += 1;
1011 Ok(Some(batch_output))
1012 } else {
1013 Ok(None)
1014 }
1015}
1016
1017pub fn render_batch(
1036 requests: Vec<BatchRenderRequest>,
1037 config: &BatchRenderConfig,
1038) -> Result<Vec<BatchRenderOutput>, RenderError> {
1039 let mut renderer = create_batch_renderer(config)?;
1040
1041 for request in requests {
1043 queue_render_request(&mut renderer, request)?;
1044 }
1045
1046 let mut results = Vec::new();
1048 while let Some(output) = render_next_in_batch(&mut renderer, config.frame_timeout_ms)? {
1049 results.push(output);
1050 }
1051
1052 Ok(results)
1053}
1054
1055pub use bevy::prelude::{Quat, Transform, Vec3};
1057
1058#[cfg(test)]
1059mod tests {
1060 use super::*;
1061
1062 #[test]
1063 fn test_object_rotation_identity() {
1064 let rot = ObjectRotation::identity();
1065 assert_eq!(rot.pitch, 0.0);
1066 assert_eq!(rot.yaw, 0.0);
1067 assert_eq!(rot.roll, 0.0);
1068 }
1069
1070 #[test]
1071 fn test_object_rotation_from_array() {
1072 let rot = ObjectRotation::from_array([10.0, 20.0, 30.0]);
1073 assert_eq!(rot.pitch, 10.0);
1074 assert_eq!(rot.yaw, 20.0);
1075 assert_eq!(rot.roll, 30.0);
1076 }
1077
1078 #[test]
1079 fn test_tbp_benchmark_rotations() {
1080 let rotations = ObjectRotation::tbp_benchmark_rotations();
1081 assert_eq!(rotations.len(), 3);
1082 assert_eq!(rotations[0], ObjectRotation::from_array([0.0, 0.0, 0.0]));
1083 assert_eq!(rotations[1], ObjectRotation::from_array([0.0, 90.0, 0.0]));
1084 assert_eq!(rotations[2], ObjectRotation::from_array([0.0, 180.0, 0.0]));
1085 }
1086
1087 #[test]
1088 fn test_tbp_known_orientations_count() {
1089 let orientations = ObjectRotation::tbp_known_orientations();
1090 assert_eq!(orientations.len(), 14);
1091 }
1092
1093 #[test]
1094 fn test_rotation_to_quat() {
1095 let rot = ObjectRotation::identity();
1096 let quat = rot.to_quat();
1097 assert!((quat.w - 1.0).abs() < 0.001);
1099 assert!(quat.x.abs() < 0.001);
1100 assert!(quat.y.abs() < 0.001);
1101 assert!(quat.z.abs() < 0.001);
1102 }
1103
1104 #[test]
1105 fn test_rotation_90_yaw() {
1106 let rot = ObjectRotation::new(0.0, 90.0, 0.0);
1107 let quat = rot.to_quat();
1108 assert!((quat.w - 0.707).abs() < 0.01);
1110 assert!((quat.y - 0.707).abs() < 0.01);
1111 }
1112
1113 #[test]
1114 fn test_viewpoint_config_default() {
1115 let config = ViewpointConfig::default();
1116 assert_eq!(config.radius, 0.5);
1117 assert_eq!(config.yaw_count, 8);
1118 assert_eq!(config.pitch_angles_deg.len(), 3);
1119 }
1120
1121 #[test]
1122 fn test_viewpoint_count() {
1123 let config = ViewpointConfig::default();
1124 assert_eq!(config.viewpoint_count(), 24); }
1126
1127 #[test]
1128 fn test_generate_viewpoints_count() {
1129 let config = ViewpointConfig::default();
1130 let viewpoints = generate_viewpoints(&config);
1131 assert_eq!(viewpoints.len(), 24);
1132 }
1133
1134 #[test]
1135 fn test_viewpoints_spherical_radius() {
1136 let config = ViewpointConfig::default();
1137 let viewpoints = generate_viewpoints(&config);
1138
1139 for (i, transform) in viewpoints.iter().enumerate() {
1140 let actual_radius = transform.translation.length();
1141 assert!(
1142 (actual_radius - config.radius).abs() < 0.001,
1143 "Viewpoint {} has incorrect radius: {} (expected {})",
1144 i,
1145 actual_radius,
1146 config.radius
1147 );
1148 }
1149 }
1150
1151 #[test]
1152 fn test_viewpoints_looking_at_origin() {
1153 let config = ViewpointConfig::default();
1154 let viewpoints = generate_viewpoints(&config);
1155
1156 for (i, transform) in viewpoints.iter().enumerate() {
1157 let forward = transform.forward();
1158 let to_origin = (Vec3::ZERO - transform.translation).normalize();
1159 let dot = forward.dot(to_origin);
1160 assert!(
1161 dot > 0.99,
1162 "Viewpoint {} not looking at origin, dot product: {}",
1163 i,
1164 dot
1165 );
1166 }
1167 }
1168
1169 #[test]
1170 fn test_sensor_config_default() {
1171 let config = SensorConfig::default();
1172 assert_eq!(config.object_rotations.len(), 1);
1173 assert_eq!(config.total_captures(), 24);
1174 }
1175
1176 #[test]
1177 fn test_sensor_config_tbp_benchmark() {
1178 let config = SensorConfig::tbp_benchmark();
1179 assert_eq!(config.object_rotations.len(), 3);
1180 assert_eq!(config.total_captures(), 72); }
1182
1183 #[test]
1184 fn test_sensor_config_tbp_full() {
1185 let config = SensorConfig::tbp_full_training();
1186 assert_eq!(config.object_rotations.len(), 14);
1187 assert_eq!(config.total_captures(), 336); }
1189
1190 #[test]
1191 fn test_ycb_representative_objects() {
1192 assert_eq!(crate::ycb::REPRESENTATIVE_OBJECTS.len(), 3);
1194 assert!(crate::ycb::REPRESENTATIVE_OBJECTS.contains(&"003_cracker_box"));
1195 }
1196
1197 #[test]
1198 fn test_ycb_ten_objects() {
1199 assert_eq!(crate::ycb::TEN_OBJECTS.len(), 10);
1201 }
1202
1203 #[test]
1204 fn test_ycb_object_mesh_path() {
1205 let path = crate::ycb::object_mesh_path("/tmp/ycb", "003_cracker_box");
1206 assert_eq!(
1207 path.to_string_lossy(),
1208 "/tmp/ycb/003_cracker_box/google_16k/textured.obj"
1209 );
1210 }
1211
1212 #[test]
1213 fn test_ycb_object_texture_path() {
1214 let path = crate::ycb::object_texture_path("/tmp/ycb", "003_cracker_box");
1215 assert_eq!(
1216 path.to_string_lossy(),
1217 "/tmp/ycb/003_cracker_box/google_16k/texture_map.png"
1218 );
1219 }
1220
1221 #[test]
1226 fn test_render_config_tbp_default() {
1227 let config = RenderConfig::tbp_default();
1228 assert_eq!(config.width, 64);
1229 assert_eq!(config.height, 64);
1230 assert_eq!(config.zoom, 1.0);
1231 assert_eq!(config.near_plane, 0.01);
1232 assert_eq!(config.far_plane, 10.0);
1233 }
1234
1235 #[test]
1236 fn test_render_config_preview() {
1237 let config = RenderConfig::preview();
1238 assert_eq!(config.width, 256);
1239 assert_eq!(config.height, 256);
1240 }
1241
1242 #[test]
1243 fn test_render_config_default_is_tbp() {
1244 let default = RenderConfig::default();
1245 let tbp = RenderConfig::tbp_default();
1246 assert_eq!(default.width, tbp.width);
1247 assert_eq!(default.height, tbp.height);
1248 }
1249
1250 #[test]
1251 fn test_render_config_fov() {
1252 let config = RenderConfig::tbp_default();
1253 let fov = config.fov_radians();
1254 assert!((fov - 1.047).abs() < 0.01);
1256
1257 let zoomed = RenderConfig {
1259 zoom: 2.0,
1260 ..config
1261 };
1262 assert!(zoomed.fov_radians() < fov);
1263 }
1264
1265 #[test]
1266 fn test_render_config_intrinsics() {
1267 let config = RenderConfig::tbp_default();
1268 let intrinsics = config.intrinsics();
1269
1270 assert_eq!(intrinsics.image_size, [64, 64]);
1271 assert_eq!(intrinsics.principal_point, [32.0, 32.0]);
1272 assert!(intrinsics.focal_length[0] > 0.0);
1274 assert!(intrinsics.focal_length[1] > 0.0);
1275 assert!((intrinsics.focal_length[0] - 55.4).abs() < 1.0);
1277 }
1278
1279 #[test]
1280 fn test_camera_intrinsics_project() {
1281 let intrinsics = CameraIntrinsics {
1282 focal_length: [100.0, 100.0],
1283 principal_point: [32.0, 32.0],
1284 image_size: [64, 64],
1285 };
1286
1287 let center = intrinsics.project(Vec3::new(0.0, 0.0, 1.0));
1289 assert!(center.is_some());
1290 let [x, y] = center.unwrap();
1291 assert!((x - 32.0).abs() < 0.001);
1292 assert!((y - 32.0).abs() < 0.001);
1293
1294 let behind = intrinsics.project(Vec3::new(0.0, 0.0, -1.0));
1296 assert!(behind.is_none());
1297 }
1298
1299 #[test]
1300 fn test_camera_intrinsics_unproject() {
1301 let intrinsics = CameraIntrinsics {
1302 focal_length: [100.0, 100.0],
1303 principal_point: [32.0, 32.0],
1304 image_size: [64, 64],
1305 };
1306
1307 let point = intrinsics.unproject([32.0, 32.0], 1.0);
1309 assert!((point[0]).abs() < 0.001); assert!((point[1]).abs() < 0.001); assert!((point[2] - 1.0).abs() < 0.001); }
1313
1314 #[test]
1315 fn test_render_output_get_rgba() {
1316 let output = RenderOutput {
1317 rgba: vec![
1318 255, 0, 0, 255, 0, 255, 0, 255, 0, 0, 255, 255, 255, 255, 255, 255,
1319 ],
1320 depth: vec![1.0, 2.0, 3.0, 4.0],
1321 width: 2,
1322 height: 2,
1323 intrinsics: RenderConfig::tbp_default().intrinsics(),
1324 camera_transform: Transform::IDENTITY,
1325 object_rotation: ObjectRotation::identity(),
1326 };
1327
1328 assert_eq!(output.get_rgba(0, 0), Some([255, 0, 0, 255]));
1330 assert_eq!(output.get_rgba(1, 0), Some([0, 255, 0, 255]));
1332 assert_eq!(output.get_rgba(0, 1), Some([0, 0, 255, 255]));
1334 assert_eq!(output.get_rgba(1, 1), Some([255, 255, 255, 255]));
1336 assert_eq!(output.get_rgba(2, 0), None);
1338 }
1339
1340 #[test]
1341 fn test_render_output_get_depth() {
1342 let output = RenderOutput {
1343 rgba: vec![0u8; 16],
1344 depth: vec![1.0, 2.0, 3.0, 4.0],
1345 width: 2,
1346 height: 2,
1347 intrinsics: RenderConfig::tbp_default().intrinsics(),
1348 camera_transform: Transform::IDENTITY,
1349 object_rotation: ObjectRotation::identity(),
1350 };
1351
1352 assert_eq!(output.get_depth(0, 0), Some(1.0));
1353 assert_eq!(output.get_depth(1, 0), Some(2.0));
1354 assert_eq!(output.get_depth(0, 1), Some(3.0));
1355 assert_eq!(output.get_depth(1, 1), Some(4.0));
1356 assert_eq!(output.get_depth(2, 0), None);
1357 }
1358
1359 #[test]
1360 fn test_render_output_to_rgb_image() {
1361 let output = RenderOutput {
1362 rgba: vec![
1363 255, 0, 0, 255, 0, 255, 0, 255, 0, 0, 255, 255, 255, 255, 255, 255,
1364 ],
1365 depth: vec![1.0, 2.0, 3.0, 4.0],
1366 width: 2,
1367 height: 2,
1368 intrinsics: RenderConfig::tbp_default().intrinsics(),
1369 camera_transform: Transform::IDENTITY,
1370 object_rotation: ObjectRotation::identity(),
1371 };
1372
1373 let image = output.to_rgb_image();
1374 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]); }
1381
1382 #[test]
1383 fn test_render_output_to_depth_image() {
1384 let output = RenderOutput {
1385 rgba: vec![0u8; 16],
1386 depth: vec![1.0, 2.0, 3.0, 4.0],
1387 width: 2,
1388 height: 2,
1389 intrinsics: RenderConfig::tbp_default().intrinsics(),
1390 camera_transform: Transform::IDENTITY,
1391 object_rotation: ObjectRotation::identity(),
1392 };
1393
1394 let depth_image = output.to_depth_image();
1395 assert_eq!(depth_image.len(), 2);
1396 assert_eq!(depth_image[0], vec![1.0, 2.0]);
1397 assert_eq!(depth_image[1], vec![3.0, 4.0]);
1398 }
1399
1400 #[test]
1401 fn test_render_error_display() {
1402 let err = RenderError::MeshNotFound("/path/to/mesh.obj".to_string());
1403 assert!(err.to_string().contains("Mesh not found"));
1404 assert!(err.to_string().contains("/path/to/mesh.obj"));
1405 }
1406
1407 #[test]
1412 fn test_object_rotation_extreme_angles() {
1413 let rot = ObjectRotation::new(450.0, -720.0, 1080.0);
1415 let quat = rot.to_quat();
1416 assert!((quat.length() - 1.0).abs() < 0.001);
1418 }
1419
1420 #[test]
1421 fn test_object_rotation_to_transform() {
1422 let rot = ObjectRotation::new(45.0, 90.0, 0.0);
1423 let transform = rot.to_transform();
1424 assert_eq!(transform.translation, Vec3::ZERO);
1426 assert!(transform.rotation != Quat::IDENTITY);
1428 }
1429
1430 #[test]
1431 fn test_viewpoint_config_single_viewpoint() {
1432 let config = ViewpointConfig {
1433 radius: 1.0,
1434 yaw_count: 1,
1435 pitch_angles_deg: vec![0.0],
1436 };
1437 assert_eq!(config.viewpoint_count(), 1);
1438 let viewpoints = generate_viewpoints(&config);
1439 assert_eq!(viewpoints.len(), 1);
1440 let pos = viewpoints[0].translation;
1442 assert!((pos.x).abs() < 0.001);
1443 assert!((pos.y).abs() < 0.001);
1444 assert!((pos.z - 1.0).abs() < 0.001);
1445 }
1446
1447 #[test]
1448 fn test_viewpoint_radius_scaling() {
1449 let config1 = ViewpointConfig {
1450 radius: 0.5,
1451 yaw_count: 4,
1452 pitch_angles_deg: vec![0.0],
1453 };
1454 let config2 = ViewpointConfig {
1455 radius: 2.0,
1456 yaw_count: 4,
1457 pitch_angles_deg: vec![0.0],
1458 };
1459
1460 let v1 = generate_viewpoints(&config1);
1461 let v2 = generate_viewpoints(&config2);
1462
1463 for (vp1, vp2) in v1.iter().zip(v2.iter()) {
1465 let ratio = vp2.translation.length() / vp1.translation.length();
1466 assert!((ratio - 4.0).abs() < 0.01); }
1468 }
1469
1470 #[test]
1471 fn test_camera_intrinsics_project_at_z_zero() {
1472 let intrinsics = CameraIntrinsics {
1473 focal_length: [100.0, 100.0],
1474 principal_point: [32.0, 32.0],
1475 image_size: [64, 64],
1476 };
1477
1478 let result = intrinsics.project(Vec3::new(1.0, 1.0, 0.0));
1480 assert!(result.is_none());
1481 }
1482
1483 #[test]
1484 fn test_camera_intrinsics_roundtrip() {
1485 let intrinsics = CameraIntrinsics {
1486 focal_length: [100.0, 100.0],
1487 principal_point: [32.0, 32.0],
1488 image_size: [64, 64],
1489 };
1490
1491 let original = Vec3::new(0.5, -0.3, 2.0);
1493 let projected = intrinsics.project(original).unwrap();
1494
1495 let unprojected = intrinsics.unproject(projected, original.z as f64);
1497
1498 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); }
1503
1504 #[test]
1505 fn test_render_output_empty() {
1506 let output = RenderOutput {
1507 rgba: vec![],
1508 depth: vec![],
1509 width: 0,
1510 height: 0,
1511 intrinsics: RenderConfig::tbp_default().intrinsics(),
1512 camera_transform: Transform::IDENTITY,
1513 object_rotation: ObjectRotation::identity(),
1514 };
1515
1516 assert_eq!(output.get_rgba(0, 0), None);
1518 assert_eq!(output.get_depth(0, 0), None);
1519 assert!(output.to_rgb_image().is_empty());
1520 assert!(output.to_depth_image().is_empty());
1521 }
1522
1523 #[test]
1524 fn test_render_output_1x1() {
1525 let output = RenderOutput {
1526 rgba: vec![128, 64, 32, 255],
1527 depth: vec![0.5],
1528 width: 1,
1529 height: 1,
1530 intrinsics: RenderConfig::tbp_default().intrinsics(),
1531 camera_transform: Transform::IDENTITY,
1532 object_rotation: ObjectRotation::identity(),
1533 };
1534
1535 assert_eq!(output.get_rgba(0, 0), Some([128, 64, 32, 255]));
1536 assert_eq!(output.get_depth(0, 0), Some(0.5));
1537 assert_eq!(output.get_rgb(0, 0), Some([128, 64, 32]));
1538
1539 let rgb_img = output.to_rgb_image();
1540 assert_eq!(rgb_img.len(), 1);
1541 assert_eq!(rgb_img[0].len(), 1);
1542 assert_eq!(rgb_img[0][0], [128, 64, 32]);
1543 }
1544
1545 #[test]
1546 fn test_render_config_high_res() {
1547 let config = RenderConfig::high_res();
1548 assert_eq!(config.width, 512);
1549 assert_eq!(config.height, 512);
1550
1551 let intrinsics = config.intrinsics();
1552 assert_eq!(intrinsics.image_size, [512, 512]);
1553 assert_eq!(intrinsics.principal_point, [256.0, 256.0]);
1554 }
1555
1556 #[test]
1557 fn test_render_config_zoom_affects_fov() {
1558 let base = RenderConfig::tbp_default();
1559 let zoomed = RenderConfig {
1560 zoom: 2.0,
1561 ..base.clone()
1562 };
1563
1564 assert!(zoomed.fov_radians() < base.fov_radians());
1566 assert!((zoomed.fov_radians() - base.fov_radians() / 2.0).abs() < 0.01);
1568 }
1569
1570 #[test]
1571 fn test_render_config_zoom_affects_intrinsics() {
1572 let base = RenderConfig::tbp_default();
1573 let zoomed = RenderConfig {
1574 zoom: 2.0,
1575 ..base.clone()
1576 };
1577
1578 let base_intrinsics = base.intrinsics();
1580 let zoomed_intrinsics = zoomed.intrinsics();
1581
1582 assert!(zoomed_intrinsics.focal_length[0] > base_intrinsics.focal_length[0]);
1583 }
1584
1585 #[test]
1586 fn test_lighting_config_variants() {
1587 let default = LightingConfig::default();
1588 let bright = LightingConfig::bright();
1589 let soft = LightingConfig::soft();
1590 let unlit = LightingConfig::unlit();
1591
1592 assert!(bright.key_light_intensity > default.key_light_intensity);
1594
1595 assert_eq!(unlit.key_light_intensity, 0.0);
1597 assert_eq!(unlit.fill_light_intensity, 0.0);
1598 assert_eq!(unlit.ambient_brightness, 1.0);
1599
1600 assert!(soft.key_light_intensity < default.key_light_intensity);
1602 }
1603
1604 #[test]
1605 fn test_all_render_error_variants() {
1606 let errors = vec![
1607 RenderError::MeshNotFound("mesh.obj".to_string()),
1608 RenderError::TextureNotFound("texture.png".to_string()),
1609 RenderError::RenderFailed("GPU error".to_string()),
1610 RenderError::InvalidConfig("bad config".to_string()),
1611 ];
1612
1613 for err in errors {
1614 let msg = err.to_string();
1616 assert!(!msg.is_empty());
1617 }
1618 }
1619
1620 #[test]
1621 fn test_tbp_known_orientations_unique() {
1622 let orientations = ObjectRotation::tbp_known_orientations();
1623
1624 let quats: Vec<Quat> = orientations.iter().map(|r| r.to_quat()).collect();
1626
1627 for (i, q1) in quats.iter().enumerate() {
1628 for (j, q2) in quats.iter().enumerate() {
1629 if i != j {
1630 let dot = q1.dot(*q2).abs();
1632 assert!(
1633 dot < 0.999,
1634 "Orientations {} and {} produce same quaternion",
1635 i,
1636 j
1637 );
1638 }
1639 }
1640 }
1641 }
1642}