1use bevy::app::{ScheduleRunnerPlugin, TerminalCtrlCHandlerPlugin};
41use bevy::asset::LoadState;
42use bevy::core_pipeline::prepass::{DepthPrepass, NormalPrepass};
43use bevy::core_pipeline::tonemapping::Tonemapping;
44use bevy::ecs::query::QueryItem;
45use bevy::log::LogPlugin;
46use bevy::prelude::*;
47use bevy::render::camera::{ExtractedCamera, RenderTarget};
48use bevy::render::render_asset::{RenderAssetUsages, RenderAssets};
49use bevy::render::render_graph::{
50 Node, NodeRunError, RenderGraphApp, RenderGraphContext, RenderLabel, ViewNode, ViewNodeRunner,
51};
52use bevy::render::render_resource::{
53 Buffer, BufferDescriptor, BufferUsages, CommandEncoderDescriptor, Extent3d, ImageCopyBuffer,
54 ImageCopyTexture, ImageDataLayout, MapMode, Origin3d, TextureAspect, TextureDimension,
55 TextureFormat, TextureUsages,
56};
57use bevy::render::renderer::RenderQueue;
58use bevy::render::renderer::{RenderContext, RenderDevice};
59use bevy::render::texture::GpuImage;
60use bevy::render::view::screenshot::{Screenshot, ScreenshotCaptured};
61use bevy::render::view::ViewDepthTexture;
62use bevy::render::{Extract, Render, RenderApp, RenderSet};
63use bevy::window::{ExitCondition, WindowPlugin};
64use bevy_obj::ObjPlugin;
65use std::fs::File;
66use std::io::Read as IoRead;
67use std::path::{Path, PathBuf};
68#[cfg(test)]
69use std::sync::atomic::{AtomicUsize, Ordering};
70use std::sync::{Arc, Mutex, OnceLock};
71use std::time::Duration;
72
73use crate::{backend::BackendConfig, ObjectRotation, RenderConfig, RenderError, RenderOutput};
74use ycbust::{GOOGLE_16K_MESH_RELATIVE, GOOGLE_16K_TEXTURE_RELATIVE};
75
76const RENDER_TIMEOUT_SECS: u64 = 180;
82
83const BATCH_WARMUP_FRAMES: u32 = 1;
92
93const PERSISTENT_WARMUP_FRAMES: u32 = 3;
116
117#[inline]
120fn render_trace_enabled() -> bool {
121 std::env::var("BEVY_SENSOR_RENDER_TRACE").is_ok()
122}
123
124#[allow(dead_code)]
128fn display_available() -> bool {
129 std::env::var("DISPLAY").is_ok() || std::env::var("WAYLAND_DISPLAY").is_ok()
130}
131
132#[allow(dead_code)]
134fn is_wsl2() -> bool {
135 if let Ok(version) = std::fs::read_to_string("/proc/version") {
136 return version.to_lowercase().contains("microsoft")
137 || version.to_lowercase().contains("wsl");
138 }
139 false
140}
141
142#[derive(Resource, Default)]
144struct RenderState {
145 frame_count: u32,
146 scene_loaded: bool,
147 texture_loaded: bool,
148 materials_applied: bool,
149 materials_applied_frame: u32,
153 capture_ready: bool,
154 screenshot_requested: bool,
155 captured: bool,
156 exit_requested: bool,
157 #[allow(dead_code)]
158 exit_frame_count: u32,
159 rgba_data: Option<Vec<u8>>,
160 depth_data: Option<Vec<f64>>,
161 image_width: u32,
162 image_height: u32,
163}
164
165#[cfg(test)]
166static HEADLESS_SCENE_SETUP_COUNT: AtomicUsize = AtomicUsize::new(0);
167
168#[cfg(test)]
169fn reset_headless_scene_setup_count() {
170 HEADLESS_SCENE_SETUP_COUNT.store(0, Ordering::SeqCst);
171}
172
173#[cfg(test)]
174fn headless_scene_setup_count() -> usize {
175 HEADLESS_SCENE_SETUP_COUNT.load(Ordering::SeqCst)
176}
177
178#[derive(Resource, Clone)]
180#[allow(clippy::type_complexity)]
181#[allow(dead_code)]
182struct SharedImageBuffer(Arc<Mutex<Option<(Vec<u8>, u32, u32)>>>);
183
184#[derive(Resource, Clone, Default)]
188#[allow(clippy::type_complexity)]
189struct SharedDepthBuffer(Arc<Mutex<Option<(Vec<f64>, u32, u32)>>>);
190
191#[derive(Resource, Default, Clone)]
197struct DepthCaptureRequest {
198 requested: bool,
199 near: f32,
200 far: f32,
201}
202
203struct PendingDepthCapture {
205 buffer: Buffer,
206 width: u32,
207 height: u32,
208 near: f32,
209 far: f32,
210}
211
212#[derive(Resource, Default)]
214struct PendingDepthCaptureQueue(Arc<Mutex<Vec<PendingDepthCapture>>>);
215
216mod depth_helpers {
221 pub const COPY_BYTES_PER_ROW_ALIGNMENT: u32 = 256;
223
224 pub fn align_byte_size(value: u32) -> u32 {
226 let remainder = value % COPY_BYTES_PER_ROW_ALIGNMENT;
227 if remainder == 0 {
228 value
229 } else {
230 value + (COPY_BYTES_PER_ROW_ALIGNMENT - remainder)
231 }
232 }
233
234 #[allow(dead_code)]
236 pub fn get_aligned_size(width: u32, height: u32, pixel_size: u32) -> u32 {
237 height * align_byte_size(width * pixel_size)
238 }
239
240 pub fn reverse_z_to_linear_depth(ndc_depth: f32, near: f32, far: f32) -> f32 {
250 if ndc_depth <= 0.0 {
252 return far; }
254 if ndc_depth >= 1.0 {
255 return near; }
257 far / (1.0 + ndc_depth * (far / near - 1.0))
259 }
260
261 pub fn extract_depth_with_alignment(data: &[u8], width: u32, height: u32) -> Vec<f32> {
263 let pixel_size = 4u32; let aligned_row_bytes = align_byte_size(width * pixel_size) as usize;
265 let actual_row_bytes = (width * pixel_size) as usize;
266
267 let mut depth_values = Vec::with_capacity((width * height) as usize);
268
269 for y in 0..height as usize {
270 let row_start = y * aligned_row_bytes;
271 let row_data = &data[row_start..row_start + actual_row_bytes];
272
273 for x in 0..width as usize {
274 let offset = x * 4;
275 let bytes: [u8; 4] = row_data[offset..offset + 4].try_into().unwrap();
276 let depth_value = f32::from_le_bytes(bytes);
277 depth_values.push(depth_value);
278 }
279 }
280
281 depth_values
282 }
283
284 pub fn convert_depth_to_linear(raw_depth: &[f32], near: f32, far: f32) -> Vec<f64> {
286 raw_depth
287 .iter()
288 .map(|&ndc| reverse_z_to_linear_depth(ndc, near, far) as f64)
289 .collect()
290 }
291
292 #[cfg(test)]
293 mod tests {
294 use super::*;
295
296 #[test]
297 fn test_align_byte_size() {
298 assert_eq!(align_byte_size(256), 256);
299 assert_eq!(align_byte_size(257), 512);
300 assert_eq!(align_byte_size(1), 256);
301 assert_eq!(align_byte_size(512), 512);
302 assert_eq!(align_byte_size(0), 0);
303 }
304
305 #[test]
306 fn test_reverse_z_to_linear_depth() {
307 let near = 0.01;
308 let far = 10.0;
309
310 let linear_near = reverse_z_to_linear_depth(1.0, near, far);
312 assert!((linear_near - near).abs() < 0.001);
313
314 let linear_mid = reverse_z_to_linear_depth(0.5, near, far);
316 assert!(linear_mid > near && linear_mid < far);
318
319 let linear_almost_far = reverse_z_to_linear_depth(0.0001, near, far);
321 assert!(linear_almost_far > 9.0);
323
324 let background = reverse_z_to_linear_depth(0.0, near, far);
326 assert_eq!(background, far);
327 }
328
329 #[test]
330 fn test_extract_depth_with_alignment() {
331 let width = 2u32;
334 let height = 2u32;
335
336 let mut data = vec![0u8; 256 * 2]; data[0..4].copy_from_slice(&0.5f32.to_le_bytes());
341 data[4..8].copy_from_slice(&0.6f32.to_le_bytes());
342 data[256..260].copy_from_slice(&0.7f32.to_le_bytes());
344 data[260..264].copy_from_slice(&0.8f32.to_le_bytes());
345
346 let depth = extract_depth_with_alignment(&data, width, height);
347 assert_eq!(depth.len(), 4);
348 assert!((depth[0] - 0.5).abs() < 0.001);
349 assert!((depth[1] - 0.6).abs() < 0.001);
350 assert!((depth[2] - 0.7).abs() < 0.001);
351 assert!((depth[3] - 0.8).abs() < 0.001);
352 }
353
354 #[test]
355 fn test_reverse_z_depth_at_near_plane() {
356 let near = 0.01;
358 let far = 100.0;
359 let depth = reverse_z_to_linear_depth(1.0, near, far);
360 assert!((depth - near).abs() < 0.0001);
361 }
362
363 #[test]
364 fn test_reverse_z_depth_at_far_plane() {
365 let near = 0.01;
367 let far = 100.0;
368 let depth = reverse_z_to_linear_depth(0.0, near, far);
369 assert!((depth - far).abs() < 0.0001);
370 }
371
372 #[test]
373 fn test_reverse_z_monotonic() {
374 let near = 0.01;
376 let far = 10.0;
377
378 let mut prev_depth = 0.0;
379 for i in (0..=100).rev() {
380 let ndc = i as f32 / 100.0;
381 let depth = reverse_z_to_linear_depth(ndc, near, far);
382 assert!(
383 depth >= prev_depth,
384 "Depth should be monotonic: ndc={}, depth={}, prev={}",
385 ndc,
386 depth,
387 prev_depth
388 );
389 prev_depth = depth;
390 }
391 }
392
393 #[test]
394 fn test_convert_depth_to_linear_batch() {
395 let near = 0.01f32;
396 let far = 10.0f32;
397 let ndc_depths = vec![1.0f32, 0.5, 0.1, 0.0];
398
399 let linear = convert_depth_to_linear(&ndc_depths, near, far);
400
401 assert_eq!(linear.len(), 4);
402 assert!((linear[0] - near as f64).abs() < 0.001);
404 assert!((linear[3] - far as f64).abs() < 0.001);
406 for d in &linear {
408 assert!(*d >= near as f64 && *d <= far as f64);
409 }
410 }
411
412 #[test]
413 fn test_align_byte_size_edge_cases() {
414 assert_eq!(align_byte_size(256), 256);
416 assert_eq!(align_byte_size(512), 512);
417 assert_eq!(align_byte_size(1024), 1024);
418
419 assert_eq!(align_byte_size(255), 256);
421 assert_eq!(align_byte_size(128), 256);
422
423 assert_eq!(align_byte_size(300), 512);
425 }
426
427 #[test]
428 fn test_extract_depth_64x64() {
429 let width = 64u32;
431 let height = 64u32;
432 let bytes_per_pixel = 4u32;
433 let padded_row = align_byte_size(width * bytes_per_pixel);
434
435 let mut data = vec![0u8; (padded_row * height) as usize];
437
438 for y in 0..height {
440 for x in 0..width {
441 let value = (y * width + x) as f32 / (width * height) as f32;
442 let offset = (y * padded_row + x * bytes_per_pixel) as usize;
443 data[offset..offset + 4].copy_from_slice(&value.to_le_bytes());
444 }
445 }
446
447 let depth = extract_depth_with_alignment(&data, width, height);
448 assert_eq!(depth.len(), (width * height) as usize);
449
450 assert!((depth[0] - 0.0).abs() < 0.001);
452 let expected_last = (width * height - 1) as f32 / (width * height) as f32;
453 assert!((depth[(width * height - 1) as usize] - expected_last).abs() < 0.001);
454 }
455 }
456}
457
458#[derive(Debug, Hash, PartialEq, Eq, Clone, bevy::render::render_graph::RenderLabel)]
464struct DepthReadbackLabel;
465
466#[derive(Default)]
469struct DepthReadbackNode;
470
471impl ViewNode for DepthReadbackNode {
472 type ViewQuery = (&'static ViewDepthTexture, &'static ExtractedCamera);
473
474 fn run<'w>(
475 &self,
476 _graph: &mut RenderGraphContext,
477 render_context: &mut RenderContext<'w>,
478 (view_depth_texture, camera): QueryItem<'w, Self::ViewQuery>,
479 world: &'w World,
480 ) -> Result<(), NodeRunError> {
481 let trace = render_trace_enabled();
482 let t0 = trace.then(std::time::Instant::now);
483
484 let Some(request) = world.get_resource::<DepthCaptureRequest>() else {
486 return Ok(());
487 };
488 if !request.requested {
489 return Ok(());
490 }
491
492 let Some(queue) = world.get_resource::<PendingDepthCaptureQueue>() else {
494 return Ok(());
495 };
496
497 let Some(physical_size) = camera.physical_target_size else {
499 return Ok(());
500 };
501 let width = physical_size.x;
502 let height = physical_size.y;
503
504 let render_device = world.resource::<RenderDevice>();
505
506 let bytes_per_pixel = 4u32; let unpadded_bytes_per_row = width * bytes_per_pixel;
509 let padded_bytes_per_row = depth_helpers::align_byte_size(unpadded_bytes_per_row);
510 let buffer_size = (padded_bytes_per_row * height) as u64;
511
512 let staging_buffer = render_device.create_buffer(&BufferDescriptor {
514 label: Some("depth_staging_buffer"),
515 size: buffer_size,
516 usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ,
517 mapped_at_creation: false,
518 });
519
520 let encoder = render_context.command_encoder();
522 encoder.copy_texture_to_buffer(
523 ImageCopyTexture {
524 texture: &view_depth_texture.texture,
525 mip_level: 0,
526 origin: Origin3d::ZERO,
527 aspect: TextureAspect::DepthOnly,
528 },
529 ImageCopyBuffer {
530 buffer: &staging_buffer,
531 layout: ImageDataLayout {
532 offset: 0,
533 bytes_per_row: Some(padded_bytes_per_row),
534 rows_per_image: Some(height),
535 },
536 },
537 Extent3d {
538 width,
539 height,
540 depth_or_array_layers: 1,
541 },
542 );
543
544 if let Ok(mut pending) = queue.0.lock() {
546 pending.push(PendingDepthCapture {
547 buffer: staging_buffer,
548 width,
549 height,
550 near: request.near,
551 far: request.far,
552 });
553 }
554
555 if let Some(t0) = t0 {
556 eprintln!(
557 "[render_trace][node] DepthReadbackNode ms={:.3}",
558 t0.elapsed().as_secs_f64() * 1000.0
559 );
560 }
561
562 Ok(())
563 }
564}
565
566struct DepthReadbackPlugin {
572 shared_depth: SharedDepthBuffer,
573 near: f32,
574 far: f32,
575}
576
577impl Plugin for DepthReadbackPlugin {
578 fn build(&self, app: &mut App) {
579 use bevy::core_pipeline::core_3d::graph::Core3d;
580 use bevy::core_pipeline::core_3d::graph::Node3d;
581
582 app.insert_resource(self.shared_depth.clone());
584 app.insert_resource(DepthCaptureRequest {
585 requested: false,
586 near: self.near,
587 far: self.far,
588 });
589
590 let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
592 eprintln!("Failed to get RenderApp for depth readback");
593 return;
594 };
595
596 render_app.insert_resource(self.shared_depth.clone());
598 render_app.init_resource::<PendingDepthCaptureQueue>();
599
600 render_app.add_systems(ExtractSchedule, extract_depth_request);
602
603 render_app.add_systems(Render, collect_depth_captures.in_set(RenderSet::Cleanup));
605
606 render_app
609 .add_render_graph_node::<ViewNodeRunner<DepthReadbackNode>>(Core3d, DepthReadbackLabel)
610 .add_render_graph_edges(
611 Core3d,
612 (Node3d::EndMainPass, DepthReadbackLabel, Node3d::Tonemapping),
613 );
614 }
615}
616
617fn extract_depth_request(mut commands: Commands, request: Extract<Res<DepthCaptureRequest>>) {
619 commands.insert_resource(DepthCaptureRequest {
620 requested: request.requested,
621 near: request.near,
622 far: request.far,
623 });
624}
625
626fn collect_depth_captures(
628 queue: Res<PendingDepthCaptureQueue>,
629 shared_depth: Res<SharedDepthBuffer>,
630 render_device: Res<RenderDevice>,
631) {
632 let trace = render_trace_enabled();
633 let t_sys = trace.then(std::time::Instant::now);
634
635 let pending_captures = {
637 let Ok(mut pending) = queue.0.lock() else {
638 return;
639 };
640 std::mem::take(&mut *pending)
641 };
642
643 if pending_captures.is_empty() {
644 if let Some(t0) = t_sys {
645 eprintln!(
646 "[render_trace][sys] collect_depth_captures empty ms={:.3}",
647 t0.elapsed().as_secs_f64() * 1000.0
648 );
649 }
650 return;
651 }
652
653 let pending_count = pending_captures.len();
654
655 for pending in pending_captures {
657 let width = pending.width;
658 let height = pending.height;
659 let near = pending.near;
660 let far = pending.far;
661 let buffer = pending.buffer;
662 let shared = shared_depth.0.clone();
663
664 let buffer_slice = buffer.slice(..);
666
667 let (tx, rx) = std::sync::mpsc::channel();
669 buffer_slice.map_async(MapMode::Read, move |result| {
670 let _ = tx.send(result);
671 });
672
673 let t_wait = trace.then(std::time::Instant::now);
674 let mut poll_iters: u32 = 0;
675
676 loop {
678 render_device.poll(bevy::render::render_resource::Maintain::Poll);
679 poll_iters += 1;
680 match rx.try_recv() {
681 Ok(Ok(())) => {
682 let data = buffer_slice.get_mapped_range();
683
684 let ndc_depth =
686 depth_helpers::extract_depth_with_alignment(&data, width, height);
687
688 drop(data);
689 buffer.unmap();
690
691 let linear_depth =
693 depth_helpers::convert_depth_to_linear(&ndc_depth, near, far);
694
695 if let Ok(mut guard) = shared.lock() {
697 *guard = Some((linear_depth, width, height));
698 }
699 break;
700 }
701 Ok(Err(e)) => {
702 eprintln!("Failed to map depth buffer: {:?}", e);
703 break;
704 }
705 Err(std::sync::mpsc::TryRecvError::Empty) => {
706 std::thread::sleep(std::time::Duration::from_millis(1));
708 }
709 Err(std::sync::mpsc::TryRecvError::Disconnected) => {
710 eprintln!("Depth buffer mapping channel disconnected");
711 break;
712 }
713 }
714 }
715
716 if let Some(t_wait) = t_wait {
717 eprintln!(
718 "[render_trace][sys] collect_depth_captures mapping_wait poll_iters={} ms={:.3}",
719 poll_iters,
720 t_wait.elapsed().as_secs_f64() * 1000.0
721 );
722 }
723 }
724
725 if let Some(t0) = t_sys {
726 eprintln!(
727 "[render_trace][sys] collect_depth_captures done pending={} ms={:.3}",
728 pending_count,
729 t0.elapsed().as_secs_f64() * 1000.0
730 );
731 }
732}
733
734#[derive(Debug, Hash, PartialEq, Eq, Clone, RenderLabel)]
740struct ImageCopyLabel;
741
742#[derive(Component, Clone)]
744struct ImageCopier {
745 src_image: Handle<Image>,
747 enabled: bool,
749}
750
751#[derive(Resource, Default)]
753struct ImageCopiers(Vec<ImageCopier>);
754
755struct PendingImageCapture {
757 buffer: Buffer,
758 width: u32,
759 height: u32,
760 padded_bytes_per_row: u32,
761}
762
763#[derive(Resource, Default)]
765struct PendingImageCaptureQueue(Arc<Mutex<Vec<PendingImageCapture>>>);
766
767#[derive(Resource, Clone, Default)]
769#[allow(clippy::type_complexity)]
770struct SharedRgbaBuffer(Arc<Mutex<Option<(Vec<u8>, u32, u32)>>>);
771
772struct ImageCopyDriver;
774
775impl Node for ImageCopyDriver {
776 fn run(
777 &self,
778 _graph: &mut RenderGraphContext,
779 _render_context: &mut RenderContext,
780 world: &World,
781 ) -> Result<(), NodeRunError> {
782 let trace = render_trace_enabled();
783 let t0 = trace.then(std::time::Instant::now);
784
785 let Some(image_copiers) = world.get_resource::<ImageCopiers>() else {
786 return Ok(());
787 };
788
789 let Some(gpu_images) = world.get_resource::<RenderAssets<GpuImage>>() else {
790 return Ok(());
791 };
792
793 let Some(queue) = world.get_resource::<PendingImageCaptureQueue>() else {
794 return Ok(());
795 };
796
797 let render_device = world.resource::<RenderDevice>();
798
799 let Some(render_queue) = world.get_resource::<RenderQueue>() else {
800 return Ok(());
801 };
802
803 for image_copier in image_copiers.0.iter() {
804 if !image_copier.enabled {
805 continue;
806 }
807
808 let Some(gpu_image) = gpu_images.get(&image_copier.src_image) else {
809 continue;
810 };
811
812 let width = gpu_image.size.x;
813 let height = gpu_image.size.y;
814
815 let block_dimensions = gpu_image.texture_format.block_dimensions();
817 let block_size = gpu_image.texture_format.block_copy_size(None).unwrap_or(4); let padded_bytes_per_row = RenderDevice::align_copy_bytes_per_row(
820 (width as usize / block_dimensions.0 as usize) * block_size as usize,
821 );
822
823 let buffer_size = (padded_bytes_per_row * height as usize) as u64;
824
825 let staging_buffer = render_device.create_buffer(&BufferDescriptor {
827 label: Some("image_copy_staging_buffer"),
828 size: buffer_size,
829 usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ,
830 mapped_at_creation: false,
831 });
832
833 let mut encoder =
835 render_device.create_command_encoder(&CommandEncoderDescriptor::default());
836
837 let texture_extent = Extent3d {
838 width,
839 height,
840 depth_or_array_layers: 1,
841 };
842
843 encoder.copy_texture_to_buffer(
845 gpu_image.texture.as_image_copy(),
846 ImageCopyBuffer {
847 buffer: &staging_buffer,
848 layout: ImageDataLayout {
849 offset: 0,
850 bytes_per_row: Some(padded_bytes_per_row as u32),
851 rows_per_image: None,
852 },
853 },
854 texture_extent,
855 );
856
857 render_queue.submit(std::iter::once(encoder.finish()));
859
860 if let Ok(mut pending) = queue.0.lock() {
862 pending.push(PendingImageCapture {
863 buffer: staging_buffer,
864 width,
865 height,
866 padded_bytes_per_row: padded_bytes_per_row as u32,
867 });
868 }
869 }
870
871 if let Some(t0) = t0 {
872 eprintln!(
873 "[render_trace][node] ImageCopyDriver ms={:.3}",
874 t0.elapsed().as_secs_f64() * 1000.0
875 );
876 }
877
878 Ok(())
879 }
880}
881
882fn extract_image_copiers(mut commands: Commands, query: Extract<Query<&ImageCopier>>) {
884 commands.insert_resource(ImageCopiers(query.iter().cloned().collect()));
885}
886
887fn collect_image_captures(
889 queue: Res<PendingImageCaptureQueue>,
890 shared_rgba: Res<SharedRgbaBuffer>,
891 render_device: Res<RenderDevice>,
892) {
893 let trace = render_trace_enabled();
894 let t_sys = trace.then(std::time::Instant::now);
895
896 let pending_captures = {
897 let Ok(mut pending) = queue.0.lock() else {
898 return;
899 };
900 std::mem::take(&mut *pending)
901 };
902
903 if pending_captures.is_empty() {
904 if let Some(t0) = t_sys {
905 eprintln!(
906 "[render_trace][sys] collect_image_captures empty ms={:.3}",
907 t0.elapsed().as_secs_f64() * 1000.0
908 );
909 }
910 return;
911 }
912
913 let pending_count = pending_captures.len();
914
915 for pending in pending_captures {
916 let width = pending.width;
917 let height = pending.height;
918 let padded_bytes_per_row = pending.padded_bytes_per_row;
919 let buffer = pending.buffer;
920 let shared = shared_rgba.0.clone();
921
922 let buffer_slice = buffer.slice(..);
924
925 let (tx, rx) = std::sync::mpsc::channel();
927 buffer_slice.map_async(MapMode::Read, move |result| {
928 let _ = tx.send(result);
929 });
930
931 let start = std::time::Instant::now();
933 let timeout = std::time::Duration::from_secs(10);
934 let mut poll_iters: u32 = 0;
935 loop {
936 render_device.poll(bevy::render::render_resource::Maintain::Poll);
937 poll_iters += 1;
938
939 if start.elapsed() > timeout {
940 eprintln!(
941 "Warning: Buffer mapping timeout after {:?}",
942 start.elapsed()
943 );
944 break;
945 }
946
947 match rx.try_recv() {
948 Ok(Ok(())) => {
949 let data = buffer_slice.get_mapped_range();
950
951 let bytes_per_pixel = 4u32;
953 let actual_row_bytes = (width * bytes_per_pixel) as usize;
954 let padded_row_bytes = padded_bytes_per_row as usize;
955
956 let mut rgba = Vec::with_capacity((width * height * 4) as usize);
957 for y in 0..height as usize {
958 let row_start = y * padded_row_bytes;
959 rgba.extend_from_slice(&data[row_start..row_start + actual_row_bytes]);
960 }
961
962 drop(data);
963 buffer.unmap();
964
965 if let Ok(mut guard) = shared.lock() {
966 *guard = Some((rgba, width, height));
967 }
968 break;
969 }
970 Ok(Err(e)) => {
971 eprintln!("Failed to map image buffer: {:?}", e);
972 break;
973 }
974 Err(std::sync::mpsc::TryRecvError::Empty) => {
975 std::thread::sleep(std::time::Duration::from_millis(1));
977 }
978 Err(std::sync::mpsc::TryRecvError::Disconnected) => {
979 eprintln!("Image buffer mapping channel disconnected");
980 break;
981 }
982 }
983 }
984
985 if trace {
986 eprintln!(
987 "[render_trace][sys] collect_image_captures mapping_wait poll_iters={} ms={:.3}",
988 poll_iters,
989 start.elapsed().as_secs_f64() * 1000.0
990 );
991 }
992 }
993
994 if let Some(t0) = t_sys {
995 eprintln!(
996 "[render_trace][sys] collect_image_captures done pending={} ms={:.3}",
997 pending_count,
998 t0.elapsed().as_secs_f64() * 1000.0
999 );
1000 }
1001}
1002
1003struct ImageCopyPlugin {
1005 shared_rgba: SharedRgbaBuffer,
1006}
1007
1008impl Plugin for ImageCopyPlugin {
1009 fn build(&self, app: &mut App) {
1010 use bevy::render::render_graph::RenderGraph;
1011
1012 app.insert_resource(self.shared_rgba.clone());
1013
1014 let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
1015 return;
1016 };
1017
1018 render_app.insert_resource(self.shared_rgba.clone());
1019 render_app.init_resource::<ImageCopiers>();
1020 render_app.init_resource::<PendingImageCaptureQueue>();
1021
1022 render_app.add_systems(ExtractSchedule, extract_image_copiers);
1023 render_app.add_systems(Render, collect_image_captures.in_set(RenderSet::Cleanup));
1024
1025 let mut graph = render_app.world_mut().resource_mut::<RenderGraph>();
1027 graph.add_node(ImageCopyLabel, ImageCopyDriver);
1028 graph.add_node_edge(bevy::render::graph::CameraDriverLabel, ImageCopyLabel);
1029 }
1030}
1031
1032#[derive(Resource, Clone)]
1038struct RenderRequest {
1039 mesh_path: String,
1040 texture_path: String,
1041 camera_transform: Transform,
1042 object_rotation: ObjectRotation,
1043 config: RenderConfig,
1044}
1045
1046#[derive(Component)]
1048struct RenderedObject;
1049
1050#[derive(Component)]
1052struct RenderCamera;
1053
1054#[derive(Resource)]
1056struct LoadedTexture(Handle<Image>);
1057
1058#[derive(Resource)]
1060struct LoadedScene(Handle<Scene>);
1061
1062#[derive(Resource, Clone)]
1064struct SharedOutput(Arc<Mutex<Option<RenderOutput>>>);
1065
1066#[derive(Resource)]
1068#[allow(dead_code)]
1069struct RenderTargetImage(Handle<Image>);
1070
1071#[derive(Resource)]
1073struct HeadlessBatchSequence {
1074 viewpoints: Vec<Transform>,
1075 current_index: usize,
1076 outputs: Vec<RenderOutput>,
1077 warmup_frames_remaining: u32,
1078 done: bool,
1079}
1080
1081impl HeadlessBatchSequence {
1082 fn new(viewpoints: Vec<Transform>) -> Self {
1083 let capacity = viewpoints.len();
1084 Self {
1085 viewpoints,
1086 current_index: 0,
1087 outputs: Vec::with_capacity(capacity),
1088 warmup_frames_remaining: 0,
1089 done: capacity == 0,
1090 }
1091 }
1092
1093 fn current_viewpoint(&self) -> Option<Transform> {
1094 self.viewpoints.get(self.current_index).cloned()
1095 }
1096}
1097
1098#[allow(dead_code)]
1107pub fn render_headless(
1108 object_dir: &Path,
1109 camera_transform: &Transform,
1110 object_rotation: &ObjectRotation,
1111 config: &RenderConfig,
1112) -> Result<RenderOutput, RenderError> {
1113 let object_dir = std::fs::canonicalize(object_dir).map_err(|e| {
1117 RenderError::RenderFailed(format!(
1118 "Cannot canonicalize object directory {}: {}",
1119 object_dir.display(),
1120 e
1121 ))
1122 })?;
1123 let mesh_path = object_dir.join(GOOGLE_16K_MESH_RELATIVE);
1124 let texture_path = object_dir.join(GOOGLE_16K_TEXTURE_RELATIVE);
1125
1126 if !mesh_path.exists() {
1127 return Err(RenderError::MeshNotFound(mesh_path.display().to_string()));
1128 }
1129 if !texture_path.exists() {
1130 return Err(RenderError::TextureNotFound(
1131 texture_path.display().to_string(),
1132 ));
1133 }
1134
1135 let request = RenderRequest {
1136 mesh_path: mesh_path.display().to_string(),
1137 texture_path: texture_path.display().to_string(),
1138 camera_transform: *camera_transform,
1139 object_rotation: object_rotation.clone(),
1140 config: config.clone(),
1141 };
1142
1143 let shared_output: SharedOutput = SharedOutput(Arc::new(Mutex::new(None)));
1144 let output_clone = shared_output.clone();
1145
1146 let shared_rgba: SharedRgbaBuffer = SharedRgbaBuffer::default();
1148
1149 let shared_depth: SharedDepthBuffer = SharedDepthBuffer::default();
1151
1152 let temp_path =
1154 std::env::temp_dir().join(format!("bevy_sensor_render_{}.bin", std::process::id()));
1155
1156 let output_poll_for_timeout = shared_output.clone();
1158 std::thread::spawn(move || {
1159 let timeout = std::time::Duration::from_secs(RENDER_TIMEOUT_SECS);
1160 let start = std::time::Instant::now();
1161 let poll_interval = std::time::Duration::from_millis(100);
1162
1163 loop {
1164 if let Ok(guard) = output_poll_for_timeout.0.lock() {
1166 if guard.is_some() {
1167 return; }
1170 }
1171
1172 if start.elapsed() > timeout {
1173 eprintln!(
1174 "Error: Render timeout after {} seconds",
1175 RENDER_TIMEOUT_SECS
1176 );
1177 eprintln!("Debug info: This may indicate GPU issues, missing assets, or insufficient system resources.");
1178 std::process::exit(1);
1180 }
1181
1182 std::thread::sleep(poll_interval);
1183 }
1184 });
1185
1186 build_headless_app(request, output_clone, shared_rgba, shared_depth).run();
1189
1190 if let Ok(guard) = shared_output.0.lock() {
1192 if let Some(output) = guard.as_ref() {
1193 return Ok(output.clone());
1194 }
1195 }
1196
1197 if temp_path.exists() {
1199 if let Ok(output) = read_output_from_file(&temp_path) {
1200 let _ = std::fs::remove_file(&temp_path);
1201 return Ok(output);
1202 }
1203 }
1204
1205 Err(RenderError::RenderFailed(
1206 "Render did not complete".to_string(),
1207 ))
1208}
1209
1210pub fn render_headless_sequence(
1215 object_dir: &Path,
1216 viewpoints: &[Transform],
1217 object_rotation: &ObjectRotation,
1218 config: &RenderConfig,
1219) -> Result<Vec<RenderOutput>, RenderError> {
1220 if viewpoints.is_empty() {
1221 return Ok(Vec::new());
1222 }
1223
1224 let object_dir = std::fs::canonicalize(object_dir).map_err(|e| {
1225 RenderError::RenderFailed(format!(
1226 "Cannot canonicalize object directory {}: {}",
1227 object_dir.display(),
1228 e
1229 ))
1230 })?;
1231 let mesh_path = object_dir.join(GOOGLE_16K_MESH_RELATIVE);
1232 let texture_path = object_dir.join(GOOGLE_16K_TEXTURE_RELATIVE);
1233
1234 if !mesh_path.exists() {
1235 return Err(RenderError::MeshNotFound(mesh_path.display().to_string()));
1236 }
1237 if !texture_path.exists() {
1238 return Err(RenderError::TextureNotFound(
1239 texture_path.display().to_string(),
1240 ));
1241 }
1242
1243 let request = RenderRequest {
1244 mesh_path: mesh_path.display().to_string(),
1245 texture_path: texture_path.display().to_string(),
1246 camera_transform: viewpoints[0],
1247 object_rotation: object_rotation.clone(),
1248 config: config.clone(),
1249 };
1250
1251 let shared_rgba: SharedRgbaBuffer = SharedRgbaBuffer::default();
1252 let rgba_clone = shared_rgba.clone();
1253
1254 let shared_depth: SharedDepthBuffer = SharedDepthBuffer::default();
1255 let depth_clone = shared_depth.clone();
1256
1257 let mut app = App::new();
1258 app.add_plugins(
1259 DefaultPlugins
1260 .set(WindowPlugin {
1261 primary_window: None,
1262 exit_condition: ExitCondition::DontExit,
1263 ..default()
1264 })
1265 .disable::<bevy::winit::WinitPlugin>()
1266 .disable::<LogPlugin>()
1267 .disable::<TerminalCtrlCHandlerPlugin>(),
1268 )
1269 .add_plugins(ObjPlugin)
1270 .add_plugins(ImageCopyPlugin {
1271 shared_rgba: rgba_clone,
1272 })
1273 .add_plugins(DepthReadbackPlugin {
1274 shared_depth: depth_clone,
1275 near: config.near_plane,
1276 far: config.far_plane,
1277 })
1278 .insert_resource(request)
1279 .insert_resource(shared_rgba)
1280 .insert_resource(HeadlessBatchSequence::new(viewpoints.to_vec()))
1281 .init_resource::<RenderState>()
1282 .add_systems(Startup, setup_headless_scene)
1283 .add_systems(
1284 Update,
1285 (
1286 check_assets_loaded,
1287 apply_materials,
1288 tick_headless_batch_warmup,
1289 request_headless_capture,
1290 check_headless_capture_ready,
1291 extract_and_continue_headless_batch,
1292 )
1293 .chain(),
1294 );
1295
1296 let trace_outer = render_trace_enabled();
1300 let t_finish = std::time::Instant::now();
1301 app.finish();
1302 let finish_ms = t_finish.elapsed().as_secs_f64() * 1000.0;
1303 let t_cleanup = std::time::Instant::now();
1304 app.cleanup();
1305 let cleanup_ms = t_cleanup.elapsed().as_secs_f64() * 1000.0;
1306 if trace_outer {
1307 eprintln!(
1308 "[render_trace][coldinit] app.finish ms={:.3} app.cleanup ms={:.3}",
1309 finish_ms, cleanup_ms
1310 );
1311 }
1312
1313 let timeout = std::time::Duration::from_secs(RENDER_TIMEOUT_SECS);
1314 let start = std::time::Instant::now();
1315
1316 let trace = std::env::var("BEVY_SENSOR_RENDER_TRACE").is_ok();
1317 let mut update_idx: u32 = 0;
1318 let mut last_completed_outputs: usize = 0;
1319 let mut viewpoint_start = std::time::Instant::now();
1320
1321 loop {
1322 if start.elapsed() > timeout {
1323 return Err(RenderError::RenderTimeout {
1324 duration_secs: RENDER_TIMEOUT_SECS,
1325 });
1326 }
1327
1328 let update_start = std::time::Instant::now();
1329 app.update();
1330 let update_elapsed_ms = update_start.elapsed().as_secs_f64() * 1000.0;
1331
1332 if trace {
1333 let batch = app.world().resource::<HeadlessBatchSequence>();
1334 let warmup = batch.warmup_frames_remaining;
1335 let current = batch.current_index;
1336 let completed = batch.outputs.len();
1337 let vp_ms = viewpoint_start.elapsed().as_secs_f64() * 1000.0;
1338 eprintln!(
1339 "[render_trace] update={update_idx} vp={current} warmup={warmup} \
1340 completed={completed} update_ms={update_elapsed_ms:.2} vp_ms={vp_ms:.2}"
1341 );
1342 if completed > last_completed_outputs {
1343 eprintln!(
1344 "[render_trace] viewpoint {} finished in {:.2} ms",
1345 completed - 1,
1346 vp_ms
1347 );
1348 last_completed_outputs = completed;
1349 viewpoint_start = std::time::Instant::now();
1350 }
1351 }
1352
1353 update_idx += 1;
1354
1355 if app.world().resource::<HeadlessBatchSequence>().done {
1356 break;
1357 }
1358 }
1359
1360 if trace {
1361 eprintln!(
1362 "[render_trace] total_wall_ms={:.2} updates={update_idx} viewpoints={}",
1363 start.elapsed().as_secs_f64() * 1000.0,
1364 viewpoints.len()
1365 );
1366 }
1367
1368 let mut batch = app.world_mut().resource_mut::<HeadlessBatchSequence>();
1369 if batch.outputs.len() != viewpoints.len() {
1370 return Err(RenderError::RenderFailed(format!(
1371 "Batch render produced {} outputs for {} viewpoints",
1372 batch.outputs.len(),
1373 viewpoints.len()
1374 )));
1375 }
1376
1377 Ok(std::mem::take(&mut batch.outputs))
1378}
1379
1380fn build_headless_app(
1382 request: RenderRequest,
1383 shared_output: SharedOutput,
1384 shared_rgba: SharedRgbaBuffer,
1385 shared_depth: SharedDepthBuffer,
1386) -> App {
1387 let near = request.config.near_plane;
1388 let far = request.config.far_plane;
1389
1390 let mut app = App::new();
1391 app.add_plugins(
1392 DefaultPlugins
1393 .set(WindowPlugin {
1394 primary_window: None,
1395 exit_condition: ExitCondition::DontExit,
1396 ..default()
1397 })
1398 .disable::<bevy::winit::WinitPlugin>()
1399 .disable::<LogPlugin>()
1400 .disable::<TerminalCtrlCHandlerPlugin>(),
1401 )
1402 .add_plugins(ScheduleRunnerPlugin::run_loop(Duration::from_secs_f64(
1403 1.0 / 60.0,
1404 )))
1405 .add_plugins(ObjPlugin)
1406 .add_plugins(ImageCopyPlugin {
1407 shared_rgba: shared_rgba.clone(),
1408 })
1409 .add_plugins(DepthReadbackPlugin {
1410 shared_depth,
1411 near,
1412 far,
1413 })
1414 .insert_resource(request)
1415 .insert_resource(shared_output)
1416 .insert_resource(shared_rgba)
1417 .init_resource::<RenderState>()
1418 .add_systems(Startup, setup_headless_scene)
1419 .add_systems(
1420 Update,
1421 (
1422 check_assets_loaded,
1423 apply_materials,
1424 request_headless_capture,
1425 check_headless_capture_ready,
1426 extract_and_exit_headless,
1427 )
1428 .chain(),
1429 );
1430 app
1431}
1432
1433#[allow(dead_code)]
1435fn serialize_output(output: &RenderOutput) -> Vec<u8> {
1436 let mut data = Vec::new();
1437
1438 data.extend_from_slice(&output.width.to_le_bytes());
1440 data.extend_from_slice(&output.height.to_le_bytes());
1441 data.extend_from_slice(&(output.rgba.len() as u32).to_le_bytes());
1442 data.extend_from_slice(&(output.depth.len() as u32).to_le_bytes());
1443
1444 data.extend_from_slice(&output.rgba);
1446
1447 for d in &output.depth {
1449 data.extend_from_slice(&d.to_le_bytes());
1450 }
1451
1452 data.extend_from_slice(&output.intrinsics.focal_length[0].to_le_bytes());
1454 data.extend_from_slice(&output.intrinsics.focal_length[1].to_le_bytes());
1455 data.extend_from_slice(&output.intrinsics.principal_point[0].to_le_bytes());
1456 data.extend_from_slice(&output.intrinsics.principal_point[1].to_le_bytes());
1457 data.extend_from_slice(&output.intrinsics.image_size[0].to_le_bytes());
1458 data.extend_from_slice(&output.intrinsics.image_size[1].to_le_bytes());
1459
1460 let t = output.camera_transform.translation;
1462 let r = output.camera_transform.rotation;
1463 data.extend_from_slice(&t.x.to_le_bytes());
1464 data.extend_from_slice(&t.y.to_le_bytes());
1465 data.extend_from_slice(&t.z.to_le_bytes());
1466 data.extend_from_slice(&r.x.to_le_bytes());
1467 data.extend_from_slice(&r.y.to_le_bytes());
1468 data.extend_from_slice(&r.z.to_le_bytes());
1469 data.extend_from_slice(&r.w.to_le_bytes());
1470
1471 let or = &output.object_rotation;
1473 data.extend_from_slice(&or.pitch.to_le_bytes());
1474 data.extend_from_slice(&or.yaw.to_le_bytes());
1475 data.extend_from_slice(&or.roll.to_le_bytes());
1476
1477 data
1478}
1479
1480fn read_output_from_file(path: &std::path::Path) -> Result<RenderOutput, RenderError> {
1482 let mut file = File::open(path).map_err(|e| RenderError::RenderFailed(e.to_string()))?;
1483 let mut data = Vec::new();
1484 file.read_to_end(&mut data)
1485 .map_err(|e| RenderError::RenderFailed(e.to_string()))?;
1486
1487 let mut cursor = 0;
1488
1489 let read_u32 = |data: &[u8], cursor: &mut usize| -> u32 {
1490 let val = u32::from_le_bytes(data[*cursor..*cursor + 4].try_into().unwrap());
1491 *cursor += 4;
1492 val
1493 };
1494
1495 let read_f32 = |data: &[u8], cursor: &mut usize| -> f32 {
1496 let val = f32::from_le_bytes(data[*cursor..*cursor + 4].try_into().unwrap());
1497 *cursor += 4;
1498 val
1499 };
1500
1501 let read_f64 = |data: &[u8], cursor: &mut usize| -> f64 {
1502 let val = f64::from_le_bytes(data[*cursor..*cursor + 8].try_into().unwrap());
1503 *cursor += 8;
1504 val
1505 };
1506
1507 let width = read_u32(&data, &mut cursor);
1508 let height = read_u32(&data, &mut cursor);
1509 let rgba_len = read_u32(&data, &mut cursor) as usize;
1510 let depth_len = read_u32(&data, &mut cursor) as usize;
1511
1512 let rgba = data[cursor..cursor + rgba_len].to_vec();
1513 cursor += rgba_len;
1514
1515 let mut depth = Vec::with_capacity(depth_len);
1517 for _ in 0..depth_len {
1518 depth.push(read_f64(&data, &mut cursor));
1519 }
1520
1521 let focal_length = [read_f64(&data, &mut cursor), read_f64(&data, &mut cursor)];
1523 let principal_point = [read_f64(&data, &mut cursor), read_f64(&data, &mut cursor)];
1524 let image_size = [read_u32(&data, &mut cursor), read_u32(&data, &mut cursor)];
1525
1526 let tx = read_f32(&data, &mut cursor);
1528 let ty = read_f32(&data, &mut cursor);
1529 let tz = read_f32(&data, &mut cursor);
1530 let rx = read_f32(&data, &mut cursor);
1531 let ry = read_f32(&data, &mut cursor);
1532 let rz = read_f32(&data, &mut cursor);
1533 let rw = read_f32(&data, &mut cursor);
1534
1535 let pitch = read_f64(&data, &mut cursor);
1537 let yaw = read_f64(&data, &mut cursor);
1538 let roll = read_f64(&data, &mut cursor);
1539
1540 Ok(RenderOutput {
1541 rgba,
1542 depth,
1543 width,
1544 height,
1545 intrinsics: crate::CameraIntrinsics {
1546 focal_length,
1547 principal_point,
1548 image_size,
1549 },
1550 camera_transform: Transform {
1551 translation: Vec3::new(tx, ty, tz),
1552 rotation: Quat::from_xyzw(rx, ry, rz, rw),
1553 scale: Vec3::ONE,
1554 },
1555 object_rotation: ObjectRotation { pitch, yaw, roll },
1556 })
1557}
1558
1559#[allow(dead_code)]
1561fn setup_scene(
1562 mut commands: Commands,
1563 asset_server: Res<AssetServer>,
1564 request: Res<RenderRequest>,
1565 mut _materials: ResMut<Assets<StandardMaterial>>,
1566) {
1567 let fov = request.config.fov_radians();
1571 commands.spawn((
1572 Camera3d::default(),
1573 Camera {
1574 hdr: true,
1575 ..default()
1576 },
1577 Projection::Perspective(PerspectiveProjection {
1578 fov,
1579 near: request.config.near_plane,
1580 far: request.config.far_plane,
1581 ..default()
1582 }),
1583 Msaa::Off,
1584 request.camera_transform,
1585 Tonemapping::None, DepthPrepass,
1587 NormalPrepass,
1588 RenderCamera,
1589 ));
1590
1591 let lighting = &request.config.lighting;
1593 commands.insert_resource(AmbientLight {
1594 color: Color::WHITE,
1595 brightness: lighting.ambient_brightness,
1596 });
1597
1598 if lighting.key_light_intensity > 0.0 {
1600 commands.spawn((
1601 PointLight {
1602 intensity: lighting.key_light_intensity,
1603 shadows_enabled: lighting.shadows_enabled,
1604 ..default()
1605 },
1606 Transform::from_xyz(
1607 lighting.key_light_position[0],
1608 lighting.key_light_position[1],
1609 lighting.key_light_position[2],
1610 ),
1611 ));
1612 }
1613
1614 if lighting.fill_light_intensity > 0.0 {
1616 commands.spawn((
1617 PointLight {
1618 intensity: lighting.fill_light_intensity,
1619 shadows_enabled: lighting.shadows_enabled,
1620 ..default()
1621 },
1622 Transform::from_xyz(
1623 lighting.fill_light_position[0],
1624 lighting.fill_light_position[1],
1625 lighting.fill_light_position[2],
1626 ),
1627 ));
1628 }
1629
1630 let scene_handle: Handle<Scene> = asset_server.load(&request.mesh_path);
1632 commands.insert_resource(LoadedScene(scene_handle.clone()));
1633
1634 let texture_handle: Handle<Image> = asset_server.load(&request.texture_path);
1636 commands.insert_resource(LoadedTexture(texture_handle.clone()));
1637
1638 let _material = _materials.add(StandardMaterial {
1640 base_color_texture: Some(texture_handle),
1641 unlit: true,
1642 ..default()
1643 });
1644
1645 commands.spawn((
1647 SceneRoot(scene_handle),
1648 Transform::from_rotation(request.object_rotation.to_quat()),
1649 RenderedObject,
1650 ));
1651
1652 println!("Scene setup complete");
1653}
1654
1655fn check_assets_loaded(
1657 mut state: ResMut<RenderState>,
1658 asset_server: Res<AssetServer>,
1659 scene: Option<Res<LoadedScene>>,
1660 texture: Option<Res<LoadedTexture>>,
1661) {
1662 let trace = render_trace_enabled();
1663 let was_scene_loaded = state.scene_loaded;
1664 let was_texture_loaded = state.texture_loaded;
1665
1666 state.frame_count += 1;
1667
1668 if state.scene_loaded && state.texture_loaded {
1669 return;
1670 }
1671
1672 if let Some(scene) = scene {
1673 match asset_server.get_load_state(&scene.0) {
1674 Some(LoadState::Loaded) => {
1675 state.scene_loaded = true;
1676 }
1677 Some(LoadState::Failed(_)) => {}
1678 _ => {}
1679 }
1680 }
1681
1682 if let Some(texture) = texture {
1683 match asset_server.get_load_state(&texture.0) {
1684 Some(LoadState::Loaded) => {
1685 state.texture_loaded = true;
1686 }
1687 Some(LoadState::Failed(_)) => {}
1688 _ => {}
1689 }
1690 }
1691
1692 if trace {
1693 if !was_scene_loaded && state.scene_loaded {
1694 eprintln!(
1695 "[render_trace][coldinit] scene_loaded frame_count={}",
1696 state.frame_count
1697 );
1698 }
1699 if !was_texture_loaded && state.texture_loaded {
1700 eprintln!(
1701 "[render_trace][coldinit] texture_loaded frame_count={}",
1702 state.frame_count
1703 );
1704 }
1705 }
1706}
1707
1708fn apply_materials(
1710 mut state: ResMut<RenderState>,
1711 texture: Option<Res<LoadedTexture>>,
1712 mut materials: ResMut<Assets<StandardMaterial>>,
1713 mut mesh_query: Query<&mut MeshMaterial3d<StandardMaterial>, With<Mesh3d>>,
1715) {
1716 if !state.scene_loaded || !state.texture_loaded || state.capture_ready {
1717 return;
1718 }
1719
1720 state.frame_count += 1;
1721
1722 let Some(tex) = texture else { return };
1723
1724 if !state.materials_applied {
1725 if mesh_query.is_empty() {
1728 return;
1729 }
1730
1731 let textured_material = materials.add(StandardMaterial {
1732 base_color_texture: Some(tex.0.clone()),
1733 unlit: true,
1734 ..default()
1735 });
1736
1737 for mut mat in mesh_query.iter_mut() {
1738 mat.0 = textured_material.clone();
1739 }
1740
1741 state.materials_applied = true;
1742 state.materials_applied_frame = state.frame_count;
1743 }
1744
1745 if state.frame_count >= state.materials_applied_frame + 2 {
1749 let was_ready = state.capture_ready;
1750 state.capture_ready = true;
1751 if render_trace_enabled() && !was_ready {
1752 eprintln!(
1753 "[render_trace][coldinit] capture_ready frame_count={}",
1754 state.frame_count
1755 );
1756 }
1757 }
1758}
1759
1760#[allow(dead_code)]
1762fn request_screenshot(
1763 mut commands: Commands,
1764 mut state: ResMut<RenderState>,
1765 shared_image: Res<SharedImageBuffer>,
1766 mut depth_request: ResMut<DepthCaptureRequest>,
1767) {
1768 if !state.capture_ready || state.screenshot_requested {
1769 return;
1770 }
1771
1772 let image_buffer = shared_image.0.clone();
1774
1775 depth_request.requested = true;
1777 println!("Depth capture requested");
1778
1779 println!("Requesting screenshot via Screenshot entity");
1781 commands.spawn(Screenshot::primary_window()).observe(
1782 move |trigger: Trigger<ScreenshotCaptured>| {
1783 let image: &Image = trigger.event();
1785
1786 let width = image.texture_descriptor.size.width;
1788 let height = image.texture_descriptor.size.height;
1789
1790 let rgba_data = image.data.clone();
1792
1793 if let Ok(mut guard) = image_buffer.lock() {
1795 *guard = Some((rgba_data, width, height));
1796 }
1797 },
1798 );
1799
1800 state.screenshot_requested = true;
1801 println!("Screenshot requested");
1802}
1803
1804#[allow(dead_code)]
1806fn check_screenshot_ready(
1807 mut state: ResMut<RenderState>,
1808 shared_image: Res<SharedImageBuffer>,
1809 shared_depth: Res<SharedDepthBuffer>,
1810 request: Res<RenderRequest>,
1811) {
1812 if !state.screenshot_requested || state.captured {
1813 return;
1814 }
1815
1816 state.frame_count += 1;
1818
1819 let rgba_ready = if let Ok(guard) = shared_image.0.lock() {
1821 if let Some((rgba_data, width, height)) = guard.as_ref() {
1822 if state.rgba_data.is_none() {
1823 state.rgba_data = Some(rgba_data.clone());
1824 state.image_width = *width;
1825 state.image_height = *height;
1826 }
1827 true
1828 } else {
1829 false
1830 }
1831 } else {
1832 false
1833 };
1834
1835 let depth_ready = if let Ok(guard) = shared_depth.0.lock() {
1837 if let Some((depth_data, _width, _height)) = guard.as_ref() {
1838 if state.depth_data.is_none() {
1839 state.depth_data = Some(depth_data.clone());
1840 }
1841 true
1842 } else {
1843 false
1844 }
1845 } else {
1846 false
1847 };
1848
1849 if rgba_ready && !depth_ready && state.frame_count > 60 {
1852 let camera_dist = request.camera_transform.translation.length() as f64;
1853 let pixel_count = (state.image_width * state.image_height) as usize;
1854 state.depth_data = Some(vec![camera_dist; pixel_count]);
1855 }
1856
1857 if state.rgba_data.is_some() && state.depth_data.is_some() {
1859 state.captured = true;
1860 }
1861}
1862
1863#[allow(dead_code)]
1865fn extract_and_exit(
1866 mut state: ResMut<RenderState>,
1867 request: Res<RenderRequest>,
1868 shared_output: Res<SharedOutput>,
1869 mut commands: Commands,
1870 windows: Query<Entity, With<bevy::window::Window>>,
1871) {
1872 if state.exit_requested {
1874 state.exit_frame_count += 1;
1875 return;
1877 }
1878
1879 if !state.captured {
1880 return;
1881 }
1882
1883 if let (Some(rgba), Some(depth)) = (&state.rgba_data, &state.depth_data) {
1884 let width = state.image_width;
1886 let height = state.image_height;
1887
1888 let intrinsics = request.config.intrinsics_for_size(width, height);
1890
1891 let output = RenderOutput {
1892 rgba: rgba.clone(),
1893 depth: depth.clone(),
1894 width,
1895 height,
1896 intrinsics,
1897 camera_transform: request.camera_transform,
1898 object_rotation: request.object_rotation.clone(),
1899 };
1900
1901 if let Ok(mut guard) = shared_output.0.lock() {
1902 *guard = Some(output);
1903 drop(guard); std::thread::sleep(std::time::Duration::from_millis(200));
1907 }
1908
1909 for window_entity in windows.iter() {
1912 commands.entity(window_entity).despawn();
1913 }
1914 state.exit_requested = true;
1915 }
1916}
1917
1918fn setup_headless_scene(
1924 mut commands: Commands,
1925 mut images: ResMut<Assets<Image>>,
1926 asset_server: Res<AssetServer>,
1927 request: Res<RenderRequest>,
1928 mut _materials: ResMut<Assets<StandardMaterial>>,
1929) {
1930 let trace = render_trace_enabled();
1931 let t0 = trace.then(std::time::Instant::now);
1932
1933 #[cfg(test)]
1934 HEADLESS_SCENE_SETUP_COUNT.fetch_add(1, Ordering::SeqCst);
1935
1936 let width = request.config.width;
1937 let height = request.config.height;
1938
1939 let size = Extent3d {
1941 width,
1942 height,
1943 depth_or_array_layers: 1,
1944 };
1945
1946 let mut render_target_image = Image::new_fill(
1947 size,
1948 TextureDimension::D2,
1949 &[0, 0, 0, 255], TextureFormat::Rgba8UnormSrgb,
1951 RenderAssetUsages::default(),
1952 );
1953
1954 render_target_image.texture_descriptor.usage =
1956 TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_SRC | TextureUsages::RENDER_ATTACHMENT;
1957
1958 let render_target_handle = images.add(render_target_image);
1959
1960 commands.insert_resource(RenderTargetImage(render_target_handle.clone()));
1962
1963 let fov = request.config.fov_radians();
1965 commands.spawn((
1966 Camera3d::default(),
1967 Camera {
1968 hdr: true,
1969 target: RenderTarget::Image(render_target_handle.clone()),
1970 ..default()
1971 },
1972 Projection::Perspective(PerspectiveProjection {
1973 fov,
1974 near: request.config.near_plane,
1975 far: request.config.far_plane,
1976 ..default()
1977 }),
1978 Msaa::Off,
1979 request.camera_transform,
1980 Tonemapping::None,
1981 DepthPrepass,
1982 NormalPrepass,
1983 RenderCamera,
1984 ImageCopier {
1986 src_image: render_target_handle,
1987 enabled: false, },
1989 ));
1990
1991 let lighting = &request.config.lighting;
1993 commands.insert_resource(AmbientLight {
1994 color: Color::WHITE,
1995 brightness: lighting.ambient_brightness,
1996 });
1997
1998 if lighting.key_light_intensity > 0.0 {
2000 commands.spawn((
2001 PointLight {
2002 intensity: lighting.key_light_intensity,
2003 shadows_enabled: lighting.shadows_enabled,
2004 ..default()
2005 },
2006 Transform::from_xyz(
2007 lighting.key_light_position[0],
2008 lighting.key_light_position[1],
2009 lighting.key_light_position[2],
2010 ),
2011 ));
2012 }
2013
2014 if lighting.fill_light_intensity > 0.0 {
2016 commands.spawn((
2017 PointLight {
2018 intensity: lighting.fill_light_intensity,
2019 shadows_enabled: lighting.shadows_enabled,
2020 ..default()
2021 },
2022 Transform::from_xyz(
2023 lighting.fill_light_position[0],
2024 lighting.fill_light_position[1],
2025 lighting.fill_light_position[2],
2026 ),
2027 ));
2028 }
2029
2030 let scene_handle: Handle<Scene> = asset_server.load(&request.mesh_path);
2032 commands.insert_resource(LoadedScene(scene_handle.clone()));
2033
2034 let texture_handle: Handle<Image> = asset_server.load(&request.texture_path);
2036 commands.insert_resource(LoadedTexture(texture_handle.clone()));
2037
2038 let _material = _materials.add(StandardMaterial {
2040 base_color_texture: Some(texture_handle),
2041 unlit: true,
2042 ..default()
2043 });
2044
2045 commands.spawn((
2047 SceneRoot(scene_handle),
2048 Transform::from_rotation(request.object_rotation.to_quat()),
2049 RenderedObject,
2050 ));
2051
2052 if let Some(t0) = t0 {
2053 eprintln!(
2054 "[render_trace][startup] setup_headless_scene ms={:.3}",
2055 t0.elapsed().as_secs_f64() * 1000.0
2056 );
2057 }
2058}
2059
2060fn request_headless_capture(
2062 mut state: ResMut<RenderState>,
2063 mut depth_request: ResMut<DepthCaptureRequest>,
2064 mut query: Query<&mut ImageCopier>,
2065 batch: Option<Res<HeadlessBatchSequence>>,
2066) {
2067 let trace = render_trace_enabled();
2068 let t0 = trace.then(std::time::Instant::now);
2069
2070 if !state.capture_ready || state.screenshot_requested {
2071 if let Some(t0) = t0 {
2072 eprintln!(
2073 "[render_trace][sys] request_headless_capture skipped(gate) ms={:.3}",
2074 t0.elapsed().as_secs_f64() * 1000.0
2075 );
2076 }
2077 return;
2078 }
2079
2080 if batch
2081 .as_ref()
2082 .is_some_and(|batch| batch.warmup_frames_remaining > 0)
2083 {
2084 if let Some(t0) = t0 {
2085 eprintln!(
2086 "[render_trace][sys] request_headless_capture skipped(warmup) ms={:.3}",
2087 t0.elapsed().as_secs_f64() * 1000.0
2088 );
2089 }
2090 return;
2091 }
2092
2093 for mut copier in query.iter_mut() {
2095 copier.enabled = true;
2096 }
2097
2098 depth_request.requested = true;
2100
2101 state.screenshot_requested = true;
2102
2103 if let Some(t0) = t0 {
2104 eprintln!(
2105 "[render_trace][sys] request_headless_capture requested ms={:.3}",
2106 t0.elapsed().as_secs_f64() * 1000.0
2107 );
2108 }
2109}
2110
2111fn check_headless_capture_ready(
2113 mut state: ResMut<RenderState>,
2114 shared_rgba: Res<SharedRgbaBuffer>,
2115 shared_depth: Res<SharedDepthBuffer>,
2116 request: Res<RenderRequest>,
2117 mut query: Query<&mut ImageCopier>,
2118) {
2119 let trace = render_trace_enabled();
2120 let t0 = trace.then(std::time::Instant::now);
2121
2122 if !state.screenshot_requested || state.captured {
2123 if let Some(t0) = t0 {
2124 eprintln!(
2125 "[render_trace][sys] check_headless_capture_ready skipped(gate) ms={:.3}",
2126 t0.elapsed().as_secs_f64() * 1000.0
2127 );
2128 }
2129 return;
2130 }
2131
2132 state.frame_count += 1;
2133
2134 let rgba_ready = if let Ok(guard) = shared_rgba.0.lock() {
2136 if let Some((rgba_data, width, height)) = guard.as_ref() {
2137 if state.rgba_data.is_none() {
2138 state.rgba_data = Some(rgba_data.clone());
2139 state.image_width = *width;
2140 state.image_height = *height;
2141 for mut copier in query.iter_mut() {
2143 copier.enabled = false;
2144 }
2145 }
2146 true
2147 } else {
2148 false
2149 }
2150 } else {
2151 false
2152 };
2153
2154 let depth_ready = if let Ok(guard) = shared_depth.0.lock() {
2156 if let Some((depth_data, _width, _height)) = guard.as_ref() {
2157 if state.depth_data.is_none() {
2158 state.depth_data = Some(depth_data.clone());
2159 }
2160 true
2161 } else {
2162 false
2163 }
2164 } else {
2165 false
2166 };
2167
2168 if rgba_ready && !depth_ready && state.frame_count > 70 {
2170 let camera_dist = request.camera_transform.translation.length() as f64;
2171 let pixel_count = (state.image_width * state.image_height) as usize;
2172 state.depth_data = Some(vec![camera_dist; pixel_count]);
2173 }
2174
2175 if state.rgba_data.is_some() && state.depth_data.is_some() {
2176 state.captured = true;
2177 }
2178
2179 if let Some(t0) = t0 {
2180 eprintln!(
2181 "[render_trace][sys] check_headless_capture_ready rgba_ready={} depth_ready={} captured={} frame_count={} ms={:.3}",
2182 rgba_ready,
2183 depth_ready,
2184 state.captured,
2185 state.frame_count,
2186 t0.elapsed().as_secs_f64() * 1000.0
2187 );
2188 }
2189}
2190
2191fn extract_and_exit_headless(
2193 mut state: ResMut<RenderState>,
2194 request: Res<RenderRequest>,
2195 shared_output: Res<SharedOutput>,
2196 mut app_exit: EventWriter<bevy::app::AppExit>,
2197 batch: Option<Res<HeadlessBatchSequence>>,
2198) {
2199 if batch.is_some() {
2200 return;
2201 }
2202
2203 if state.exit_requested {
2204 return;
2205 }
2206
2207 if !state.captured {
2208 return;
2209 }
2210
2211 if let (Some(rgba), Some(depth)) = (&state.rgba_data, &state.depth_data) {
2212 let width = state.image_width;
2213 let height = state.image_height;
2214
2215 let intrinsics = request.config.intrinsics_for_size(width, height);
2217
2218 let output = RenderOutput {
2219 rgba: rgba.clone(),
2220 depth: depth.clone(),
2221 width,
2222 height,
2223 intrinsics,
2224 camera_transform: request.camera_transform,
2225 object_rotation: request.object_rotation.clone(),
2226 };
2227
2228 if let Ok(mut guard) = shared_output.0.lock() {
2229 *guard = Some(output);
2230 drop(guard);
2231 std::thread::sleep(std::time::Duration::from_millis(200));
2232 }
2233
2234 app_exit.send(bevy::app::AppExit::Success);
2236 state.exit_requested = true;
2237 }
2238}
2239
2240fn tick_headless_batch_warmup(batch: Option<ResMut<HeadlessBatchSequence>>) {
2242 let Some(mut batch) = batch else {
2243 return;
2244 };
2245
2246 if batch.warmup_frames_remaining > 0 {
2247 batch.warmup_frames_remaining -= 1;
2248 }
2249}
2250
2251fn extract_and_continue_headless_batch(
2253 mut state: ResMut<RenderState>,
2254 request: Res<RenderRequest>,
2255 buffers: (Res<SharedRgbaBuffer>, Res<SharedDepthBuffer>),
2256 batch: Option<ResMut<HeadlessBatchSequence>>,
2257 mut camera_query: Query<&mut Transform, With<RenderCamera>>,
2258 mut depth_request: ResMut<DepthCaptureRequest>,
2259 mut image_copiers: Query<&mut ImageCopier>,
2260) {
2261 let trace = render_trace_enabled();
2262 let t0 = trace.then(std::time::Instant::now);
2263
2264 let (shared_rgba, shared_depth) = buffers;
2265 let Some(mut batch) = batch else {
2266 if let Some(t0) = t0 {
2267 eprintln!(
2268 "[render_trace][sys] extract_and_continue_headless_batch skipped(no_batch) ms={:.3}",
2269 t0.elapsed().as_secs_f64() * 1000.0
2270 );
2271 }
2272 return;
2273 };
2274
2275 if state.exit_requested || !state.captured || batch.done {
2276 if let Some(t0) = t0 {
2277 eprintln!(
2278 "[render_trace][sys] extract_and_continue_headless_batch skipped(gate) captured={} done={} ms={:.3}",
2279 state.captured,
2280 batch.done,
2281 t0.elapsed().as_secs_f64() * 1000.0
2282 );
2283 }
2284 return;
2285 }
2286
2287 if let (Some(rgba), Some(depth)) = (&state.rgba_data, &state.depth_data) {
2288 let width = state.image_width;
2289 let height = state.image_height;
2290
2291 let intrinsics = request.config.intrinsics_for_size(width, height);
2292
2293 let output = RenderOutput {
2294 rgba: rgba.clone(),
2295 depth: depth.clone(),
2296 width,
2297 height,
2298 intrinsics,
2299 camera_transform: batch
2300 .current_viewpoint()
2301 .unwrap_or(request.camera_transform),
2302 object_rotation: request.object_rotation.clone(),
2303 };
2304 batch.outputs.push(output);
2305
2306 let next_index = batch.current_index + 1;
2307 if next_index >= batch.viewpoints.len() {
2308 batch.done = true;
2309 state.exit_requested = true;
2310 return;
2311 }
2312
2313 batch.current_index = next_index;
2314 batch.warmup_frames_remaining = BATCH_WARMUP_FRAMES;
2315
2316 if let Some(next_viewpoint) = batch.current_viewpoint() {
2317 for mut camera_transform in camera_query.iter_mut() {
2318 *camera_transform = next_viewpoint;
2319 }
2320 }
2321
2322 if let Ok(mut guard) = shared_rgba.0.lock() {
2323 *guard = None;
2324 }
2325 if let Ok(mut guard) = shared_depth.0.lock() {
2326 *guard = None;
2327 }
2328
2329 for mut copier in image_copiers.iter_mut() {
2330 copier.enabled = false;
2331 }
2332
2333 depth_request.requested = false;
2334 state.frame_count = 0;
2335 state.capture_ready = true;
2336 state.screenshot_requested = false;
2337 state.captured = false;
2338 state.rgba_data = None;
2339 state.depth_data = None;
2340 state.image_width = 0;
2341 state.image_height = 0;
2342
2343 if let Some(t0) = t0 {
2344 eprintln!(
2345 "[render_trace][sys] extract_and_continue_headless_batch extracted vp={} next={} done={} ms={:.3}",
2346 batch.current_index.saturating_sub(1),
2347 batch.current_index,
2348 batch.done,
2349 t0.elapsed().as_secs_f64() * 1000.0
2350 );
2351 }
2352 } else if let Some(t0) = t0 {
2353 eprintln!(
2354 "[render_trace][sys] extract_and_continue_headless_batch no_data ms={:.3}",
2355 t0.elapsed().as_secs_f64() * 1000.0
2356 );
2357 }
2358}
2359
2360#[derive(Component)]
2374struct SessionScene;
2375
2376fn setup_session_persistent_scene(
2381 mut commands: Commands,
2382 mut images: ResMut<Assets<Image>>,
2383 config: Res<SessionRenderConfig>,
2384) {
2385 let width = config.0.width;
2386 let height = config.0.height;
2387
2388 let size = Extent3d {
2389 width,
2390 height,
2391 depth_or_array_layers: 1,
2392 };
2393
2394 let mut render_target_image = Image::new_fill(
2395 size,
2396 TextureDimension::D2,
2397 &[0, 0, 0, 255],
2398 TextureFormat::Rgba8UnormSrgb,
2399 RenderAssetUsages::default(),
2400 );
2401 render_target_image.texture_descriptor.usage =
2402 TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_SRC | TextureUsages::RENDER_ATTACHMENT;
2403
2404 let render_target_handle = images.add(render_target_image);
2405 commands.insert_resource(RenderTargetImage(render_target_handle.clone()));
2406
2407 let fov = config.0.fov_radians();
2408 commands.spawn((
2409 Camera3d::default(),
2410 Camera {
2411 hdr: true,
2412 target: RenderTarget::Image(render_target_handle.clone()),
2413 ..default()
2414 },
2415 Projection::Perspective(PerspectiveProjection {
2416 fov,
2417 near: config.0.near_plane,
2418 far: config.0.far_plane,
2419 ..default()
2420 }),
2421 Msaa::Off,
2422 Transform::default(),
2423 Tonemapping::None,
2424 DepthPrepass,
2425 NormalPrepass,
2426 RenderCamera,
2427 ImageCopier {
2428 src_image: render_target_handle,
2429 enabled: false,
2430 },
2431 ));
2432
2433 let lighting = &config.0.lighting;
2434 commands.insert_resource(AmbientLight {
2435 color: Color::WHITE,
2436 brightness: lighting.ambient_brightness,
2437 });
2438
2439 if lighting.key_light_intensity > 0.0 {
2440 commands.spawn((
2441 PointLight {
2442 intensity: lighting.key_light_intensity,
2443 shadows_enabled: lighting.shadows_enabled,
2444 ..default()
2445 },
2446 Transform::from_xyz(
2447 lighting.key_light_position[0],
2448 lighting.key_light_position[1],
2449 lighting.key_light_position[2],
2450 ),
2451 ));
2452 }
2453
2454 if lighting.fill_light_intensity > 0.0 {
2455 commands.spawn((
2456 PointLight {
2457 intensity: lighting.fill_light_intensity,
2458 shadows_enabled: lighting.shadows_enabled,
2459 ..default()
2460 },
2461 Transform::from_xyz(
2462 lighting.fill_light_position[0],
2463 lighting.fill_light_position[1],
2464 lighting.fill_light_position[2],
2465 ),
2466 ));
2467 }
2468}
2469
2470#[derive(Resource)]
2473struct SessionRenderConfig(RenderConfig);
2474
2475pub struct RenderSession {
2499 app: App,
2500 render_config: RenderConfig,
2501 shared_rgba: SharedRgbaBuffer,
2502 shared_depth: SharedDepthBuffer,
2503 _not_send_sync: std::marker::PhantomData<*const ()>,
2504}
2505
2506impl RenderSession {
2507 pub fn new(render_config: &crate::RenderConfig) -> Result<Self, crate::RenderError> {
2512 let shared_rgba: SharedRgbaBuffer = SharedRgbaBuffer::default();
2513 let shared_depth: SharedDepthBuffer = SharedDepthBuffer::default();
2514
2515 let mut app = App::new();
2516 app.add_plugins(
2517 DefaultPlugins
2518 .set(WindowPlugin {
2519 primary_window: None,
2520 exit_condition: ExitCondition::DontExit,
2521 ..default()
2522 })
2523 .disable::<bevy::winit::WinitPlugin>()
2524 .disable::<LogPlugin>()
2525 .disable::<TerminalCtrlCHandlerPlugin>(),
2526 )
2527 .add_plugins(ObjPlugin)
2528 .add_plugins(ImageCopyPlugin {
2529 shared_rgba: shared_rgba.clone(),
2530 })
2531 .add_plugins(DepthReadbackPlugin {
2532 shared_depth: shared_depth.clone(),
2533 near: render_config.near_plane,
2534 far: render_config.far_plane,
2535 })
2536 .insert_resource(SessionRenderConfig(render_config.clone()))
2537 .insert_resource(shared_rgba.clone())
2538 .init_resource::<RenderState>()
2539 .add_systems(Startup, setup_session_persistent_scene)
2540 .add_systems(
2541 Update,
2542 (
2543 check_assets_loaded,
2544 apply_materials,
2545 tick_headless_batch_warmup,
2546 request_headless_capture,
2547 check_headless_capture_ready,
2548 extract_and_continue_headless_batch,
2549 )
2550 .chain()
2551 .run_if(bevy::ecs::schedule::common_conditions::resource_exists::<RenderRequest>),
2558 );
2559
2560 app.finish();
2561 app.cleanup();
2562
2563 app.update();
2569
2570 Ok(Self {
2571 app,
2572 render_config: render_config.clone(),
2573 shared_rgba,
2574 shared_depth,
2575 _not_send_sync: std::marker::PhantomData,
2576 })
2577 }
2578
2579 pub fn render(
2587 &mut self,
2588 requests: &[crate::BatchRenderRequest],
2589 ) -> Result<Vec<crate::BatchRenderOutput>, crate::BatchRenderError> {
2590 use crate::{BatchRenderError, BatchRenderOutput};
2591
2592 if requests.is_empty() {
2593 return Ok(Vec::new());
2594 }
2595
2596 let first = &requests[0];
2598 if first.render_config != self.render_config {
2599 return Err(BatchRenderError::InvalidConfig(
2600 "RenderSession render_config mismatch: session was constructed with a different \
2601 RenderConfig than the first request carries. Session config cannot change after \
2602 `new()`; construct a new session if you need a different resolution/camera."
2603 .to_string(),
2604 ));
2605 }
2606 for r in &requests[1..] {
2607 if r.object_dir != first.object_dir
2608 || r.object_rotation != first.object_rotation
2609 || r.render_config != first.render_config
2610 {
2611 return Err(BatchRenderError::InvalidConfig(
2612 "Phase 1 RenderSession::render requires homogeneous requests \
2613 (same object_dir, object_rotation, and render_config across the batch). \
2614 Call render() once per group instead."
2615 .to_string(),
2616 ));
2617 }
2618 }
2619
2620 let object_dir = std::fs::canonicalize(&first.object_dir).map_err(|e| {
2624 BatchRenderError::InvalidConfig(format!(
2625 "Cannot canonicalize object directory {}: {}",
2626 first.object_dir.display(),
2627 e
2628 ))
2629 })?;
2630 let mesh_path = object_dir.join(GOOGLE_16K_MESH_RELATIVE);
2631 let texture_path = object_dir.join(GOOGLE_16K_TEXTURE_RELATIVE);
2632 if !mesh_path.exists() {
2633 return Err(BatchRenderError::InvalidConfig(format!(
2634 "Mesh not found: {}",
2635 mesh_path.display()
2636 )));
2637 }
2638 if !texture_path.exists() {
2639 return Err(BatchRenderError::InvalidConfig(format!(
2640 "Texture not found: {}",
2641 texture_path.display()
2642 )));
2643 }
2644
2645 let viewpoints: Vec<Transform> = requests.iter().map(|r| r.viewpoint).collect();
2646
2647 {
2649 let world = self.app.world_mut();
2650
2651 let stale: Vec<Entity> = world
2653 .query_filtered::<Entity, With<SessionScene>>()
2654 .iter(world)
2655 .collect();
2656 for entity in stale {
2657 world.entity_mut(entity).despawn_recursive();
2658 }
2659
2660 if let Ok(mut guard) = self.shared_rgba.0.lock() {
2663 *guard = None;
2664 }
2665 if let Ok(mut guard) = self.shared_depth.0.lock() {
2666 *guard = None;
2667 }
2668
2669 *world.resource_mut::<RenderState>() = RenderState::default();
2672
2673 let new_request = RenderRequest {
2676 mesh_path: mesh_path.display().to_string(),
2677 texture_path: texture_path.display().to_string(),
2678 camera_transform: viewpoints[0],
2679 object_rotation: first.object_rotation.clone(),
2680 config: self.render_config.clone(),
2681 };
2682 world.insert_resource(new_request);
2683
2684 let asset_server = world.resource::<AssetServer>().clone();
2687 let scene_handle: Handle<Scene> = asset_server.load(mesh_path.display().to_string());
2688 let texture_handle: Handle<Image> =
2689 asset_server.load(texture_path.display().to_string());
2690 world.insert_resource(LoadedScene(scene_handle.clone()));
2691 world.insert_resource(LoadedTexture(texture_handle));
2692
2693 world.spawn((
2696 SceneRoot(scene_handle),
2697 Transform::from_rotation(first.object_rotation.to_quat()),
2698 RenderedObject,
2699 SessionScene,
2700 ));
2701
2702 let camera_entity = world
2706 .query_filtered::<Entity, With<RenderCamera>>()
2707 .iter(world)
2708 .next();
2709 if let Some(cam) = camera_entity {
2710 if let Some(mut transform) = world.entity_mut(cam).get_mut::<Transform>() {
2711 *transform = viewpoints[0];
2712 }
2713 }
2714
2715 world.insert_resource(HeadlessBatchSequence::new(viewpoints.clone()));
2717 }
2718
2719 let timeout = std::time::Duration::from_secs(RENDER_TIMEOUT_SECS);
2721 let start = std::time::Instant::now();
2722 loop {
2723 if start.elapsed() > timeout {
2724 return Err(BatchRenderError::TotalFailure(format!(
2725 "RenderSession::render timed out after {}s",
2726 RENDER_TIMEOUT_SECS
2727 )));
2728 }
2729
2730 self.app.update();
2731
2732 if self.app.world().resource::<HeadlessBatchSequence>().done {
2733 break;
2734 }
2735 }
2736
2737 let mut sequence = self.app.world_mut().resource_mut::<HeadlessBatchSequence>();
2740 if sequence.outputs.len() != requests.len() {
2741 return Err(BatchRenderError::TotalFailure(format!(
2742 "RenderSession produced {} outputs for {} requests",
2743 sequence.outputs.len(),
2744 requests.len()
2745 )));
2746 }
2747 let outputs = std::mem::take(&mut sequence.outputs);
2748
2749 Ok(requests
2750 .iter()
2751 .cloned()
2752 .zip(outputs)
2753 .map(|(req, out)| BatchRenderOutput::from_render_output(req, out))
2754 .collect())
2755 }
2756}
2757
2758#[derive(Component)]
2779struct PersistentScene;
2780
2781pub struct PersistentRenderer {
2799 app: App,
2800 object_dir: PathBuf,
2801 render_config: RenderConfig,
2802 shared_rgba: SharedRgbaBuffer,
2803 shared_depth: SharedDepthBuffer,
2804 _not_send_sync: std::marker::PhantomData<*const ()>,
2805}
2806
2807impl PersistentRenderer {
2808 pub fn new(
2813 object_dir: &Path,
2814 render_config: &RenderConfig,
2815 ) -> Result<Self, crate::RenderError> {
2816 let object_dir =
2817 std::fs::canonicalize(object_dir).map_err(|e| crate::RenderError::FileNotFound {
2818 path: object_dir.display().to_string(),
2819 reason: e.to_string(),
2820 })?;
2821 let mesh_path = object_dir.join(GOOGLE_16K_MESH_RELATIVE);
2822 let texture_path = object_dir.join(GOOGLE_16K_TEXTURE_RELATIVE);
2823 if !mesh_path.exists() {
2824 return Err(crate::RenderError::MeshNotFound(
2825 mesh_path.display().to_string(),
2826 ));
2827 }
2828 if !texture_path.exists() {
2829 return Err(crate::RenderError::TextureNotFound(
2830 texture_path.display().to_string(),
2831 ));
2832 }
2833
2834 let shared_rgba: SharedRgbaBuffer = SharedRgbaBuffer::default();
2835 let shared_depth: SharedDepthBuffer = SharedDepthBuffer::default();
2836
2837 let mut app = App::new();
2838 app.add_plugins(
2839 DefaultPlugins
2840 .set(WindowPlugin {
2841 primary_window: None,
2842 exit_condition: ExitCondition::DontExit,
2843 ..default()
2844 })
2845 .disable::<bevy::winit::WinitPlugin>()
2846 .disable::<LogPlugin>()
2847 .disable::<TerminalCtrlCHandlerPlugin>(),
2848 )
2849 .add_plugins(ObjPlugin)
2850 .add_plugins(ImageCopyPlugin {
2851 shared_rgba: shared_rgba.clone(),
2852 })
2853 .add_plugins(DepthReadbackPlugin {
2854 shared_depth: shared_depth.clone(),
2855 near: render_config.near_plane,
2856 far: render_config.far_plane,
2857 })
2858 .insert_resource(SessionRenderConfig(render_config.clone()))
2859 .insert_resource(shared_rgba.clone())
2860 .init_resource::<RenderState>()
2861 .add_systems(Startup, setup_session_persistent_scene)
2862 .add_systems(
2863 Update,
2864 (
2865 check_assets_loaded,
2866 apply_materials,
2867 tick_headless_batch_warmup,
2868 request_headless_capture,
2869 check_headless_capture_ready,
2870 extract_and_continue_headless_batch,
2871 )
2872 .chain()
2873 .run_if(bevy::ecs::schedule::common_conditions::resource_exists::<RenderRequest>),
2877 );
2878
2879 app.finish();
2880 app.cleanup();
2881 app.update();
2883
2884 let initial_request = RenderRequest {
2888 mesh_path: mesh_path.display().to_string(),
2889 texture_path: texture_path.display().to_string(),
2890 camera_transform: Transform::default(),
2891 object_rotation: ObjectRotation::identity(),
2892 config: render_config.clone(),
2893 };
2894
2895 {
2896 let world = app.world_mut();
2897 let asset_server = world.resource::<AssetServer>().clone();
2898 let scene_handle: Handle<Scene> = asset_server.load(mesh_path.display().to_string());
2899 let texture_handle: Handle<Image> =
2900 asset_server.load(texture_path.display().to_string());
2901 world.insert_resource(LoadedScene(scene_handle.clone()));
2902 world.insert_resource(LoadedTexture(texture_handle));
2903 world.insert_resource(initial_request);
2904 world.spawn((
2905 SceneRoot(scene_handle),
2906 Transform::from_rotation(ObjectRotation::identity().to_quat()),
2907 RenderedObject,
2908 PersistentScene,
2909 ));
2910 world.insert_resource(HeadlessBatchSequence::new(vec![Transform::default()]));
2911 }
2912
2913 let timeout = std::time::Duration::from_secs(RENDER_TIMEOUT_SECS);
2915 let start = std::time::Instant::now();
2916 loop {
2917 if start.elapsed() > timeout {
2918 return Err(crate::RenderError::RenderFailed(format!(
2919 "PersistentRenderer::new warmup render timed out after {RENDER_TIMEOUT_SECS}s"
2920 )));
2921 }
2922 app.update();
2923 if app.world().resource::<HeadlessBatchSequence>().done {
2924 break;
2925 }
2926 }
2927 app.world_mut()
2930 .resource_mut::<HeadlessBatchSequence>()
2931 .outputs
2932 .clear();
2933
2934 Ok(Self {
2935 app,
2936 object_dir,
2937 render_config: render_config.clone(),
2938 shared_rgba,
2939 shared_depth,
2940 _not_send_sync: std::marker::PhantomData,
2941 })
2942 }
2943
2944 pub fn render(
2947 &mut self,
2948 camera_transform: &Transform,
2949 object_rotation: &ObjectRotation,
2950 ) -> Result<RenderOutput, crate::RenderError> {
2951 let camera_transform = *camera_transform;
2952 let object_rotation_owned = object_rotation.clone();
2953
2954 {
2955 let world = self.app.world_mut();
2956
2957 let scene_entity = world
2961 .query_filtered::<Entity, With<PersistentScene>>()
2962 .iter(world)
2963 .next();
2964 if let Some(entity) = scene_entity {
2965 if let Some(mut transform) = world.entity_mut(entity).get_mut::<Transform>() {
2966 *transform = Transform::from_rotation(object_rotation_owned.to_quat());
2967 }
2968 }
2969
2970 let cam_entity = world
2972 .query_filtered::<Entity, With<RenderCamera>>()
2973 .iter(world)
2974 .next();
2975 if let Some(cam) = cam_entity {
2976 if let Some(mut transform) = world.entity_mut(cam).get_mut::<Transform>() {
2977 *transform = camera_transform;
2978 }
2979 }
2980
2981 {
2996 let mut state = world.resource_mut::<RenderState>();
2997 state.exit_requested = false;
2998 state.screenshot_requested = false;
2999 state.captured = false;
3000 state.rgba_data = None;
3001 state.depth_data = None;
3002 state.frame_count = 0;
3003 state.image_width = 0;
3004 state.image_height = 0;
3005 state.capture_ready = true;
3006 }
3007
3008 if let Ok(mut guard) = self.shared_rgba.0.lock() {
3011 *guard = None;
3012 }
3013 if let Ok(mut guard) = self.shared_depth.0.lock() {
3014 *guard = None;
3015 }
3016
3017 {
3020 let mut req = world.resource_mut::<RenderRequest>();
3021 req.camera_transform = camera_transform;
3022 req.object_rotation = object_rotation_owned.clone();
3023 }
3024
3025 let mut batch = HeadlessBatchSequence::new(vec![camera_transform]);
3029 batch.warmup_frames_remaining = PERSISTENT_WARMUP_FRAMES;
3030 world.insert_resource(batch);
3031 }
3032
3033 let timeout = std::time::Duration::from_secs(RENDER_TIMEOUT_SECS);
3034 let start = std::time::Instant::now();
3035 loop {
3036 if start.elapsed() > timeout {
3037 return Err(crate::RenderError::RenderFailed(format!(
3038 "PersistentRenderer::render timed out after {RENDER_TIMEOUT_SECS}s"
3039 )));
3040 }
3041 self.app.update();
3042 if self.app.world().resource::<HeadlessBatchSequence>().done {
3043 break;
3044 }
3045 }
3046
3047 let mut sequence = self.app.world_mut().resource_mut::<HeadlessBatchSequence>();
3048 let mut outputs = std::mem::take(&mut sequence.outputs);
3049 if outputs.len() != 1 {
3050 return Err(crate::RenderError::RenderFailed(format!(
3051 "PersistentRenderer::render expected 1 output, got {}",
3052 outputs.len()
3053 )));
3054 }
3055
3056 Ok(outputs.remove(0))
3057 }
3058
3059 pub fn object_dir(&self) -> &Path {
3061 &self.object_dir
3062 }
3063
3064 pub fn render_config(&self) -> &RenderConfig {
3066 &self.render_config
3067 }
3068
3069 pub fn close(self) {
3072 }
3074}
3075
3076pub fn render_to_files(
3081 object_dir: &Path,
3082 camera_transform: &Transform,
3083 object_rotation: &ObjectRotation,
3084 config: &RenderConfig,
3085 rgba_path: &Path,
3086 depth_path: &Path,
3087) -> Result<(), RenderError> {
3088 let mesh_path = object_dir.join(GOOGLE_16K_MESH_RELATIVE);
3089 let texture_path = object_dir.join(GOOGLE_16K_TEXTURE_RELATIVE);
3090
3091 if !mesh_path.exists() {
3092 return Err(RenderError::MeshNotFound(mesh_path.display().to_string()));
3093 }
3094 if !texture_path.exists() {
3095 return Err(RenderError::TextureNotFound(
3096 texture_path.display().to_string(),
3097 ));
3098 }
3099
3100 let request = RenderRequest {
3101 mesh_path: mesh_path.display().to_string(),
3102 texture_path: texture_path.display().to_string(),
3103 camera_transform: *camera_transform,
3104 object_rotation: object_rotation.clone(),
3105 config: config.clone(),
3106 };
3107
3108 let shared_output: SharedOutput = SharedOutput(Arc::new(Mutex::new(None)));
3110 let output_poll = shared_output.clone();
3111
3112 let rgba_path = rgba_path.to_path_buf();
3114 let depth_path = depth_path.to_path_buf();
3115
3116 let shared_rgba: SharedRgbaBuffer = SharedRgbaBuffer::default();
3118
3119 let shared_depth: SharedDepthBuffer = SharedDepthBuffer::default();
3121
3122 std::thread::spawn(move || {
3124 let timeout = std::time::Duration::from_secs(RENDER_TIMEOUT_SECS);
3125 let start = std::time::Instant::now();
3126 let poll_interval = std::time::Duration::from_millis(100);
3127
3128 loop {
3129 if let Ok(guard) = output_poll.0.lock() {
3130 if let Some(output) = guard.as_ref() {
3131 if let Err(e) =
3133 save_rgba_to_png(&output.rgba, output.width, output.height, &rgba_path)
3134 {
3135 eprintln!("Failed to save RGBA: {:?}", e);
3136 std::process::exit(1);
3137 }
3138
3139 if let Err(e) = save_depth_to_binary(&output.depth, &depth_path) {
3141 eprintln!("Failed to save depth: {:?}", e);
3142 std::process::exit(1);
3143 }
3144
3145 std::process::exit(0);
3146 }
3147 }
3148
3149 if start.elapsed() > timeout {
3150 eprintln!(
3151 "Error: Render timeout after {} seconds",
3152 RENDER_TIMEOUT_SECS
3153 );
3154 eprintln!("Debug info: This may indicate GPU issues, missing assets, or insufficient system resources.");
3155 std::process::exit(1);
3156 }
3157
3158 std::thread::sleep(poll_interval);
3159 }
3160 });
3161
3162 static BACKEND_INIT: OnceLock<()> = OnceLock::new();
3168 BACKEND_INIT.get_or_init(|| {
3169 let backend_config = BackendConfig::headless();
3170 backend_config.apply_env();
3171 });
3172
3173 build_headless_app(request, shared_output, shared_rgba, shared_depth).run();
3175
3176 Err(RenderError::RenderFailed(
3178 "Render did not complete".to_string(),
3179 ))
3180}
3181
3182fn save_rgba_to_png(rgba: &[u8], width: u32, height: u32, path: &Path) -> Result<(), String> {
3184 use image::{ImageBuffer, Rgba};
3185
3186 if let Some(parent) = path.parent() {
3188 std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
3189 }
3190
3191 let img: ImageBuffer<Rgba<u8>, Vec<u8>> =
3192 ImageBuffer::from_raw(width, height, rgba.to_vec())
3193 .ok_or_else(|| "Failed to create image buffer".to_string())?;
3194
3195 img.save(path).map_err(|e| e.to_string())
3196}
3197
3198fn save_depth_to_binary(depth: &[f64], path: &Path) -> Result<(), String> {
3200 if let Some(parent) = path.parent() {
3202 std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
3203 }
3204
3205 let bytes: Vec<u8> = depth.iter().flat_map(|f| f.to_le_bytes()).collect();
3206 std::fs::write(path, &bytes).map_err(|e| e.to_string())
3207}
3208
3209#[cfg(test)]
3210mod smoke_tests {
3211 use super::{headless_scene_setup_count, reset_headless_scene_setup_count};
3212 use crate::{
3213 BatchRenderConfig, BatchRenderRequest, ObjectRotation, RenderConfig, ViewpointConfig,
3214 };
3215 use image::{ImageBuffer, Rgba};
3216 use tempfile::TempDir;
3217
3218 fn write_synthetic_object() -> TempDir {
3219 let temp_dir = TempDir::new().expect("create temp dir for synthetic object");
3220 let object_dir = temp_dir.path().join("synthetic_cube").join("google_16k");
3221 std::fs::create_dir_all(&object_dir).expect("create synthetic google_16k dir");
3222
3223 let obj = r#"o SyntheticCube
3226v -0.10 -0.10 0.10
3227v 0.10 -0.10 0.10
3228v 0.10 0.10 0.10
3229v -0.10 0.10 0.10
3230v -0.10 -0.10 -0.10
3231v 0.10 -0.10 -0.10
3232v 0.10 0.10 -0.10
3233v -0.10 0.10 -0.10
3234vt 0.0 0.0
3235vt 1.0 0.0
3236vt 1.0 1.0
3237vt 0.0 1.0
3238f 1/1 2/2 3/3
3239f 1/1 3/3 4/4
3240f 6/1 5/2 8/3
3241f 6/1 8/3 7/4
3242f 2/1 6/2 7/3
3243f 2/1 7/3 3/4
3244f 5/1 1/2 4/3
3245f 5/1 4/3 8/4
3246f 4/1 3/2 7/3
3247f 4/1 7/3 8/4
3248f 5/1 6/2 2/3
3249f 5/1 2/3 1/4
3250"#;
3251 std::fs::write(object_dir.join("textured.obj"), obj).expect("write synthetic obj");
3252
3253 let texture = ImageBuffer::from_fn(2, 2, |x, y| match (x, y) {
3254 (0, 0) => Rgba([255u8, 48, 48, 255]),
3255 (1, 0) => Rgba([48u8, 255, 48, 255]),
3256 (0, 1) => Rgba([48u8, 48, 255, 255]),
3257 _ => Rgba([255u8, 255, 64, 255]),
3258 });
3259 texture
3260 .save(object_dir.join("texture_map.png"))
3261 .expect("write synthetic texture");
3262
3263 temp_dir
3264 }
3265
3266 #[test]
3267 #[ignore = "headless throughput smoke check is opt-in because it needs a local render backend"]
3268 fn test_headless_batch_throughput_smoke() {
3269 crate::initialize();
3270 reset_headless_scene_setup_count();
3271
3272 let object_root = write_synthetic_object();
3273 let object_dir = object_root.path().join("synthetic_cube");
3274 let viewpoints = crate::generate_viewpoints(&ViewpointConfig::default());
3275 let request_count = 5usize;
3276 let config = RenderConfig::tbp_default();
3277
3278 let requests: Vec<_> = viewpoints
3279 .iter()
3280 .take(request_count)
3281 .copied()
3282 .map(|viewpoint| BatchRenderRequest {
3283 object_dir: object_dir.clone(),
3284 viewpoint,
3285 object_rotation: ObjectRotation::identity(),
3286 render_config: config.clone(),
3287 })
3288 .collect();
3289
3290 let start = std::time::Instant::now();
3291 let outputs = crate::render_batch(requests, &BatchRenderConfig::default())
3292 .expect("synthetic headless batch render should succeed");
3293 let elapsed = start.elapsed();
3294
3295 assert_eq!(outputs.len(), request_count);
3296 assert_eq!(
3300 headless_scene_setup_count(),
3301 1,
3302 "homogeneous batch smoke check should reuse one headless app setup"
3303 );
3304
3305 for (idx, output) in outputs.iter().enumerate() {
3306 assert_eq!(output.width, config.width, "output {idx} width mismatch");
3307 assert_eq!(output.height, config.height, "output {idx} height mismatch");
3308 assert_eq!(
3309 output.rgba.len(),
3310 (config.width * config.height * 4) as usize,
3311 "output {idx} rgba size mismatch"
3312 );
3313 assert_eq!(
3314 output.depth.len(),
3315 (config.width * config.height) as usize,
3316 "output {idx} depth size mismatch"
3317 );
3318 assert!(
3319 output
3320 .rgba
3321 .chunks_exact(4)
3322 .any(|px| px[0] != 0 || px[1] != 0 || px[2] != 0),
3323 "output {idx} should contain visible color"
3324 );
3325 }
3326
3327 assert!(
3331 elapsed < std::time::Duration::from_secs(8),
3332 "5 synthetic headless captures took {:.2}s, expected < 8.0s",
3333 elapsed.as_secs_f64()
3334 );
3335 }
3336}