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 missing_objects<P: AsRef<Path>>(output_dir: P, object_ids: &[&str]) -> Vec<String> {
134 ycbust::validate_objects(output_dir.as_ref(), object_ids)
135 .into_iter()
136 .filter(|validation| !validation.is_complete())
137 .map(|validation| validation.name)
138 .collect()
139 }
140
141 pub fn objects_exist<P: AsRef<Path>>(output_dir: P, object_ids: &[&str]) -> bool {
143 missing_objects(output_dir, object_ids).is_empty()
144 }
145
146 pub fn models_exist<P: AsRef<Path>>(output_dir: P) -> bool {
148 objects_exist(output_dir, REPRESENTATIVE_OBJECTS)
149 }
150
151 pub fn object_mesh_path<P: AsRef<Path>>(output_dir: P, object_id: &str) -> std::path::PathBuf {
153 ycbust::object_mesh_path(output_dir.as_ref(), object_id)
154 }
155
156 pub fn object_texture_path<P: AsRef<Path>>(
158 output_dir: P,
159 object_id: &str,
160 ) -> std::path::PathBuf {
161 ycbust::object_texture_path(output_dir.as_ref(), object_id)
162 }
163}
164
165pub fn initialize() {
199 use std::sync::atomic::{AtomicBool, Ordering};
201 static INITIALIZED: AtomicBool = AtomicBool::new(false);
202
203 if !INITIALIZED.swap(true, Ordering::SeqCst) {
204 let config = backend::BackendConfig::new();
206 config.apply_env();
207 }
208}
209
210#[derive(Clone, Debug, PartialEq)]
213pub struct ObjectRotation {
214 pub pitch: f64,
216 pub yaw: f64,
218 pub roll: f64,
220}
221
222impl ObjectRotation {
223 pub fn new(pitch: f64, yaw: f64, roll: f64) -> Self {
225 Self { pitch, yaw, roll }
226 }
227
228 pub fn from_array(arr: [f64; 3]) -> Self {
230 Self {
231 pitch: arr[0],
232 yaw: arr[1],
233 roll: arr[2],
234 }
235 }
236
237 pub fn identity() -> Self {
239 Self::new(0.0, 0.0, 0.0)
240 }
241
242 pub fn tbp_benchmark_rotations() -> Vec<Self> {
245 vec![
246 Self::from_array([0.0, 0.0, 0.0]),
247 Self::from_array([0.0, 90.0, 0.0]),
248 Self::from_array([0.0, 180.0, 0.0]),
249 ]
250 }
251
252 pub fn tbp_known_orientations() -> Vec<Self> {
255 vec![
256 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]),
265 Self::from_array([45.0, 135.0, 0.0]),
266 Self::from_array([45.0, 225.0, 0.0]),
267 Self::from_array([45.0, 315.0, 0.0]),
268 Self::from_array([-45.0, 45.0, 0.0]),
269 Self::from_array([-45.0, 135.0, 0.0]),
270 Self::from_array([-45.0, 225.0, 0.0]),
271 Self::from_array([-45.0, 315.0, 0.0]),
272 ]
273 }
274
275 pub fn to_quat(&self) -> Quat {
277 Quat::from_euler(
278 EulerRot::XYZ,
279 (self.pitch as f32).to_radians(),
280 (self.yaw as f32).to_radians(),
281 (self.roll as f32).to_radians(),
282 )
283 }
284
285 pub fn to_transform(&self) -> Transform {
287 Transform::from_rotation(self.to_quat())
288 }
289}
290
291impl Default for ObjectRotation {
292 fn default() -> Self {
293 Self::identity()
294 }
295}
296
297#[derive(Clone, Debug)]
300pub struct ViewpointConfig {
301 pub radius: f32,
303 pub yaw_count: usize,
305 pub pitch_angles_deg: Vec<f32>,
307}
308
309impl Default for ViewpointConfig {
310 fn default() -> Self {
311 Self {
312 radius: 0.5,
313 yaw_count: 8,
314 pitch_angles_deg: vec![-30.0, 0.0, 30.0],
317 }
318 }
319}
320
321impl ViewpointConfig {
322 pub fn viewpoint_count(&self) -> usize {
324 self.yaw_count * self.pitch_angles_deg.len()
325 }
326}
327
328#[derive(Clone, Debug, Resource)]
330pub struct SensorConfig {
331 pub viewpoints: ViewpointConfig,
333 pub object_rotations: Vec<ObjectRotation>,
335 pub output_dir: String,
337 pub filename_pattern: String,
339}
340
341impl Default for SensorConfig {
342 fn default() -> Self {
343 Self {
344 viewpoints: ViewpointConfig::default(),
345 object_rotations: vec![ObjectRotation::identity()],
346 output_dir: ".".to_string(),
347 filename_pattern: "capture_{rot}_{view}.png".to_string(),
348 }
349 }
350}
351
352impl SensorConfig {
353 pub fn tbp_benchmark() -> Self {
355 Self {
356 viewpoints: ViewpointConfig::default(),
357 object_rotations: ObjectRotation::tbp_benchmark_rotations(),
358 output_dir: ".".to_string(),
359 filename_pattern: "capture_{rot}_{view}.png".to_string(),
360 }
361 }
362
363 pub fn tbp_full_training() -> Self {
365 Self {
366 viewpoints: ViewpointConfig::default(),
367 object_rotations: ObjectRotation::tbp_known_orientations(),
368 output_dir: ".".to_string(),
369 filename_pattern: "capture_{rot}_{view}.png".to_string(),
370 }
371 }
372
373 pub fn total_captures(&self) -> usize {
375 self.viewpoints.viewpoint_count() * self.object_rotations.len()
376 }
377}
378
379pub fn generate_viewpoints(config: &ViewpointConfig) -> Vec<Transform> {
386 let mut views = Vec::with_capacity(config.viewpoint_count());
387
388 for pitch_deg in &config.pitch_angles_deg {
389 let pitch = pitch_deg.to_radians();
390
391 for i in 0..config.yaw_count {
392 let yaw = (i as f32) * 2.0 * PI / (config.yaw_count as f32);
393
394 let x = config.radius * pitch.cos() * yaw.sin();
399 let y = config.radius * pitch.sin();
400 let z = config.radius * pitch.cos() * yaw.cos();
401
402 let transform = Transform::from_xyz(x, y, z).looking_at(Vec3::ZERO, Vec3::Y);
403 views.push(transform);
404 }
405 }
406 views
407}
408
409#[derive(Component)]
411pub struct CaptureTarget;
412
413#[derive(Component)]
415pub struct CaptureCamera;
416
417#[derive(Clone, Debug, PartialEq)]
425pub struct RenderConfig {
426 pub width: u32,
428 pub height: u32,
430 pub zoom: f32,
433 pub near_plane: f32,
435 pub far_plane: f32,
437 pub lighting: LightingConfig,
439}
440
441#[derive(Clone, Debug, PartialEq)]
445pub struct LightingConfig {
446 pub ambient_brightness: f32,
448 pub key_light_intensity: f32,
450 pub key_light_position: [f32; 3],
452 pub fill_light_intensity: f32,
454 pub fill_light_position: [f32; 3],
456 pub shadows_enabled: bool,
458}
459
460impl Default for LightingConfig {
461 fn default() -> Self {
462 Self {
463 ambient_brightness: 0.3,
464 key_light_intensity: 1500.0,
465 key_light_position: [4.0, 8.0, 4.0],
466 fill_light_intensity: 500.0,
467 fill_light_position: [-4.0, 2.0, -4.0],
468 shadows_enabled: false,
469 }
470 }
471}
472
473impl LightingConfig {
474 pub fn bright() -> Self {
476 Self {
477 ambient_brightness: 0.5,
478 key_light_intensity: 2000.0,
479 key_light_position: [4.0, 8.0, 4.0],
480 fill_light_intensity: 800.0,
481 fill_light_position: [-4.0, 2.0, -4.0],
482 shadows_enabled: false,
483 }
484 }
485
486 pub fn soft() -> Self {
488 Self {
489 ambient_brightness: 0.4,
490 key_light_intensity: 1000.0,
491 key_light_position: [3.0, 6.0, 3.0],
492 fill_light_intensity: 600.0,
493 fill_light_position: [-3.0, 3.0, -3.0],
494 shadows_enabled: false,
495 }
496 }
497
498 pub fn unlit() -> Self {
500 Self {
501 ambient_brightness: 1.0,
502 key_light_intensity: 0.0,
503 key_light_position: [0.0, 0.0, 0.0],
504 fill_light_intensity: 0.0,
505 fill_light_position: [0.0, 0.0, 0.0],
506 shadows_enabled: false,
507 }
508 }
509}
510
511impl Default for RenderConfig {
512 fn default() -> Self {
513 Self::tbp_default()
514 }
515}
516
517impl RenderConfig {
518 pub fn tbp_default() -> Self {
526 Self {
527 width: 64,
528 height: 64,
529 zoom: 4.0,
530 near_plane: 0.01,
531 far_plane: 10.0,
532 lighting: LightingConfig::default(),
533 }
534 }
535
536 pub fn preview() -> Self {
538 Self {
539 width: 256,
540 height: 256,
541 zoom: 1.0,
542 near_plane: 0.01,
543 far_plane: 10.0,
544 lighting: LightingConfig::default(),
545 }
546 }
547
548 pub fn high_res() -> Self {
550 Self {
551 width: 512,
552 height: 512,
553 zoom: 1.0,
554 near_plane: 0.01,
555 far_plane: 10.0,
556 lighting: LightingConfig::default(),
557 }
558 }
559
560 pub fn fov_radians(&self) -> f32 {
567 let base_hfov_rad = 90.0_f32.to_radians();
568 let half_tan = (base_hfov_rad / 2.0).tan() / self.zoom;
569 2.0 * half_tan.atan()
570 }
571
572 pub fn intrinsics(&self) -> CameraIntrinsics {
580 self.intrinsics_for_size(self.width, self.height)
581 }
582
583 pub fn intrinsics_for_size(&self, width: u32, height: u32) -> CameraIntrinsics {
588 let base_hfov_rad = 90.0_f64.to_radians();
589 let fx_norm = (base_hfov_rad / 2.0).tan() / self.zoom as f64;
591 let fx = (width as f64 / 2.0) / fx_norm;
593 let fy = fx; CameraIntrinsics {
596 focal_length: [fx, fy],
597 principal_point: [width as f64 / 2.0, height as f64 / 2.0],
598 image_size: [width, height],
599 }
600 }
601}
602
603#[derive(Clone, Debug, PartialEq)]
608pub struct CameraIntrinsics {
609 pub focal_length: [f64; 2],
611 pub principal_point: [f64; 2],
613 pub image_size: [u32; 2],
615}
616
617impl CameraIntrinsics {
618 pub fn project(&self, point: Vec3) -> Option<[f64; 2]> {
620 if point.z <= 0.0 {
621 return None;
622 }
623 let x = (point.x as f64 / point.z as f64) * self.focal_length[0] + self.principal_point[0];
624 let y = (point.y as f64 / point.z as f64) * self.focal_length[1] + self.principal_point[1];
625 Some([x, y])
626 }
627
628 pub fn unproject(&self, pixel: [f64; 2], depth: f64) -> [f64; 3] {
630 let x = (pixel[0] - self.principal_point[0]) / self.focal_length[0] * depth;
631 let y = (pixel[1] - self.principal_point[1]) / self.focal_length[1] * depth;
632 [x, y, depth]
633 }
634}
635
636#[derive(Clone, Debug)]
638pub struct RenderOutput {
639 pub rgba: Vec<u8>,
641 pub depth: Vec<f64>,
645 pub width: u32,
647 pub height: u32,
649 pub intrinsics: CameraIntrinsics,
651 pub camera_transform: Transform,
653 pub object_rotation: ObjectRotation,
655}
656
657impl RenderOutput {
658 pub fn get_rgba(&self, x: u32, y: u32) -> Option<[u8; 4]> {
660 if x >= self.width || y >= self.height {
661 return None;
662 }
663 let idx = ((y * self.width + x) * 4) as usize;
664 Some([
665 self.rgba[idx],
666 self.rgba[idx + 1],
667 self.rgba[idx + 2],
668 self.rgba[idx + 3],
669 ])
670 }
671
672 pub fn get_depth(&self, x: u32, y: u32) -> Option<f64> {
674 if x >= self.width || y >= self.height {
675 return None;
676 }
677 let idx = (y * self.width + x) as usize;
678 Some(self.depth[idx])
679 }
680
681 pub fn get_rgb(&self, x: u32, y: u32) -> Option<[u8; 3]> {
683 self.get_rgba(x, y).map(|rgba| [rgba[0], rgba[1], rgba[2]])
684 }
685
686 pub fn to_rgb_image(&self) -> Vec<Vec<[u8; 3]>> {
688 let mut image = Vec::with_capacity(self.height as usize);
689 for y in 0..self.height {
690 let mut row = Vec::with_capacity(self.width as usize);
691 for x in 0..self.width {
692 row.push(self.get_rgb(x, y).unwrap_or([0, 0, 0]));
693 }
694 image.push(row);
695 }
696 image
697 }
698
699 pub fn to_depth_image(&self) -> Vec<Vec<f64>> {
701 let mut image = Vec::with_capacity(self.height as usize);
702 for y in 0..self.height {
703 let mut row = Vec::with_capacity(self.width as usize);
704 for x in 0..self.width {
705 row.push(self.get_depth(x, y).unwrap_or(0.0));
706 }
707 image.push(row);
708 }
709 image
710 }
711}
712
713#[derive(Debug, Clone)]
715pub enum RenderError {
716 MeshNotFound(String),
718 TextureNotFound(String),
720 FileNotFound { path: String, reason: String },
722 FileWriteFailed { path: String, reason: String },
724 DirectoryCreationFailed { path: String, reason: String },
726 RenderFailed(String),
728 InvalidConfig(String),
730 InvalidInput(String),
732 SerializationError(String),
734 DataParsingError(String),
736 RenderTimeout { duration_secs: u64 },
738}
739
740impl std::fmt::Display for RenderError {
741 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
742 match self {
743 RenderError::MeshNotFound(path) => write!(f, "Mesh not found: {}", path),
744 RenderError::TextureNotFound(path) => write!(f, "Texture not found: {}", path),
745 RenderError::FileNotFound { path, reason } => {
746 write!(f, "File not found at {}: {}", path, reason)
747 }
748 RenderError::FileWriteFailed { path, reason } => {
749 write!(f, "Failed to write file {}: {}", path, reason)
750 }
751 RenderError::DirectoryCreationFailed { path, reason } => {
752 write!(f, "Failed to create directory {}: {}", path, reason)
753 }
754 RenderError::RenderFailed(msg) => write!(f, "Render failed: {}", msg),
755 RenderError::InvalidConfig(msg) => write!(f, "Invalid config: {}", msg),
756 RenderError::InvalidInput(msg) => write!(f, "Invalid input: {}", msg),
757 RenderError::SerializationError(msg) => write!(f, "Serialization error: {}", msg),
758 RenderError::DataParsingError(msg) => write!(f, "Data parsing error: {}", msg),
759 RenderError::RenderTimeout { duration_secs } => {
760 write!(f, "Render timeout after {} seconds", duration_secs)
761 }
762 }
763 }
764}
765
766impl std::error::Error for RenderError {}
767
768pub fn render_to_buffer(
793 object_dir: &Path,
794 camera_transform: &Transform,
795 object_rotation: &ObjectRotation,
796 config: &RenderConfig,
797) -> Result<RenderOutput, RenderError> {
798 render::render_headless(object_dir, camera_transform, object_rotation, config)
800}
801
802pub fn render_all_viewpoints(
815 object_dir: &Path,
816 viewpoint_config: &ViewpointConfig,
817 rotations: &[ObjectRotation],
818 render_config: &RenderConfig,
819) -> Result<Vec<RenderOutput>, RenderError> {
820 let viewpoints = generate_viewpoints(viewpoint_config);
821 let mut outputs = Vec::with_capacity(viewpoints.len() * rotations.len());
822
823 for rotation in rotations {
824 for viewpoint in &viewpoints {
825 let output = render_to_buffer(object_dir, viewpoint, rotation, render_config)?;
826 outputs.push(output);
827 }
828 }
829
830 Ok(outputs)
831}
832
833pub fn render_to_buffer_cached(
903 object_dir: &Path,
904 camera_transform: &Transform,
905 object_rotation: &ObjectRotation,
906 config: &RenderConfig,
907 cache: &mut cache::ModelCache,
908) -> Result<RenderOutput, RenderError> {
909 let mesh_path = object_dir.join("google_16k/textured.obj");
910 let texture_path = object_dir.join("google_16k/texture_map.png");
911
912 cache.cache_scene(mesh_path.clone());
914 cache.cache_texture(texture_path.clone());
915
916 render::render_headless(object_dir, camera_transform, object_rotation, config)
918}
919
920pub fn render_to_files(
937 object_dir: &Path,
938 camera_transform: &Transform,
939 object_rotation: &ObjectRotation,
940 config: &RenderConfig,
941 rgba_path: &Path,
942 depth_path: &Path,
943) -> Result<(), RenderError> {
944 render::render_to_files(
945 object_dir,
946 camera_transform,
947 object_rotation,
948 config,
949 rgba_path,
950 depth_path,
951 )
952}
953
954pub use batch::{
956 BatchRenderConfig, BatchRenderError, BatchRenderOutput, BatchRenderRequest, BatchRenderer,
957 BatchState, RenderStatus,
958};
959
960pub use render::RenderSession;
963
964pub use render::PersistentRenderer;
970
971pub fn create_batch_renderer(config: &BatchRenderConfig) -> Result<BatchRenderer, RenderError> {
989 Ok(BatchRenderer::new(config.clone()))
990}
991
992pub fn queue_render_request(
1017 renderer: &mut BatchRenderer,
1018 request: BatchRenderRequest,
1019) -> Result<(), RenderError> {
1020 renderer
1021 .queue_request(request)
1022 .map_err(|e| RenderError::RenderFailed(e.to_string()))
1023}
1024
1025pub fn render_next_in_batch(
1047 renderer: &mut BatchRenderer,
1048 _timeout_ms: u32,
1049) -> Result<Option<BatchRenderOutput>, RenderError> {
1050 if let Some(request) = renderer.pending_requests.pop_front() {
1051 let output = render_to_buffer(
1052 &request.object_dir,
1053 &request.viewpoint,
1054 &request.object_rotation,
1055 &request.render_config,
1056 )?;
1057 let batch_output = BatchRenderOutput::from_render_output(request, output);
1058 renderer.completed_results.push(batch_output.clone());
1059 renderer.renders_processed += 1;
1060 Ok(Some(batch_output))
1061 } else {
1062 Ok(None)
1063 }
1064}
1065
1066pub fn render_batch(
1085 requests: Vec<BatchRenderRequest>,
1086 config: &BatchRenderConfig,
1087) -> Result<Vec<BatchRenderOutput>, RenderError> {
1088 if requests.is_empty() {
1089 return Ok(Vec::new());
1090 }
1091
1092 if requests.len() > 1 && requests_share_batch_context(&requests) {
1093 let first_request = requests[0].clone();
1094 let viewpoints: Vec<Transform> = requests.iter().map(|request| request.viewpoint).collect();
1095 let outputs = render::render_headless_sequence(
1096 &first_request.object_dir,
1097 &viewpoints,
1098 &first_request.object_rotation,
1099 &first_request.render_config,
1100 )?;
1101
1102 return Ok(requests
1103 .into_iter()
1104 .zip(outputs)
1105 .map(|(request, output)| BatchRenderOutput::from_render_output(request, output))
1106 .collect());
1107 }
1108
1109 let mut renderer = create_batch_renderer(config)?;
1110
1111 for request in requests {
1113 queue_render_request(&mut renderer, request)?;
1114 }
1115
1116 let mut results = Vec::new();
1118 while let Some(output) = render_next_in_batch(&mut renderer, config.frame_timeout_ms)? {
1119 results.push(output);
1120 }
1121
1122 Ok(results)
1123}
1124
1125fn requests_share_batch_context(requests: &[BatchRenderRequest]) -> bool {
1126 let Some(first) = requests.first() else {
1127 return true;
1128 };
1129
1130 requests.iter().all(|request| {
1131 request.object_dir == first.object_dir
1132 && request.object_rotation == first.object_rotation
1133 && request.render_config == first.render_config
1134 })
1135}
1136
1137pub use bevy::prelude::{Quat, Transform, Vec3};
1139
1140#[cfg(test)]
1141mod tests {
1142 use super::*;
1143
1144 #[test]
1145 fn test_object_rotation_identity() {
1146 let rot = ObjectRotation::identity();
1147 assert_eq!(rot.pitch, 0.0);
1148 assert_eq!(rot.yaw, 0.0);
1149 assert_eq!(rot.roll, 0.0);
1150 }
1151
1152 #[test]
1153 fn test_object_rotation_from_array() {
1154 let rot = ObjectRotation::from_array([10.0, 20.0, 30.0]);
1155 assert_eq!(rot.pitch, 10.0);
1156 assert_eq!(rot.yaw, 20.0);
1157 assert_eq!(rot.roll, 30.0);
1158 }
1159
1160 #[test]
1161 fn test_requests_share_batch_context_for_homogeneous_batch() {
1162 let config = RenderConfig::tbp_default();
1163 let request = BatchRenderRequest {
1164 object_dir: "/tmp/ycb/003_cracker_box".into(),
1165 viewpoint: Transform::IDENTITY,
1166 object_rotation: ObjectRotation::identity(),
1167 render_config: config.clone(),
1168 };
1169
1170 assert!(requests_share_batch_context(&[
1171 request.clone(),
1172 BatchRenderRequest {
1173 viewpoint: Transform::from_xyz(1.0, 0.0, 0.0),
1174 ..request
1175 },
1176 ]));
1177 }
1178
1179 #[test]
1180 fn test_requests_share_batch_context_rejects_mixed_objects() {
1181 let config = RenderConfig::tbp_default();
1182 let request = BatchRenderRequest {
1183 object_dir: "/tmp/ycb/003_cracker_box".into(),
1184 viewpoint: Transform::IDENTITY,
1185 object_rotation: ObjectRotation::identity(),
1186 render_config: config.clone(),
1187 };
1188
1189 assert!(!requests_share_batch_context(&[
1190 request.clone(),
1191 BatchRenderRequest {
1192 object_dir: "/tmp/ycb/005_tomato_soup_can".into(),
1193 ..request
1194 },
1195 ]));
1196 }
1197
1198 #[test]
1199 fn test_tbp_benchmark_rotations() {
1200 let rotations = ObjectRotation::tbp_benchmark_rotations();
1201 assert_eq!(rotations.len(), 3);
1202 assert_eq!(rotations[0], ObjectRotation::from_array([0.0, 0.0, 0.0]));
1203 assert_eq!(rotations[1], ObjectRotation::from_array([0.0, 90.0, 0.0]));
1204 assert_eq!(rotations[2], ObjectRotation::from_array([0.0, 180.0, 0.0]));
1205 }
1206
1207 #[test]
1208 fn test_tbp_known_orientations_count() {
1209 let orientations = ObjectRotation::tbp_known_orientations();
1210 assert_eq!(orientations.len(), 14);
1211 }
1212
1213 #[test]
1214 fn test_rotation_to_quat() {
1215 let rot = ObjectRotation::identity();
1216 let quat = rot.to_quat();
1217 assert!((quat.w - 1.0).abs() < 0.001);
1219 assert!(quat.x.abs() < 0.001);
1220 assert!(quat.y.abs() < 0.001);
1221 assert!(quat.z.abs() < 0.001);
1222 }
1223
1224 #[test]
1225 fn test_rotation_90_yaw() {
1226 let rot = ObjectRotation::new(0.0, 90.0, 0.0);
1227 let quat = rot.to_quat();
1228 assert!((quat.w - 0.707).abs() < 0.01);
1230 assert!((quat.y - 0.707).abs() < 0.01);
1231 }
1232
1233 #[test]
1234 fn test_viewpoint_config_default() {
1235 let config = ViewpointConfig::default();
1236 assert_eq!(config.radius, 0.5);
1237 assert_eq!(config.yaw_count, 8);
1238 assert_eq!(config.pitch_angles_deg.len(), 3);
1239 }
1240
1241 #[test]
1242 fn test_viewpoint_count() {
1243 let config = ViewpointConfig::default();
1244 assert_eq!(config.viewpoint_count(), 24); }
1246
1247 #[test]
1248 fn test_generate_viewpoints_count() {
1249 let config = ViewpointConfig::default();
1250 let viewpoints = generate_viewpoints(&config);
1251 assert_eq!(viewpoints.len(), 24);
1252 }
1253
1254 #[test]
1255 fn test_viewpoints_spherical_radius() {
1256 let config = ViewpointConfig::default();
1257 let viewpoints = generate_viewpoints(&config);
1258
1259 for (i, transform) in viewpoints.iter().enumerate() {
1260 let actual_radius = transform.translation.length();
1261 assert!(
1262 (actual_radius - config.radius).abs() < 0.001,
1263 "Viewpoint {} has incorrect radius: {} (expected {})",
1264 i,
1265 actual_radius,
1266 config.radius
1267 );
1268 }
1269 }
1270
1271 #[test]
1272 fn test_viewpoints_looking_at_origin() {
1273 let config = ViewpointConfig::default();
1274 let viewpoints = generate_viewpoints(&config);
1275
1276 for (i, transform) in viewpoints.iter().enumerate() {
1277 let forward = transform.forward();
1278 let to_origin = (Vec3::ZERO - transform.translation).normalize();
1279 let dot = forward.dot(to_origin);
1280 assert!(
1281 dot > 0.99,
1282 "Viewpoint {} not looking at origin, dot product: {}",
1283 i,
1284 dot
1285 );
1286 }
1287 }
1288
1289 #[test]
1290 fn test_sensor_config_default() {
1291 let config = SensorConfig::default();
1292 assert_eq!(config.object_rotations.len(), 1);
1293 assert_eq!(config.total_captures(), 24);
1294 }
1295
1296 #[test]
1297 fn test_sensor_config_tbp_benchmark() {
1298 let config = SensorConfig::tbp_benchmark();
1299 assert_eq!(config.object_rotations.len(), 3);
1300 assert_eq!(config.total_captures(), 72); }
1302
1303 #[test]
1304 fn test_sensor_config_tbp_full() {
1305 let config = SensorConfig::tbp_full_training();
1306 assert_eq!(config.object_rotations.len(), 14);
1307 assert_eq!(config.total_captures(), 336); }
1309
1310 #[test]
1311 fn test_ycb_representative_objects() {
1312 assert_eq!(crate::ycb::REPRESENTATIVE_OBJECTS.len(), 3);
1314 assert!(crate::ycb::REPRESENTATIVE_OBJECTS.contains(&"003_cracker_box"));
1315 }
1316
1317 #[test]
1318 fn test_ycb_tbp_standard_objects() {
1319 assert_eq!(crate::ycb::TBP_STANDARD_OBJECTS.len(), 10);
1320 assert!(crate::ycb::TBP_STANDARD_OBJECTS.contains(&"025_mug"));
1321 }
1322
1323 #[test]
1324 fn test_ycb_tbp_similar_objects() {
1325 assert_eq!(crate::ycb::TBP_SIMILAR_OBJECTS.len(), 10);
1326 assert!(crate::ycb::TBP_SIMILAR_OBJECTS.contains(&"003_cracker_box"));
1327 }
1328
1329 #[test]
1330 fn test_ycb_object_mesh_path() {
1331 let path = crate::ycb::object_mesh_path("/tmp/ycb", "003_cracker_box");
1332 assert_eq!(
1333 path,
1334 std::path::Path::new("/tmp/ycb")
1335 .join("003_cracker_box")
1336 .join("google_16k")
1337 .join("textured.obj")
1338 );
1339 }
1340
1341 #[test]
1342 fn test_ycb_object_texture_path() {
1343 let path = crate::ycb::object_texture_path("/tmp/ycb", "003_cracker_box");
1344 assert_eq!(
1345 path,
1346 std::path::Path::new("/tmp/ycb")
1347 .join("003_cracker_box")
1348 .join("google_16k")
1349 .join("texture_map.png")
1350 );
1351 }
1352
1353 #[test]
1358 fn test_render_config_tbp_default() {
1359 let config = RenderConfig::tbp_default();
1360 assert_eq!(config.width, 64);
1362 assert_eq!(config.height, 64);
1363 assert!(config.zoom > 0.0);
1365 assert!(config.near_plane > 0.0);
1367 assert!(config.far_plane > config.near_plane);
1368 }
1369
1370 #[test]
1371 fn test_render_config_preview() {
1372 let config = RenderConfig::preview();
1373 assert_eq!(config.width, 256);
1374 assert_eq!(config.height, 256);
1375 }
1376
1377 #[test]
1378 fn test_render_config_default_is_tbp() {
1379 let default = RenderConfig::default();
1380 let tbp = RenderConfig::tbp_default();
1381 assert_eq!(default.width, tbp.width);
1382 assert_eq!(default.height, tbp.height);
1383 }
1384
1385 #[test]
1386 fn test_render_config_fov() {
1387 let config = RenderConfig::tbp_default();
1388 let fov = config.fov_radians();
1389 assert!(fov > 0.0);
1392 assert!(fov < PI);
1393
1394 let zoomed = RenderConfig {
1396 zoom: config.zoom * 2.0,
1397 ..config
1398 };
1399 assert!(zoomed.fov_radians() < fov);
1400 }
1401
1402 #[test]
1403 fn test_render_config_intrinsics() {
1404 let config = RenderConfig::tbp_default();
1405 let intrinsics = config.intrinsics();
1406
1407 assert_eq!(intrinsics.image_size, [config.width, config.height]);
1409 assert_eq!(
1410 intrinsics.principal_point,
1411 [config.width as f64 / 2.0, config.height as f64 / 2.0]
1412 );
1413 assert_eq!(intrinsics.focal_length[0], intrinsics.focal_length[1]);
1415 assert!(intrinsics.focal_length[0] > 0.0);
1416 }
1417
1418 #[test]
1419 fn test_render_config_intrinsics_for_size_uses_tbp_zoom_formula() {
1420 let config = RenderConfig {
1421 width: 64,
1422 height: 64,
1423 zoom: 4.0,
1424 ..RenderConfig::tbp_default()
1425 };
1426
1427 let intrinsics = config.intrinsics_for_size(64, 64);
1428
1429 assert!((intrinsics.focal_length[0] - 128.0).abs() < 1e-9);
1432 assert!((intrinsics.focal_length[1] - 128.0).abs() < 1e-9);
1433 assert_ne!(intrinsics.focal_length[0], 64.0 * config.zoom as f64);
1434 assert_eq!(intrinsics.principal_point, [32.0, 32.0]);
1435 assert_eq!(intrinsics.image_size, [64, 64]);
1436 }
1437
1438 #[test]
1439 fn test_render_config_intrinsics_for_size_tracks_actual_readback_size() {
1440 let config = RenderConfig {
1441 width: 64,
1442 height: 64,
1443 zoom: 4.0,
1444 ..RenderConfig::tbp_default()
1445 };
1446
1447 let intrinsics = config.intrinsics_for_size(128, 96);
1448
1449 assert!((intrinsics.focal_length[0] - 256.0).abs() < 1e-9);
1450 assert!((intrinsics.focal_length[1] - 256.0).abs() < 1e-9);
1451 assert_eq!(intrinsics.principal_point, [64.0, 48.0]);
1452 assert_eq!(intrinsics.image_size, [128, 96]);
1453 }
1454
1455 #[test]
1456 fn test_camera_intrinsics_project() {
1457 let intrinsics = CameraIntrinsics {
1458 focal_length: [100.0, 100.0],
1459 principal_point: [32.0, 32.0],
1460 image_size: [64, 64],
1461 };
1462
1463 let center = intrinsics.project(Vec3::new(0.0, 0.0, 1.0));
1465 assert!(center.is_some());
1466 let [x, y] = center.unwrap();
1467 assert!((x - 32.0).abs() < 0.001);
1468 assert!((y - 32.0).abs() < 0.001);
1469
1470 let behind = intrinsics.project(Vec3::new(0.0, 0.0, -1.0));
1472 assert!(behind.is_none());
1473 }
1474
1475 #[test]
1476 fn test_camera_intrinsics_unproject() {
1477 let intrinsics = CameraIntrinsics {
1478 focal_length: [100.0, 100.0],
1479 principal_point: [32.0, 32.0],
1480 image_size: [64, 64],
1481 };
1482
1483 let point = intrinsics.unproject([32.0, 32.0], 1.0);
1485 assert!((point[0]).abs() < 0.001); assert!((point[1]).abs() < 0.001); assert!((point[2] - 1.0).abs() < 0.001); }
1489
1490 #[test]
1491 fn test_render_output_get_rgba() {
1492 let output = RenderOutput {
1493 rgba: vec![
1494 255, 0, 0, 255, 0, 255, 0, 255, 0, 0, 255, 255, 255, 255, 255, 255,
1495 ],
1496 depth: vec![1.0, 2.0, 3.0, 4.0],
1497 width: 2,
1498 height: 2,
1499 intrinsics: RenderConfig::tbp_default().intrinsics(),
1500 camera_transform: Transform::IDENTITY,
1501 object_rotation: ObjectRotation::identity(),
1502 };
1503
1504 assert_eq!(output.get_rgba(0, 0), Some([255, 0, 0, 255]));
1506 assert_eq!(output.get_rgba(1, 0), Some([0, 255, 0, 255]));
1508 assert_eq!(output.get_rgba(0, 1), Some([0, 0, 255, 255]));
1510 assert_eq!(output.get_rgba(1, 1), Some([255, 255, 255, 255]));
1512 assert_eq!(output.get_rgba(2, 0), None);
1514 }
1515
1516 #[test]
1517 fn test_render_output_get_depth() {
1518 let output = RenderOutput {
1519 rgba: vec![0u8; 16],
1520 depth: vec![1.0, 2.0, 3.0, 4.0],
1521 width: 2,
1522 height: 2,
1523 intrinsics: RenderConfig::tbp_default().intrinsics(),
1524 camera_transform: Transform::IDENTITY,
1525 object_rotation: ObjectRotation::identity(),
1526 };
1527
1528 assert_eq!(output.get_depth(0, 0), Some(1.0));
1529 assert_eq!(output.get_depth(1, 0), Some(2.0));
1530 assert_eq!(output.get_depth(0, 1), Some(3.0));
1531 assert_eq!(output.get_depth(1, 1), Some(4.0));
1532 assert_eq!(output.get_depth(2, 0), None);
1533 }
1534
1535 #[test]
1536 fn test_render_output_to_rgb_image() {
1537 let output = RenderOutput {
1538 rgba: vec![
1539 255, 0, 0, 255, 0, 255, 0, 255, 0, 0, 255, 255, 255, 255, 255, 255,
1540 ],
1541 depth: vec![1.0, 2.0, 3.0, 4.0],
1542 width: 2,
1543 height: 2,
1544 intrinsics: RenderConfig::tbp_default().intrinsics(),
1545 camera_transform: Transform::IDENTITY,
1546 object_rotation: ObjectRotation::identity(),
1547 };
1548
1549 let image = output.to_rgb_image();
1550 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]); }
1557
1558 #[test]
1559 fn test_render_output_to_depth_image() {
1560 let output = RenderOutput {
1561 rgba: vec![0u8; 16],
1562 depth: vec![1.0, 2.0, 3.0, 4.0],
1563 width: 2,
1564 height: 2,
1565 intrinsics: RenderConfig::tbp_default().intrinsics(),
1566 camera_transform: Transform::IDENTITY,
1567 object_rotation: ObjectRotation::identity(),
1568 };
1569
1570 let depth_image = output.to_depth_image();
1571 assert_eq!(depth_image.len(), 2);
1572 assert_eq!(depth_image[0], vec![1.0, 2.0]);
1573 assert_eq!(depth_image[1], vec![3.0, 4.0]);
1574 }
1575
1576 #[test]
1577 fn test_render_error_display() {
1578 let err = RenderError::MeshNotFound("/path/to/mesh.obj".to_string());
1579 assert!(err.to_string().contains("Mesh not found"));
1580 assert!(err.to_string().contains("/path/to/mesh.obj"));
1581 }
1582
1583 #[test]
1588 fn test_object_rotation_extreme_angles() {
1589 let rot = ObjectRotation::new(450.0, -720.0, 1080.0);
1591 let quat = rot.to_quat();
1592 assert!((quat.length() - 1.0).abs() < 0.001);
1594 }
1595
1596 #[test]
1597 fn test_object_rotation_to_transform() {
1598 let rot = ObjectRotation::new(45.0, 90.0, 0.0);
1599 let transform = rot.to_transform();
1600 assert_eq!(transform.translation, Vec3::ZERO);
1602 assert!(transform.rotation != Quat::IDENTITY);
1604 }
1605
1606 #[test]
1607 fn test_viewpoint_config_single_viewpoint() {
1608 let config = ViewpointConfig {
1609 radius: 1.0,
1610 yaw_count: 1,
1611 pitch_angles_deg: vec![0.0],
1612 };
1613 assert_eq!(config.viewpoint_count(), 1);
1614 let viewpoints = generate_viewpoints(&config);
1615 assert_eq!(viewpoints.len(), 1);
1616 let pos = viewpoints[0].translation;
1618 assert!((pos.x).abs() < 0.001);
1619 assert!((pos.y).abs() < 0.001);
1620 assert!((pos.z - 1.0).abs() < 0.001);
1621 }
1622
1623 #[test]
1624 fn test_viewpoint_radius_scaling() {
1625 let config1 = ViewpointConfig {
1626 radius: 0.5,
1627 yaw_count: 4,
1628 pitch_angles_deg: vec![0.0],
1629 };
1630 let config2 = ViewpointConfig {
1631 radius: 2.0,
1632 yaw_count: 4,
1633 pitch_angles_deg: vec![0.0],
1634 };
1635
1636 let v1 = generate_viewpoints(&config1);
1637 let v2 = generate_viewpoints(&config2);
1638
1639 for (vp1, vp2) in v1.iter().zip(v2.iter()) {
1641 let ratio = vp2.translation.length() / vp1.translation.length();
1642 assert!((ratio - 4.0).abs() < 0.01); }
1644 }
1645
1646 #[test]
1647 fn test_camera_intrinsics_project_at_z_zero() {
1648 let intrinsics = CameraIntrinsics {
1649 focal_length: [100.0, 100.0],
1650 principal_point: [32.0, 32.0],
1651 image_size: [64, 64],
1652 };
1653
1654 let result = intrinsics.project(Vec3::new(1.0, 1.0, 0.0));
1656 assert!(result.is_none());
1657 }
1658
1659 #[test]
1660 fn test_camera_intrinsics_roundtrip() {
1661 let intrinsics = CameraIntrinsics {
1662 focal_length: [100.0, 100.0],
1663 principal_point: [32.0, 32.0],
1664 image_size: [64, 64],
1665 };
1666
1667 let original = Vec3::new(0.5, -0.3, 2.0);
1669 let projected = intrinsics.project(original).unwrap();
1670
1671 let unprojected = intrinsics.unproject(projected, original.z as f64);
1673
1674 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); }
1679
1680 #[test]
1681 fn test_render_output_empty() {
1682 let output = RenderOutput {
1683 rgba: vec![],
1684 depth: vec![],
1685 width: 0,
1686 height: 0,
1687 intrinsics: RenderConfig::tbp_default().intrinsics(),
1688 camera_transform: Transform::IDENTITY,
1689 object_rotation: ObjectRotation::identity(),
1690 };
1691
1692 assert_eq!(output.get_rgba(0, 0), None);
1694 assert_eq!(output.get_depth(0, 0), None);
1695 assert!(output.to_rgb_image().is_empty());
1696 assert!(output.to_depth_image().is_empty());
1697 }
1698
1699 #[test]
1700 fn test_render_output_1x1() {
1701 let output = RenderOutput {
1702 rgba: vec![128, 64, 32, 255],
1703 depth: vec![0.5],
1704 width: 1,
1705 height: 1,
1706 intrinsics: RenderConfig::tbp_default().intrinsics(),
1707 camera_transform: Transform::IDENTITY,
1708 object_rotation: ObjectRotation::identity(),
1709 };
1710
1711 assert_eq!(output.get_rgba(0, 0), Some([128, 64, 32, 255]));
1712 assert_eq!(output.get_depth(0, 0), Some(0.5));
1713 assert_eq!(output.get_rgb(0, 0), Some([128, 64, 32]));
1714
1715 let rgb_img = output.to_rgb_image();
1716 assert_eq!(rgb_img.len(), 1);
1717 assert_eq!(rgb_img[0].len(), 1);
1718 assert_eq!(rgb_img[0][0], [128, 64, 32]);
1719 }
1720
1721 #[test]
1722 fn test_render_config_high_res() {
1723 let config = RenderConfig::high_res();
1724 assert_eq!(config.width, 512);
1725 assert_eq!(config.height, 512);
1726
1727 let intrinsics = config.intrinsics();
1728 assert_eq!(intrinsics.image_size, [512, 512]);
1729 assert_eq!(intrinsics.principal_point, [256.0, 256.0]);
1730 }
1731
1732 #[test]
1733 fn test_render_config_zoom_affects_fov() {
1734 let base = RenderConfig {
1739 zoom: 2.0,
1740 ..RenderConfig::tbp_default()
1741 };
1742 let doubled = RenderConfig {
1743 zoom: 4.0,
1744 ..RenderConfig::tbp_default()
1745 };
1746
1747 assert!(doubled.fov_radians() < base.fov_radians());
1749
1750 let base_half_tan = (base.fov_radians() / 2.0).tan();
1752 let doubled_half_tan = (doubled.fov_radians() / 2.0).tan();
1753 assert!((base_half_tan / doubled_half_tan - 2.0).abs() < 1e-4);
1754 }
1755
1756 #[test]
1757 fn test_render_config_zoom_affects_intrinsics() {
1758 let a = RenderConfig {
1761 zoom: 2.0,
1762 ..RenderConfig::tbp_default()
1763 };
1764 let b = RenderConfig {
1765 zoom: 4.0,
1766 ..RenderConfig::tbp_default()
1767 };
1768
1769 let fx_a = a.intrinsics().focal_length[0];
1770 let fx_b = b.intrinsics().focal_length[0];
1771
1772 assert!(fx_b > fx_a);
1774
1775 assert!((fx_a / a.zoom as f64 - fx_b / b.zoom as f64).abs() < 1e-9);
1777 }
1778
1779 #[test]
1780 fn test_lighting_config_variants() {
1781 let default = LightingConfig::default();
1782 let bright = LightingConfig::bright();
1783 let soft = LightingConfig::soft();
1784 let unlit = LightingConfig::unlit();
1785
1786 assert!(bright.key_light_intensity > default.key_light_intensity);
1788
1789 assert_eq!(unlit.key_light_intensity, 0.0);
1791 assert_eq!(unlit.fill_light_intensity, 0.0);
1792 assert_eq!(unlit.ambient_brightness, 1.0);
1793
1794 assert!(soft.key_light_intensity < default.key_light_intensity);
1796 }
1797
1798 #[test]
1799 fn test_all_render_error_variants() {
1800 let errors = vec![
1801 RenderError::MeshNotFound("mesh.obj".to_string()),
1802 RenderError::TextureNotFound("texture.png".to_string()),
1803 RenderError::RenderFailed("GPU error".to_string()),
1804 RenderError::InvalidConfig("bad config".to_string()),
1805 ];
1806
1807 for err in errors {
1808 let msg = err.to_string();
1810 assert!(!msg.is_empty());
1811 }
1812 }
1813
1814 #[test]
1815 fn test_tbp_known_orientations_unique() {
1816 let orientations = ObjectRotation::tbp_known_orientations();
1817
1818 let quats: Vec<Quat> = orientations.iter().map(|r| r.to_quat()).collect();
1820
1821 for (i, q1) in quats.iter().enumerate() {
1822 for (j, q2) in quats.iter().enumerate() {
1823 if i != j {
1824 let dot = q1.dot(*q2).abs();
1826 assert!(
1827 dot < 0.999,
1828 "Orientations {} and {} produce same quaternion",
1829 i,
1830 j
1831 );
1832 }
1833 }
1834 }
1835 }
1836}