use bevy::app::{ScheduleRunnerPlugin, TerminalCtrlCHandlerPlugin};
use bevy::asset::LoadState;
use bevy::core_pipeline::prepass::{DepthPrepass, NormalPrepass};
use bevy::core_pipeline::tonemapping::Tonemapping;
use bevy::ecs::query::QueryItem;
use bevy::log::LogPlugin;
use bevy::prelude::*;
use bevy::render::camera::{ExtractedCamera, RenderTarget};
use bevy::render::render_asset::{RenderAssetUsages, RenderAssets};
use bevy::render::render_graph::{
Node, NodeRunError, RenderGraphApp, RenderGraphContext, RenderLabel, ViewNode, ViewNodeRunner,
};
use bevy::render::render_resource::{
Buffer, BufferDescriptor, BufferUsages, CommandEncoderDescriptor, Extent3d, ImageCopyBuffer,
ImageCopyTexture, ImageDataLayout, MapMode, Origin3d, TextureAspect, TextureDimension,
TextureFormat, TextureUsages,
};
use bevy::render::renderer::RenderQueue;
use bevy::render::renderer::{RenderContext, RenderDevice};
use bevy::render::texture::GpuImage;
use bevy::render::view::screenshot::{Screenshot, ScreenshotCaptured};
use bevy::render::view::ViewDepthTexture;
use bevy::render::{Extract, Render, RenderApp, RenderSet};
use bevy::window::{ExitCondition, WindowPlugin};
use bevy_obj::ObjPlugin;
use std::fs::File;
use std::io::Read as IoRead;
use std::path::{Path, PathBuf};
#[cfg(test)]
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;
use crate::{
backend::BackendConfig, ObjectRotation, RenderConfig, RenderError, RenderOutput,
TargetingPolicy,
};
use ycbust::{GOOGLE_16K_MESH_RELATIVE, GOOGLE_16K_TEXTURE_RELATIVE};
const RENDER_TIMEOUT_SECS: u64 = 180;
const BATCH_WARMUP_FRAMES: u32 = 1;
const PERSISTENT_WARMUP_FRAMES: u32 = 3;
#[inline]
fn render_trace_enabled() -> bool {
std::env::var("BEVY_SENSOR_RENDER_TRACE").is_ok()
}
#[allow(dead_code)]
fn display_available() -> bool {
std::env::var("DISPLAY").is_ok() || std::env::var("WAYLAND_DISPLAY").is_ok()
}
#[allow(dead_code)]
fn is_wsl2() -> bool {
if let Ok(version) = std::fs::read_to_string("/proc/version") {
return version.to_lowercase().contains("microsoft")
|| version.to_lowercase().contains("wsl");
}
false
}
#[derive(Resource, Default)]
struct RenderState {
frame_count: u32,
scene_loaded: bool,
texture_loaded: bool,
materials_applied: bool,
materials_applied_frame: u32,
capture_ready: bool,
screenshot_requested: bool,
captured: bool,
exit_requested: bool,
#[allow(dead_code)]
exit_frame_count: u32,
rgba_data: Option<Vec<u8>>,
depth_data: Option<Vec<f64>>,
image_width: u32,
image_height: u32,
}
#[cfg(test)]
static HEADLESS_SCENE_SETUP_COUNT: AtomicUsize = AtomicUsize::new(0);
#[cfg(test)]
fn reset_headless_scene_setup_count() {
HEADLESS_SCENE_SETUP_COUNT.store(0, Ordering::SeqCst);
}
#[cfg(test)]
fn headless_scene_setup_count() -> usize {
HEADLESS_SCENE_SETUP_COUNT.load(Ordering::SeqCst)
}
#[derive(Resource, Clone)]
#[allow(clippy::type_complexity)]
#[allow(dead_code)]
struct SharedImageBuffer(Arc<Mutex<Option<(Vec<u8>, u32, u32)>>>);
#[derive(Resource, Clone, Default)]
#[allow(clippy::type_complexity)]
struct SharedDepthBuffer(Arc<Mutex<Option<(Vec<f64>, u32, u32)>>>);
#[derive(Resource, Default, Clone)]
struct DepthCaptureRequest {
requested: bool,
near: f32,
far: f32,
}
struct PendingDepthCapture {
buffer: Buffer,
width: u32,
height: u32,
near: f32,
far: f32,
}
#[derive(Resource, Default)]
struct PendingDepthCaptureQueue(Arc<Mutex<Vec<PendingDepthCapture>>>);
mod depth_helpers {
pub const COPY_BYTES_PER_ROW_ALIGNMENT: u32 = 256;
pub fn align_byte_size(value: u32) -> u32 {
let remainder = value % COPY_BYTES_PER_ROW_ALIGNMENT;
if remainder == 0 {
value
} else {
value + (COPY_BYTES_PER_ROW_ALIGNMENT - remainder)
}
}
#[allow(dead_code)]
pub fn get_aligned_size(width: u32, height: u32, pixel_size: u32) -> u32 {
height * align_byte_size(width * pixel_size)
}
pub fn reverse_z_to_linear_depth(ndc_depth: f32, near: f32, far: f32) -> f32 {
if ndc_depth <= 0.0 {
return far; }
if ndc_depth >= 1.0 {
return near; }
far / (1.0 + ndc_depth * (far / near - 1.0))
}
pub fn extract_depth_with_alignment(data: &[u8], width: u32, height: u32) -> Vec<f32> {
let pixel_size = 4u32; let aligned_row_bytes = align_byte_size(width * pixel_size) as usize;
let actual_row_bytes = (width * pixel_size) as usize;
let mut depth_values = Vec::with_capacity((width * height) as usize);
for y in 0..height as usize {
let row_start = y * aligned_row_bytes;
let row_data = &data[row_start..row_start + actual_row_bytes];
for x in 0..width as usize {
let offset = x * 4;
let bytes: [u8; 4] = row_data[offset..offset + 4].try_into().unwrap();
let depth_value = f32::from_le_bytes(bytes);
depth_values.push(depth_value);
}
}
depth_values
}
pub fn convert_depth_to_linear(raw_depth: &[f32], near: f32, far: f32) -> Vec<f64> {
raw_depth
.iter()
.map(|&ndc| reverse_z_to_linear_depth(ndc, near, far) as f64)
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_align_byte_size() {
assert_eq!(align_byte_size(256), 256);
assert_eq!(align_byte_size(257), 512);
assert_eq!(align_byte_size(1), 256);
assert_eq!(align_byte_size(512), 512);
assert_eq!(align_byte_size(0), 0);
}
#[test]
fn test_reverse_z_to_linear_depth() {
let near = 0.01;
let far = 10.0;
let linear_near = reverse_z_to_linear_depth(1.0, near, far);
assert!((linear_near - near).abs() < 0.001);
let linear_mid = reverse_z_to_linear_depth(0.5, near, far);
assert!(linear_mid > near && linear_mid < far);
let linear_almost_far = reverse_z_to_linear_depth(0.0001, near, far);
assert!(linear_almost_far > 9.0);
let background = reverse_z_to_linear_depth(0.0, near, far);
assert_eq!(background, far);
}
#[test]
fn test_extract_depth_with_alignment() {
let width = 2u32;
let height = 2u32;
let mut data = vec![0u8; 256 * 2];
data[0..4].copy_from_slice(&0.5f32.to_le_bytes());
data[4..8].copy_from_slice(&0.6f32.to_le_bytes());
data[256..260].copy_from_slice(&0.7f32.to_le_bytes());
data[260..264].copy_from_slice(&0.8f32.to_le_bytes());
let depth = extract_depth_with_alignment(&data, width, height);
assert_eq!(depth.len(), 4);
assert!((depth[0] - 0.5).abs() < 0.001);
assert!((depth[1] - 0.6).abs() < 0.001);
assert!((depth[2] - 0.7).abs() < 0.001);
assert!((depth[3] - 0.8).abs() < 0.001);
}
#[test]
fn test_reverse_z_depth_at_near_plane() {
let near = 0.01;
let far = 100.0;
let depth = reverse_z_to_linear_depth(1.0, near, far);
assert!((depth - near).abs() < 0.0001);
}
#[test]
fn test_reverse_z_depth_at_far_plane() {
let near = 0.01;
let far = 100.0;
let depth = reverse_z_to_linear_depth(0.0, near, far);
assert!((depth - far).abs() < 0.0001);
}
#[test]
fn test_reverse_z_monotonic() {
let near = 0.01;
let far = 10.0;
let mut prev_depth = 0.0;
for i in (0..=100).rev() {
let ndc = i as f32 / 100.0;
let depth = reverse_z_to_linear_depth(ndc, near, far);
assert!(
depth >= prev_depth,
"Depth should be monotonic: ndc={}, depth={}, prev={}",
ndc,
depth,
prev_depth
);
prev_depth = depth;
}
}
#[test]
fn test_convert_depth_to_linear_batch() {
let near = 0.01f32;
let far = 10.0f32;
let ndc_depths = vec![1.0f32, 0.5, 0.1, 0.0];
let linear = convert_depth_to_linear(&ndc_depths, near, far);
assert_eq!(linear.len(), 4);
assert!((linear[0] - near as f64).abs() < 0.001);
assert!((linear[3] - far as f64).abs() < 0.001);
for d in &linear {
assert!(*d >= near as f64 && *d <= far as f64);
}
}
#[test]
fn test_align_byte_size_edge_cases() {
assert_eq!(align_byte_size(256), 256);
assert_eq!(align_byte_size(512), 512);
assert_eq!(align_byte_size(1024), 1024);
assert_eq!(align_byte_size(255), 256);
assert_eq!(align_byte_size(128), 256);
assert_eq!(align_byte_size(300), 512);
}
#[test]
fn test_extract_depth_64x64() {
let width = 64u32;
let height = 64u32;
let bytes_per_pixel = 4u32;
let padded_row = align_byte_size(width * bytes_per_pixel);
let mut data = vec![0u8; (padded_row * height) as usize];
for y in 0..height {
for x in 0..width {
let value = (y * width + x) as f32 / (width * height) as f32;
let offset = (y * padded_row + x * bytes_per_pixel) as usize;
data[offset..offset + 4].copy_from_slice(&value.to_le_bytes());
}
}
let depth = extract_depth_with_alignment(&data, width, height);
assert_eq!(depth.len(), (width * height) as usize);
assert!((depth[0] - 0.0).abs() < 0.001);
let expected_last = (width * height - 1) as f32 / (width * height) as f32;
assert!((depth[(width * height - 1) as usize] - expected_last).abs() < 0.001);
}
}
}
#[derive(Debug, Hash, PartialEq, Eq, Clone, bevy::render::render_graph::RenderLabel)]
struct DepthReadbackLabel;
#[derive(Default)]
struct DepthReadbackNode;
impl ViewNode for DepthReadbackNode {
type ViewQuery = (&'static ViewDepthTexture, &'static ExtractedCamera);
fn run<'w>(
&self,
_graph: &mut RenderGraphContext,
render_context: &mut RenderContext<'w>,
(view_depth_texture, camera): QueryItem<'w, Self::ViewQuery>,
world: &'w World,
) -> Result<(), NodeRunError> {
let trace = render_trace_enabled();
let t0 = trace.then(std::time::Instant::now);
let Some(request) = world.get_resource::<DepthCaptureRequest>() else {
return Ok(());
};
if !request.requested {
return Ok(());
}
let Some(queue) = world.get_resource::<PendingDepthCaptureQueue>() else {
return Ok(());
};
let Some(physical_size) = camera.physical_target_size else {
return Ok(());
};
let width = physical_size.x;
let height = physical_size.y;
let render_device = world.resource::<RenderDevice>();
let bytes_per_pixel = 4u32; let unpadded_bytes_per_row = width * bytes_per_pixel;
let padded_bytes_per_row = depth_helpers::align_byte_size(unpadded_bytes_per_row);
let buffer_size = (padded_bytes_per_row * height) as u64;
let staging_buffer = render_device.create_buffer(&BufferDescriptor {
label: Some("depth_staging_buffer"),
size: buffer_size,
usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ,
mapped_at_creation: false,
});
let encoder = render_context.command_encoder();
encoder.copy_texture_to_buffer(
ImageCopyTexture {
texture: &view_depth_texture.texture,
mip_level: 0,
origin: Origin3d::ZERO,
aspect: TextureAspect::DepthOnly,
},
ImageCopyBuffer {
buffer: &staging_buffer,
layout: ImageDataLayout {
offset: 0,
bytes_per_row: Some(padded_bytes_per_row),
rows_per_image: Some(height),
},
},
Extent3d {
width,
height,
depth_or_array_layers: 1,
},
);
if let Ok(mut pending) = queue.0.lock() {
pending.push(PendingDepthCapture {
buffer: staging_buffer,
width,
height,
near: request.near,
far: request.far,
});
}
if let Some(t0) = t0 {
eprintln!(
"[render_trace][node] DepthReadbackNode ms={:.3}",
t0.elapsed().as_secs_f64() * 1000.0
);
}
Ok(())
}
}
struct DepthReadbackPlugin {
shared_depth: SharedDepthBuffer,
near: f32,
far: f32,
}
impl Plugin for DepthReadbackPlugin {
fn build(&self, app: &mut App) {
use bevy::core_pipeline::core_3d::graph::Core3d;
use bevy::core_pipeline::core_3d::graph::Node3d;
app.insert_resource(self.shared_depth.clone());
app.insert_resource(DepthCaptureRequest {
requested: false,
near: self.near,
far: self.far,
});
let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
eprintln!("Failed to get RenderApp for depth readback");
return;
};
render_app.insert_resource(self.shared_depth.clone());
render_app.init_resource::<PendingDepthCaptureQueue>();
render_app.add_systems(ExtractSchedule, extract_depth_request);
render_app.add_systems(Render, collect_depth_captures.in_set(RenderSet::Cleanup));
render_app
.add_render_graph_node::<ViewNodeRunner<DepthReadbackNode>>(Core3d, DepthReadbackLabel)
.add_render_graph_edges(
Core3d,
(Node3d::EndMainPass, DepthReadbackLabel, Node3d::Tonemapping),
);
}
}
fn extract_depth_request(mut commands: Commands, request: Extract<Res<DepthCaptureRequest>>) {
commands.insert_resource(DepthCaptureRequest {
requested: request.requested,
near: request.near,
far: request.far,
});
}
fn collect_depth_captures(
queue: Res<PendingDepthCaptureQueue>,
shared_depth: Res<SharedDepthBuffer>,
render_device: Res<RenderDevice>,
) {
let trace = render_trace_enabled();
let t_sys = trace.then(std::time::Instant::now);
let pending_captures = {
let Ok(mut pending) = queue.0.lock() else {
return;
};
std::mem::take(&mut *pending)
};
if pending_captures.is_empty() {
if let Some(t0) = t_sys {
eprintln!(
"[render_trace][sys] collect_depth_captures empty ms={:.3}",
t0.elapsed().as_secs_f64() * 1000.0
);
}
return;
}
let pending_count = pending_captures.len();
for pending in pending_captures {
let width = pending.width;
let height = pending.height;
let near = pending.near;
let far = pending.far;
let buffer = pending.buffer;
let shared = shared_depth.0.clone();
let buffer_slice = buffer.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
buffer_slice.map_async(MapMode::Read, move |result| {
let _ = tx.send(result);
});
let t_wait = trace.then(std::time::Instant::now);
let mut poll_iters: u32 = 0;
loop {
render_device.poll(bevy::render::render_resource::Maintain::Poll);
poll_iters += 1;
match rx.try_recv() {
Ok(Ok(())) => {
let data = buffer_slice.get_mapped_range();
let ndc_depth =
depth_helpers::extract_depth_with_alignment(&data, width, height);
drop(data);
buffer.unmap();
let linear_depth =
depth_helpers::convert_depth_to_linear(&ndc_depth, near, far);
if let Ok(mut guard) = shared.lock() {
*guard = Some((linear_depth, width, height));
}
break;
}
Ok(Err(e)) => {
eprintln!("Failed to map depth buffer: {:?}", e);
break;
}
Err(std::sync::mpsc::TryRecvError::Empty) => {
std::thread::sleep(std::time::Duration::from_millis(1));
}
Err(std::sync::mpsc::TryRecvError::Disconnected) => {
eprintln!("Depth buffer mapping channel disconnected");
break;
}
}
}
if let Some(t_wait) = t_wait {
eprintln!(
"[render_trace][sys] collect_depth_captures mapping_wait poll_iters={} ms={:.3}",
poll_iters,
t_wait.elapsed().as_secs_f64() * 1000.0
);
}
}
if let Some(t0) = t_sys {
eprintln!(
"[render_trace][sys] collect_depth_captures done pending={} ms={:.3}",
pending_count,
t0.elapsed().as_secs_f64() * 1000.0
);
}
}
#[derive(Debug, Hash, PartialEq, Eq, Clone, RenderLabel)]
struct ImageCopyLabel;
#[derive(Component, Clone)]
struct ImageCopier {
src_image: Handle<Image>,
enabled: bool,
}
#[derive(Resource, Default)]
struct ImageCopiers(Vec<ImageCopier>);
struct PendingImageCapture {
buffer: Buffer,
width: u32,
height: u32,
padded_bytes_per_row: u32,
}
#[derive(Resource, Default)]
struct PendingImageCaptureQueue(Arc<Mutex<Vec<PendingImageCapture>>>);
#[derive(Resource, Clone, Default)]
#[allow(clippy::type_complexity)]
struct SharedRgbaBuffer(Arc<Mutex<Option<(Vec<u8>, u32, u32)>>>);
struct ImageCopyDriver;
impl Node for ImageCopyDriver {
fn run(
&self,
_graph: &mut RenderGraphContext,
_render_context: &mut RenderContext,
world: &World,
) -> Result<(), NodeRunError> {
let trace = render_trace_enabled();
let t0 = trace.then(std::time::Instant::now);
let Some(image_copiers) = world.get_resource::<ImageCopiers>() else {
return Ok(());
};
let Some(gpu_images) = world.get_resource::<RenderAssets<GpuImage>>() else {
return Ok(());
};
let Some(queue) = world.get_resource::<PendingImageCaptureQueue>() else {
return Ok(());
};
let render_device = world.resource::<RenderDevice>();
let Some(render_queue) = world.get_resource::<RenderQueue>() else {
return Ok(());
};
for image_copier in image_copiers.0.iter() {
if !image_copier.enabled {
continue;
}
let Some(gpu_image) = gpu_images.get(&image_copier.src_image) else {
continue;
};
let width = gpu_image.size.x;
let height = gpu_image.size.y;
let block_dimensions = gpu_image.texture_format.block_dimensions();
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(
(width as usize / block_dimensions.0 as usize) * block_size as usize,
);
let buffer_size = (padded_bytes_per_row * height as usize) as u64;
let staging_buffer = render_device.create_buffer(&BufferDescriptor {
label: Some("image_copy_staging_buffer"),
size: buffer_size,
usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ,
mapped_at_creation: false,
});
let mut encoder =
render_device.create_command_encoder(&CommandEncoderDescriptor::default());
let texture_extent = Extent3d {
width,
height,
depth_or_array_layers: 1,
};
encoder.copy_texture_to_buffer(
gpu_image.texture.as_image_copy(),
ImageCopyBuffer {
buffer: &staging_buffer,
layout: ImageDataLayout {
offset: 0,
bytes_per_row: Some(padded_bytes_per_row as u32),
rows_per_image: None,
},
},
texture_extent,
);
render_queue.submit(std::iter::once(encoder.finish()));
if let Ok(mut pending) = queue.0.lock() {
pending.push(PendingImageCapture {
buffer: staging_buffer,
width,
height,
padded_bytes_per_row: padded_bytes_per_row as u32,
});
}
}
if let Some(t0) = t0 {
eprintln!(
"[render_trace][node] ImageCopyDriver ms={:.3}",
t0.elapsed().as_secs_f64() * 1000.0
);
}
Ok(())
}
}
fn extract_image_copiers(mut commands: Commands, query: Extract<Query<&ImageCopier>>) {
commands.insert_resource(ImageCopiers(query.iter().cloned().collect()));
}
fn collect_image_captures(
queue: Res<PendingImageCaptureQueue>,
shared_rgba: Res<SharedRgbaBuffer>,
render_device: Res<RenderDevice>,
) {
let trace = render_trace_enabled();
let t_sys = trace.then(std::time::Instant::now);
let pending_captures = {
let Ok(mut pending) = queue.0.lock() else {
return;
};
std::mem::take(&mut *pending)
};
if pending_captures.is_empty() {
if let Some(t0) = t_sys {
eprintln!(
"[render_trace][sys] collect_image_captures empty ms={:.3}",
t0.elapsed().as_secs_f64() * 1000.0
);
}
return;
}
let pending_count = pending_captures.len();
for pending in pending_captures {
let width = pending.width;
let height = pending.height;
let padded_bytes_per_row = pending.padded_bytes_per_row;
let buffer = pending.buffer;
let shared = shared_rgba.0.clone();
let buffer_slice = buffer.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
buffer_slice.map_async(MapMode::Read, move |result| {
let _ = tx.send(result);
});
let start = std::time::Instant::now();
let timeout = std::time::Duration::from_secs(10);
let mut poll_iters: u32 = 0;
loop {
render_device.poll(bevy::render::render_resource::Maintain::Poll);
poll_iters += 1;
if start.elapsed() > timeout {
eprintln!(
"Warning: Buffer mapping timeout after {:?}",
start.elapsed()
);
break;
}
match rx.try_recv() {
Ok(Ok(())) => {
let data = buffer_slice.get_mapped_range();
let bytes_per_pixel = 4u32;
let actual_row_bytes = (width * bytes_per_pixel) as usize;
let padded_row_bytes = padded_bytes_per_row as usize;
let mut rgba = Vec::with_capacity((width * height * 4) as usize);
for y in 0..height as usize {
let row_start = y * padded_row_bytes;
rgba.extend_from_slice(&data[row_start..row_start + actual_row_bytes]);
}
drop(data);
buffer.unmap();
if let Ok(mut guard) = shared.lock() {
*guard = Some((rgba, width, height));
}
break;
}
Ok(Err(e)) => {
eprintln!("Failed to map image buffer: {:?}", e);
break;
}
Err(std::sync::mpsc::TryRecvError::Empty) => {
std::thread::sleep(std::time::Duration::from_millis(1));
}
Err(std::sync::mpsc::TryRecvError::Disconnected) => {
eprintln!("Image buffer mapping channel disconnected");
break;
}
}
}
if trace {
eprintln!(
"[render_trace][sys] collect_image_captures mapping_wait poll_iters={} ms={:.3}",
poll_iters,
start.elapsed().as_secs_f64() * 1000.0
);
}
}
if let Some(t0) = t_sys {
eprintln!(
"[render_trace][sys] collect_image_captures done pending={} ms={:.3}",
pending_count,
t0.elapsed().as_secs_f64() * 1000.0
);
}
}
struct ImageCopyPlugin {
shared_rgba: SharedRgbaBuffer,
}
impl Plugin for ImageCopyPlugin {
fn build(&self, app: &mut App) {
use bevy::render::render_graph::RenderGraph;
app.insert_resource(self.shared_rgba.clone());
let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
return;
};
render_app.insert_resource(self.shared_rgba.clone());
render_app.init_resource::<ImageCopiers>();
render_app.init_resource::<PendingImageCaptureQueue>();
render_app.add_systems(ExtractSchedule, extract_image_copiers);
render_app.add_systems(Render, collect_image_captures.in_set(RenderSet::Cleanup));
let mut graph = render_app.world_mut().resource_mut::<RenderGraph>();
graph.add_node(ImageCopyLabel, ImageCopyDriver);
graph.add_node_edge(bevy::render::graph::CameraDriverLabel, ImageCopyLabel);
}
}
#[derive(Resource, Clone)]
struct RenderRequest {
mesh_path: String,
texture_path: String,
camera_transform: Transform,
object_rotation: ObjectRotation,
config: RenderConfig,
}
#[derive(Component)]
struct RenderedObject;
#[derive(Component)]
struct RenderCamera;
#[derive(Resource)]
struct LoadedTexture(Handle<Image>);
#[derive(Resource)]
struct LoadedScene(Handle<Scene>);
#[derive(Resource, Clone)]
struct SharedOutput(Arc<Mutex<Option<RenderOutput>>>);
#[derive(Resource)]
#[allow(dead_code)]
struct RenderTargetImage(Handle<Image>);
#[derive(Resource)]
struct HeadlessBatchSequence {
viewpoints: Vec<Transform>,
current_index: usize,
outputs: Vec<RenderOutput>,
warmup_frames_remaining: u32,
done: bool,
}
impl HeadlessBatchSequence {
fn new(viewpoints: Vec<Transform>) -> Self {
let capacity = viewpoints.len();
Self {
viewpoints,
current_index: 0,
outputs: Vec::with_capacity(capacity),
warmup_frames_remaining: 0,
done: capacity == 0,
}
}
fn current_viewpoint(&self) -> Option<Transform> {
self.viewpoints.get(self.current_index).cloned()
}
}
#[allow(dead_code)]
pub fn render_headless(
object_dir: &Path,
camera_transform: &Transform,
object_rotation: &ObjectRotation,
config: &RenderConfig,
) -> Result<RenderOutput, RenderError> {
let object_dir = std::fs::canonicalize(object_dir).map_err(|e| {
RenderError::RenderFailed(format!(
"Cannot canonicalize object directory {}: {}",
object_dir.display(),
e
))
})?;
let mesh_path = object_dir.join(GOOGLE_16K_MESH_RELATIVE);
let texture_path = object_dir.join(GOOGLE_16K_TEXTURE_RELATIVE);
if !mesh_path.exists() {
return Err(RenderError::MeshNotFound(mesh_path.display().to_string()));
}
if !texture_path.exists() {
return Err(RenderError::TextureNotFound(
texture_path.display().to_string(),
));
}
let request = RenderRequest {
mesh_path: mesh_path.display().to_string(),
texture_path: texture_path.display().to_string(),
camera_transform: *camera_transform,
object_rotation: object_rotation.clone(),
config: config.clone(),
};
let shared_output: SharedOutput = SharedOutput(Arc::new(Mutex::new(None)));
let output_clone = shared_output.clone();
let shared_rgba: SharedRgbaBuffer = SharedRgbaBuffer::default();
let shared_depth: SharedDepthBuffer = SharedDepthBuffer::default();
let temp_path =
std::env::temp_dir().join(format!("bevy_sensor_render_{}.bin", std::process::id()));
let output_poll_for_timeout = shared_output.clone();
std::thread::spawn(move || {
let timeout = std::time::Duration::from_secs(RENDER_TIMEOUT_SECS);
let start = std::time::Instant::now();
let poll_interval = std::time::Duration::from_millis(100);
loop {
if let Ok(guard) = output_poll_for_timeout.0.lock() {
if guard.is_some() {
return; }
}
if start.elapsed() > timeout {
eprintln!(
"Error: Render timeout after {} seconds",
RENDER_TIMEOUT_SECS
);
eprintln!("Debug info: This may indicate GPU issues, missing assets, or insufficient system resources.");
std::process::exit(1);
}
std::thread::sleep(poll_interval);
}
});
build_headless_app(request, output_clone, shared_rgba, shared_depth).run();
if let Ok(guard) = shared_output.0.lock() {
if let Some(output) = guard.as_ref() {
return Ok(output.clone());
}
}
if temp_path.exists() {
if let Ok(output) = read_output_from_file(&temp_path) {
let _ = std::fs::remove_file(&temp_path);
return Ok(output);
}
}
Err(RenderError::RenderFailed(
"Render did not complete".to_string(),
))
}
pub fn render_headless_sequence(
object_dir: &Path,
viewpoints: &[Transform],
object_rotation: &ObjectRotation,
config: &RenderConfig,
) -> Result<Vec<RenderOutput>, RenderError> {
if viewpoints.is_empty() {
return Ok(Vec::new());
}
let object_dir = std::fs::canonicalize(object_dir).map_err(|e| {
RenderError::RenderFailed(format!(
"Cannot canonicalize object directory {}: {}",
object_dir.display(),
e
))
})?;
let mesh_path = object_dir.join(GOOGLE_16K_MESH_RELATIVE);
let texture_path = object_dir.join(GOOGLE_16K_TEXTURE_RELATIVE);
if !mesh_path.exists() {
return Err(RenderError::MeshNotFound(mesh_path.display().to_string()));
}
if !texture_path.exists() {
return Err(RenderError::TextureNotFound(
texture_path.display().to_string(),
));
}
let request = RenderRequest {
mesh_path: mesh_path.display().to_string(),
texture_path: texture_path.display().to_string(),
camera_transform: viewpoints[0],
object_rotation: object_rotation.clone(),
config: config.clone(),
};
let shared_rgba: SharedRgbaBuffer = SharedRgbaBuffer::default();
let rgba_clone = shared_rgba.clone();
let shared_depth: SharedDepthBuffer = SharedDepthBuffer::default();
let depth_clone = shared_depth.clone();
let mut app = App::new();
app.add_plugins(
DefaultPlugins
.set(WindowPlugin {
primary_window: None,
exit_condition: ExitCondition::DontExit,
..default()
})
.disable::<bevy::winit::WinitPlugin>()
.disable::<LogPlugin>()
.disable::<TerminalCtrlCHandlerPlugin>(),
)
.add_plugins(ObjPlugin)
.add_plugins(ImageCopyPlugin {
shared_rgba: rgba_clone,
})
.add_plugins(DepthReadbackPlugin {
shared_depth: depth_clone,
near: config.near_plane,
far: config.far_plane,
})
.insert_resource(request)
.insert_resource(shared_rgba)
.insert_resource(HeadlessBatchSequence::new(viewpoints.to_vec()))
.init_resource::<RenderState>()
.add_systems(Startup, setup_headless_scene)
.add_systems(
Update,
(
check_assets_loaded,
apply_materials,
tick_headless_batch_warmup,
request_headless_capture,
check_headless_capture_ready,
extract_and_continue_headless_batch,
)
.chain(),
);
let trace_outer = render_trace_enabled();
let t_finish = std::time::Instant::now();
app.finish();
let finish_ms = t_finish.elapsed().as_secs_f64() * 1000.0;
let t_cleanup = std::time::Instant::now();
app.cleanup();
let cleanup_ms = t_cleanup.elapsed().as_secs_f64() * 1000.0;
if trace_outer {
eprintln!(
"[render_trace][coldinit] app.finish ms={:.3} app.cleanup ms={:.3}",
finish_ms, cleanup_ms
);
}
let timeout = std::time::Duration::from_secs(RENDER_TIMEOUT_SECS);
let start = std::time::Instant::now();
let trace = std::env::var("BEVY_SENSOR_RENDER_TRACE").is_ok();
let mut update_idx: u32 = 0;
let mut last_completed_outputs: usize = 0;
let mut viewpoint_start = std::time::Instant::now();
loop {
if start.elapsed() > timeout {
return Err(RenderError::RenderTimeout {
duration_secs: RENDER_TIMEOUT_SECS,
});
}
let update_start = std::time::Instant::now();
app.update();
let update_elapsed_ms = update_start.elapsed().as_secs_f64() * 1000.0;
if trace {
let batch = app.world().resource::<HeadlessBatchSequence>();
let warmup = batch.warmup_frames_remaining;
let current = batch.current_index;
let completed = batch.outputs.len();
let vp_ms = viewpoint_start.elapsed().as_secs_f64() * 1000.0;
eprintln!(
"[render_trace] update={update_idx} vp={current} warmup={warmup} \
completed={completed} update_ms={update_elapsed_ms:.2} vp_ms={vp_ms:.2}"
);
if completed > last_completed_outputs {
eprintln!(
"[render_trace] viewpoint {} finished in {:.2} ms",
completed - 1,
vp_ms
);
last_completed_outputs = completed;
viewpoint_start = std::time::Instant::now();
}
}
update_idx += 1;
if app.world().resource::<HeadlessBatchSequence>().done {
break;
}
}
if trace {
eprintln!(
"[render_trace] total_wall_ms={:.2} updates={update_idx} viewpoints={}",
start.elapsed().as_secs_f64() * 1000.0,
viewpoints.len()
);
}
let mut batch = app.world_mut().resource_mut::<HeadlessBatchSequence>();
if batch.outputs.len() != viewpoints.len() {
return Err(RenderError::RenderFailed(format!(
"Batch render produced {} outputs for {} viewpoints",
batch.outputs.len(),
viewpoints.len()
)));
}
Ok(std::mem::take(&mut batch.outputs))
}
fn build_headless_app(
request: RenderRequest,
shared_output: SharedOutput,
shared_rgba: SharedRgbaBuffer,
shared_depth: SharedDepthBuffer,
) -> App {
let near = request.config.near_plane;
let far = request.config.far_plane;
let mut app = App::new();
app.add_plugins(
DefaultPlugins
.set(WindowPlugin {
primary_window: None,
exit_condition: ExitCondition::DontExit,
..default()
})
.disable::<bevy::winit::WinitPlugin>()
.disable::<LogPlugin>()
.disable::<TerminalCtrlCHandlerPlugin>(),
)
.add_plugins(ScheduleRunnerPlugin::run_loop(Duration::from_secs_f64(
1.0 / 60.0,
)))
.add_plugins(ObjPlugin)
.add_plugins(ImageCopyPlugin {
shared_rgba: shared_rgba.clone(),
})
.add_plugins(DepthReadbackPlugin {
shared_depth,
near,
far,
})
.insert_resource(request)
.insert_resource(shared_output)
.insert_resource(shared_rgba)
.init_resource::<RenderState>()
.add_systems(Startup, setup_headless_scene)
.add_systems(
Update,
(
check_assets_loaded,
apply_materials,
request_headless_capture,
check_headless_capture_ready,
extract_and_exit_headless,
)
.chain(),
);
app
}
#[allow(dead_code)]
fn serialize_output(output: &RenderOutput) -> Vec<u8> {
let mut data = Vec::new();
data.extend_from_slice(&output.width.to_le_bytes());
data.extend_from_slice(&output.height.to_le_bytes());
data.extend_from_slice(&(output.rgba.len() as u32).to_le_bytes());
data.extend_from_slice(&(output.depth.len() as u32).to_le_bytes());
data.extend_from_slice(&output.rgba);
for d in &output.depth {
data.extend_from_slice(&d.to_le_bytes());
}
data.extend_from_slice(&output.intrinsics.focal_length[0].to_le_bytes());
data.extend_from_slice(&output.intrinsics.focal_length[1].to_le_bytes());
data.extend_from_slice(&output.intrinsics.principal_point[0].to_le_bytes());
data.extend_from_slice(&output.intrinsics.principal_point[1].to_le_bytes());
data.extend_from_slice(&output.intrinsics.image_size[0].to_le_bytes());
data.extend_from_slice(&output.intrinsics.image_size[1].to_le_bytes());
let t = output.camera_transform.translation;
let r = output.camera_transform.rotation;
data.extend_from_slice(&t.x.to_le_bytes());
data.extend_from_slice(&t.y.to_le_bytes());
data.extend_from_slice(&t.z.to_le_bytes());
data.extend_from_slice(&r.x.to_le_bytes());
data.extend_from_slice(&r.y.to_le_bytes());
data.extend_from_slice(&r.z.to_le_bytes());
data.extend_from_slice(&r.w.to_le_bytes());
let or = &output.object_rotation;
data.extend_from_slice(&or.pitch.to_le_bytes());
data.extend_from_slice(&or.yaw.to_le_bytes());
data.extend_from_slice(&or.roll.to_le_bytes());
data
}
fn read_output_from_file(path: &std::path::Path) -> Result<RenderOutput, RenderError> {
let mut file = File::open(path).map_err(|e| RenderError::RenderFailed(e.to_string()))?;
let mut data = Vec::new();
file.read_to_end(&mut data)
.map_err(|e| RenderError::RenderFailed(e.to_string()))?;
let mut cursor = 0;
let read_u32 = |data: &[u8], cursor: &mut usize| -> u32 {
let val = u32::from_le_bytes(data[*cursor..*cursor + 4].try_into().unwrap());
*cursor += 4;
val
};
let read_f32 = |data: &[u8], cursor: &mut usize| -> f32 {
let val = f32::from_le_bytes(data[*cursor..*cursor + 4].try_into().unwrap());
*cursor += 4;
val
};
let read_f64 = |data: &[u8], cursor: &mut usize| -> f64 {
let val = f64::from_le_bytes(data[*cursor..*cursor + 8].try_into().unwrap());
*cursor += 8;
val
};
let width = read_u32(&data, &mut cursor);
let height = read_u32(&data, &mut cursor);
let rgba_len = read_u32(&data, &mut cursor) as usize;
let depth_len = read_u32(&data, &mut cursor) as usize;
let rgba = data[cursor..cursor + rgba_len].to_vec();
cursor += rgba_len;
let mut depth = Vec::with_capacity(depth_len);
for _ in 0..depth_len {
depth.push(read_f64(&data, &mut cursor));
}
let focal_length = [read_f64(&data, &mut cursor), read_f64(&data, &mut cursor)];
let principal_point = [read_f64(&data, &mut cursor), read_f64(&data, &mut cursor)];
let image_size = [read_u32(&data, &mut cursor), read_u32(&data, &mut cursor)];
let tx = read_f32(&data, &mut cursor);
let ty = read_f32(&data, &mut cursor);
let tz = read_f32(&data, &mut cursor);
let rx = read_f32(&data, &mut cursor);
let ry = read_f32(&data, &mut cursor);
let rz = read_f32(&data, &mut cursor);
let rw = read_f32(&data, &mut cursor);
let pitch = read_f64(&data, &mut cursor);
let yaw = read_f64(&data, &mut cursor);
let roll = read_f64(&data, &mut cursor);
Ok(RenderOutput {
rgba,
depth,
width,
height,
intrinsics: crate::CameraIntrinsics {
focal_length,
principal_point,
image_size,
},
camera_transform: Transform {
translation: Vec3::new(tx, ty, tz),
rotation: Quat::from_xyzw(rx, ry, rz, rw),
scale: Vec3::ONE,
},
object_rotation: ObjectRotation { pitch, yaw, roll },
target_point: Vec3::ZERO,
targeting_policy: TargetingPolicy::Origin,
})
}
#[allow(dead_code)]
fn setup_scene(
mut commands: Commands,
asset_server: Res<AssetServer>,
request: Res<RenderRequest>,
mut _materials: ResMut<Assets<StandardMaterial>>,
) {
let fov = request.config.fov_radians();
commands.spawn((
Camera3d::default(),
Camera {
hdr: true,
..default()
},
Projection::Perspective(PerspectiveProjection {
fov,
near: request.config.near_plane,
far: request.config.far_plane,
..default()
}),
Msaa::Off,
request.camera_transform,
Tonemapping::None, DepthPrepass,
NormalPrepass,
RenderCamera,
));
let lighting = &request.config.lighting;
commands.insert_resource(AmbientLight {
color: Color::WHITE,
brightness: lighting.ambient_brightness,
});
if lighting.key_light_intensity > 0.0 {
commands.spawn((
PointLight {
intensity: lighting.key_light_intensity,
shadows_enabled: lighting.shadows_enabled,
..default()
},
Transform::from_xyz(
lighting.key_light_position[0],
lighting.key_light_position[1],
lighting.key_light_position[2],
),
));
}
if lighting.fill_light_intensity > 0.0 {
commands.spawn((
PointLight {
intensity: lighting.fill_light_intensity,
shadows_enabled: lighting.shadows_enabled,
..default()
},
Transform::from_xyz(
lighting.fill_light_position[0],
lighting.fill_light_position[1],
lighting.fill_light_position[2],
),
));
}
let scene_handle: Handle<Scene> = asset_server.load(&request.mesh_path);
commands.insert_resource(LoadedScene(scene_handle.clone()));
let texture_handle: Handle<Image> = asset_server.load(&request.texture_path);
commands.insert_resource(LoadedTexture(texture_handle.clone()));
let _material = _materials.add(StandardMaterial {
base_color_texture: Some(texture_handle),
unlit: true,
..default()
});
commands.spawn((
SceneRoot(scene_handle),
Transform::from_rotation(request.object_rotation.to_quat()),
RenderedObject,
));
println!("Scene setup complete");
}
fn check_assets_loaded(
mut state: ResMut<RenderState>,
asset_server: Res<AssetServer>,
scene: Option<Res<LoadedScene>>,
texture: Option<Res<LoadedTexture>>,
) {
let trace = render_trace_enabled();
let was_scene_loaded = state.scene_loaded;
let was_texture_loaded = state.texture_loaded;
state.frame_count += 1;
if state.scene_loaded && state.texture_loaded {
return;
}
if let Some(scene) = scene {
match asset_server.get_load_state(&scene.0) {
Some(LoadState::Loaded) => {
state.scene_loaded = true;
}
Some(LoadState::Failed(_)) => {}
_ => {}
}
}
if let Some(texture) = texture {
match asset_server.get_load_state(&texture.0) {
Some(LoadState::Loaded) => {
state.texture_loaded = true;
}
Some(LoadState::Failed(_)) => {}
_ => {}
}
}
if trace {
if !was_scene_loaded && state.scene_loaded {
eprintln!(
"[render_trace][coldinit] scene_loaded frame_count={}",
state.frame_count
);
}
if !was_texture_loaded && state.texture_loaded {
eprintln!(
"[render_trace][coldinit] texture_loaded frame_count={}",
state.frame_count
);
}
}
}
fn apply_materials(
mut state: ResMut<RenderState>,
texture: Option<Res<LoadedTexture>>,
mut materials: ResMut<Assets<StandardMaterial>>,
mut mesh_query: Query<&mut MeshMaterial3d<StandardMaterial>, With<Mesh3d>>,
) {
if !state.scene_loaded || !state.texture_loaded || state.capture_ready {
return;
}
state.frame_count += 1;
let Some(tex) = texture else { return };
if !state.materials_applied {
if mesh_query.is_empty() {
return;
}
let textured_material = materials.add(StandardMaterial {
base_color_texture: Some(tex.0.clone()),
unlit: true,
..default()
});
for mut mat in mesh_query.iter_mut() {
mat.0 = textured_material.clone();
}
state.materials_applied = true;
state.materials_applied_frame = state.frame_count;
}
if state.frame_count >= state.materials_applied_frame + 2 {
let was_ready = state.capture_ready;
state.capture_ready = true;
if render_trace_enabled() && !was_ready {
eprintln!(
"[render_trace][coldinit] capture_ready frame_count={}",
state.frame_count
);
}
}
}
#[allow(dead_code)]
fn request_screenshot(
mut commands: Commands,
mut state: ResMut<RenderState>,
shared_image: Res<SharedImageBuffer>,
mut depth_request: ResMut<DepthCaptureRequest>,
) {
if !state.capture_ready || state.screenshot_requested {
return;
}
let image_buffer = shared_image.0.clone();
depth_request.requested = true;
println!("Depth capture requested");
println!("Requesting screenshot via Screenshot entity");
commands.spawn(Screenshot::primary_window()).observe(
move |trigger: Trigger<ScreenshotCaptured>| {
let image: &Image = trigger.event();
let width = image.texture_descriptor.size.width;
let height = image.texture_descriptor.size.height;
let rgba_data = image.data.clone();
if let Ok(mut guard) = image_buffer.lock() {
*guard = Some((rgba_data, width, height));
}
},
);
state.screenshot_requested = true;
println!("Screenshot requested");
}
#[allow(dead_code)]
fn check_screenshot_ready(
mut state: ResMut<RenderState>,
shared_image: Res<SharedImageBuffer>,
shared_depth: Res<SharedDepthBuffer>,
request: Res<RenderRequest>,
) {
if !state.screenshot_requested || state.captured {
return;
}
state.frame_count += 1;
let rgba_ready = if let Ok(guard) = shared_image.0.lock() {
if let Some((rgba_data, width, height)) = guard.as_ref() {
if state.rgba_data.is_none() {
state.rgba_data = Some(rgba_data.clone());
state.image_width = *width;
state.image_height = *height;
}
true
} else {
false
}
} else {
false
};
let depth_ready = if let Ok(guard) = shared_depth.0.lock() {
if let Some((depth_data, _width, _height)) = guard.as_ref() {
if state.depth_data.is_none() {
state.depth_data = Some(depth_data.clone());
}
true
} else {
false
}
} else {
false
};
if rgba_ready && !depth_ready && state.frame_count > 60 {
let camera_dist = request.camera_transform.translation.length() as f64;
let pixel_count = (state.image_width * state.image_height) as usize;
state.depth_data = Some(vec![camera_dist; pixel_count]);
}
if state.rgba_data.is_some() && state.depth_data.is_some() {
state.captured = true;
}
}
#[allow(dead_code)]
fn extract_and_exit(
mut state: ResMut<RenderState>,
request: Res<RenderRequest>,
shared_output: Res<SharedOutput>,
mut commands: Commands,
windows: Query<Entity, With<bevy::window::Window>>,
) {
if state.exit_requested {
state.exit_frame_count += 1;
return;
}
if !state.captured {
return;
}
if let (Some(rgba), Some(depth)) = (&state.rgba_data, &state.depth_data) {
let width = state.image_width;
let height = state.image_height;
let intrinsics = request.config.intrinsics_for_size(width, height);
let output = RenderOutput {
rgba: rgba.clone(),
depth: depth.clone(),
width,
height,
intrinsics,
camera_transform: request.camera_transform,
object_rotation: request.object_rotation.clone(),
target_point: Vec3::ZERO,
targeting_policy: TargetingPolicy::Origin,
};
if let Ok(mut guard) = shared_output.0.lock() {
*guard = Some(output);
drop(guard);
std::thread::sleep(std::time::Duration::from_millis(200));
}
for window_entity in windows.iter() {
commands.entity(window_entity).despawn();
}
state.exit_requested = true;
}
}
fn setup_headless_scene(
mut commands: Commands,
mut images: ResMut<Assets<Image>>,
asset_server: Res<AssetServer>,
request: Res<RenderRequest>,
mut _materials: ResMut<Assets<StandardMaterial>>,
) {
let trace = render_trace_enabled();
let t0 = trace.then(std::time::Instant::now);
#[cfg(test)]
HEADLESS_SCENE_SETUP_COUNT.fetch_add(1, Ordering::SeqCst);
let width = request.config.width;
let height = request.config.height;
let size = Extent3d {
width,
height,
depth_or_array_layers: 1,
};
let mut render_target_image = Image::new_fill(
size,
TextureDimension::D2,
&[0, 0, 0, 255], TextureFormat::Rgba8UnormSrgb,
RenderAssetUsages::default(),
);
render_target_image.texture_descriptor.usage =
TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_SRC | TextureUsages::RENDER_ATTACHMENT;
let render_target_handle = images.add(render_target_image);
commands.insert_resource(RenderTargetImage(render_target_handle.clone()));
let fov = request.config.fov_radians();
commands.spawn((
Camera3d::default(),
Camera {
hdr: true,
target: RenderTarget::Image(render_target_handle.clone()),
..default()
},
Projection::Perspective(PerspectiveProjection {
fov,
near: request.config.near_plane,
far: request.config.far_plane,
..default()
}),
Msaa::Off,
request.camera_transform,
Tonemapping::None,
DepthPrepass,
NormalPrepass,
RenderCamera,
ImageCopier {
src_image: render_target_handle,
enabled: false, },
));
let lighting = &request.config.lighting;
commands.insert_resource(AmbientLight {
color: Color::WHITE,
brightness: lighting.ambient_brightness,
});
if lighting.key_light_intensity > 0.0 {
commands.spawn((
PointLight {
intensity: lighting.key_light_intensity,
shadows_enabled: lighting.shadows_enabled,
..default()
},
Transform::from_xyz(
lighting.key_light_position[0],
lighting.key_light_position[1],
lighting.key_light_position[2],
),
));
}
if lighting.fill_light_intensity > 0.0 {
commands.spawn((
PointLight {
intensity: lighting.fill_light_intensity,
shadows_enabled: lighting.shadows_enabled,
..default()
},
Transform::from_xyz(
lighting.fill_light_position[0],
lighting.fill_light_position[1],
lighting.fill_light_position[2],
),
));
}
let scene_handle: Handle<Scene> = asset_server.load(&request.mesh_path);
commands.insert_resource(LoadedScene(scene_handle.clone()));
let texture_handle: Handle<Image> = asset_server.load(&request.texture_path);
commands.insert_resource(LoadedTexture(texture_handle.clone()));
let _material = _materials.add(StandardMaterial {
base_color_texture: Some(texture_handle),
unlit: true,
..default()
});
commands.spawn((
SceneRoot(scene_handle),
Transform::from_rotation(request.object_rotation.to_quat()),
RenderedObject,
));
if let Some(t0) = t0 {
eprintln!(
"[render_trace][startup] setup_headless_scene ms={:.3}",
t0.elapsed().as_secs_f64() * 1000.0
);
}
}
fn request_headless_capture(
mut state: ResMut<RenderState>,
mut depth_request: ResMut<DepthCaptureRequest>,
mut query: Query<&mut ImageCopier>,
batch: Option<Res<HeadlessBatchSequence>>,
) {
let trace = render_trace_enabled();
let t0 = trace.then(std::time::Instant::now);
if !state.capture_ready || state.screenshot_requested {
if let Some(t0) = t0 {
eprintln!(
"[render_trace][sys] request_headless_capture skipped(gate) ms={:.3}",
t0.elapsed().as_secs_f64() * 1000.0
);
}
return;
}
if batch
.as_ref()
.is_some_and(|batch| batch.warmup_frames_remaining > 0)
{
if let Some(t0) = t0 {
eprintln!(
"[render_trace][sys] request_headless_capture skipped(warmup) ms={:.3}",
t0.elapsed().as_secs_f64() * 1000.0
);
}
return;
}
for mut copier in query.iter_mut() {
copier.enabled = true;
}
depth_request.requested = true;
state.screenshot_requested = true;
if let Some(t0) = t0 {
eprintln!(
"[render_trace][sys] request_headless_capture requested ms={:.3}",
t0.elapsed().as_secs_f64() * 1000.0
);
}
}
fn check_headless_capture_ready(
mut state: ResMut<RenderState>,
shared_rgba: Res<SharedRgbaBuffer>,
shared_depth: Res<SharedDepthBuffer>,
request: Res<RenderRequest>,
mut query: Query<&mut ImageCopier>,
) {
let trace = render_trace_enabled();
let t0 = trace.then(std::time::Instant::now);
if !state.screenshot_requested || state.captured {
if let Some(t0) = t0 {
eprintln!(
"[render_trace][sys] check_headless_capture_ready skipped(gate) ms={:.3}",
t0.elapsed().as_secs_f64() * 1000.0
);
}
return;
}
state.frame_count += 1;
let rgba_ready = if let Ok(guard) = shared_rgba.0.lock() {
if let Some((rgba_data, width, height)) = guard.as_ref() {
if state.rgba_data.is_none() {
state.rgba_data = Some(rgba_data.clone());
state.image_width = *width;
state.image_height = *height;
for mut copier in query.iter_mut() {
copier.enabled = false;
}
}
true
} else {
false
}
} else {
false
};
let depth_ready = if let Ok(guard) = shared_depth.0.lock() {
if let Some((depth_data, _width, _height)) = guard.as_ref() {
if state.depth_data.is_none() {
state.depth_data = Some(depth_data.clone());
}
true
} else {
false
}
} else {
false
};
if rgba_ready && !depth_ready && state.frame_count > 70 {
let camera_dist = request.camera_transform.translation.length() as f64;
let pixel_count = (state.image_width * state.image_height) as usize;
state.depth_data = Some(vec![camera_dist; pixel_count]);
}
if state.rgba_data.is_some() && state.depth_data.is_some() {
state.captured = true;
}
if let Some(t0) = t0 {
eprintln!(
"[render_trace][sys] check_headless_capture_ready rgba_ready={} depth_ready={} captured={} frame_count={} ms={:.3}",
rgba_ready,
depth_ready,
state.captured,
state.frame_count,
t0.elapsed().as_secs_f64() * 1000.0
);
}
}
fn extract_and_exit_headless(
mut state: ResMut<RenderState>,
request: Res<RenderRequest>,
shared_output: Res<SharedOutput>,
mut app_exit: EventWriter<bevy::app::AppExit>,
batch: Option<Res<HeadlessBatchSequence>>,
) {
if batch.is_some() {
return;
}
if state.exit_requested {
return;
}
if !state.captured {
return;
}
if let (Some(rgba), Some(depth)) = (&state.rgba_data, &state.depth_data) {
let width = state.image_width;
let height = state.image_height;
let intrinsics = request.config.intrinsics_for_size(width, height);
let output = RenderOutput {
rgba: rgba.clone(),
depth: depth.clone(),
width,
height,
intrinsics,
camera_transform: request.camera_transform,
object_rotation: request.object_rotation.clone(),
target_point: Vec3::ZERO,
targeting_policy: TargetingPolicy::Origin,
};
if let Ok(mut guard) = shared_output.0.lock() {
*guard = Some(output);
drop(guard);
std::thread::sleep(std::time::Duration::from_millis(200));
}
app_exit.send(bevy::app::AppExit::Success);
state.exit_requested = true;
}
}
fn tick_headless_batch_warmup(batch: Option<ResMut<HeadlessBatchSequence>>) {
let Some(mut batch) = batch else {
return;
};
if batch.warmup_frames_remaining > 0 {
batch.warmup_frames_remaining -= 1;
}
}
fn extract_and_continue_headless_batch(
mut state: ResMut<RenderState>,
request: Res<RenderRequest>,
buffers: (Res<SharedRgbaBuffer>, Res<SharedDepthBuffer>),
batch: Option<ResMut<HeadlessBatchSequence>>,
mut camera_query: Query<&mut Transform, With<RenderCamera>>,
mut depth_request: ResMut<DepthCaptureRequest>,
mut image_copiers: Query<&mut ImageCopier>,
) {
let trace = render_trace_enabled();
let t0 = trace.then(std::time::Instant::now);
let (shared_rgba, shared_depth) = buffers;
let Some(mut batch) = batch else {
if let Some(t0) = t0 {
eprintln!(
"[render_trace][sys] extract_and_continue_headless_batch skipped(no_batch) ms={:.3}",
t0.elapsed().as_secs_f64() * 1000.0
);
}
return;
};
if state.exit_requested || !state.captured || batch.done {
if let Some(t0) = t0 {
eprintln!(
"[render_trace][sys] extract_and_continue_headless_batch skipped(gate) captured={} done={} ms={:.3}",
state.captured,
batch.done,
t0.elapsed().as_secs_f64() * 1000.0
);
}
return;
}
if let (Some(rgba), Some(depth)) = (&state.rgba_data, &state.depth_data) {
let width = state.image_width;
let height = state.image_height;
let intrinsics = request.config.intrinsics_for_size(width, height);
let output = RenderOutput {
rgba: rgba.clone(),
depth: depth.clone(),
width,
height,
intrinsics,
camera_transform: batch
.current_viewpoint()
.unwrap_or(request.camera_transform),
object_rotation: request.object_rotation.clone(),
target_point: Vec3::ZERO,
targeting_policy: TargetingPolicy::Origin,
};
batch.outputs.push(output);
let next_index = batch.current_index + 1;
if next_index >= batch.viewpoints.len() {
batch.done = true;
state.exit_requested = true;
return;
}
batch.current_index = next_index;
batch.warmup_frames_remaining = BATCH_WARMUP_FRAMES;
if let Some(next_viewpoint) = batch.current_viewpoint() {
for mut camera_transform in camera_query.iter_mut() {
*camera_transform = next_viewpoint;
}
}
if let Ok(mut guard) = shared_rgba.0.lock() {
*guard = None;
}
if let Ok(mut guard) = shared_depth.0.lock() {
*guard = None;
}
for mut copier in image_copiers.iter_mut() {
copier.enabled = false;
}
depth_request.requested = false;
state.frame_count = 0;
state.capture_ready = true;
state.screenshot_requested = false;
state.captured = false;
state.rgba_data = None;
state.depth_data = None;
state.image_width = 0;
state.image_height = 0;
if let Some(t0) = t0 {
eprintln!(
"[render_trace][sys] extract_and_continue_headless_batch extracted vp={} next={} done={} ms={:.3}",
batch.current_index.saturating_sub(1),
batch.current_index,
batch.done,
t0.elapsed().as_secs_f64() * 1000.0
);
}
} else if let Some(t0) = t0 {
eprintln!(
"[render_trace][sys] extract_and_continue_headless_batch no_data ms={:.3}",
t0.elapsed().as_secs_f64() * 1000.0
);
}
}
#[derive(Component)]
struct SessionScene;
fn setup_session_persistent_scene(
mut commands: Commands,
mut images: ResMut<Assets<Image>>,
config: Res<SessionRenderConfig>,
) {
let width = config.0.width;
let height = config.0.height;
let size = Extent3d {
width,
height,
depth_or_array_layers: 1,
};
let mut render_target_image = Image::new_fill(
size,
TextureDimension::D2,
&[0, 0, 0, 255],
TextureFormat::Rgba8UnormSrgb,
RenderAssetUsages::default(),
);
render_target_image.texture_descriptor.usage =
TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_SRC | TextureUsages::RENDER_ATTACHMENT;
let render_target_handle = images.add(render_target_image);
commands.insert_resource(RenderTargetImage(render_target_handle.clone()));
let fov = config.0.fov_radians();
commands.spawn((
Camera3d::default(),
Camera {
hdr: true,
target: RenderTarget::Image(render_target_handle.clone()),
..default()
},
Projection::Perspective(PerspectiveProjection {
fov,
near: config.0.near_plane,
far: config.0.far_plane,
..default()
}),
Msaa::Off,
Transform::default(),
Tonemapping::None,
DepthPrepass,
NormalPrepass,
RenderCamera,
ImageCopier {
src_image: render_target_handle,
enabled: false,
},
));
let lighting = &config.0.lighting;
commands.insert_resource(AmbientLight {
color: Color::WHITE,
brightness: lighting.ambient_brightness,
});
if lighting.key_light_intensity > 0.0 {
commands.spawn((
PointLight {
intensity: lighting.key_light_intensity,
shadows_enabled: lighting.shadows_enabled,
..default()
},
Transform::from_xyz(
lighting.key_light_position[0],
lighting.key_light_position[1],
lighting.key_light_position[2],
),
));
}
if lighting.fill_light_intensity > 0.0 {
commands.spawn((
PointLight {
intensity: lighting.fill_light_intensity,
shadows_enabled: lighting.shadows_enabled,
..default()
},
Transform::from_xyz(
lighting.fill_light_position[0],
lighting.fill_light_position[1],
lighting.fill_light_position[2],
),
));
}
}
#[derive(Resource)]
struct SessionRenderConfig(RenderConfig);
pub struct RenderSession {
app: App,
render_config: RenderConfig,
shared_rgba: SharedRgbaBuffer,
shared_depth: SharedDepthBuffer,
_not_send_sync: std::marker::PhantomData<*const ()>,
}
impl RenderSession {
pub fn new(render_config: &crate::RenderConfig) -> Result<Self, crate::RenderError> {
let shared_rgba: SharedRgbaBuffer = SharedRgbaBuffer::default();
let shared_depth: SharedDepthBuffer = SharedDepthBuffer::default();
let mut app = App::new();
app.add_plugins(
DefaultPlugins
.set(WindowPlugin {
primary_window: None,
exit_condition: ExitCondition::DontExit,
..default()
})
.disable::<bevy::winit::WinitPlugin>()
.disable::<LogPlugin>()
.disable::<TerminalCtrlCHandlerPlugin>(),
)
.add_plugins(ObjPlugin)
.add_plugins(ImageCopyPlugin {
shared_rgba: shared_rgba.clone(),
})
.add_plugins(DepthReadbackPlugin {
shared_depth: shared_depth.clone(),
near: render_config.near_plane,
far: render_config.far_plane,
})
.insert_resource(SessionRenderConfig(render_config.clone()))
.insert_resource(shared_rgba.clone())
.init_resource::<RenderState>()
.add_systems(Startup, setup_session_persistent_scene)
.add_systems(
Update,
(
check_assets_loaded,
apply_materials,
tick_headless_batch_warmup,
request_headless_capture,
check_headless_capture_ready,
extract_and_continue_headless_batch,
)
.chain()
.run_if(bevy::ecs::schedule::common_conditions::resource_exists::<RenderRequest>),
);
app.finish();
app.cleanup();
app.update();
Ok(Self {
app,
render_config: render_config.clone(),
shared_rgba,
shared_depth,
_not_send_sync: std::marker::PhantomData,
})
}
pub fn render(
&mut self,
requests: &[crate::BatchRenderRequest],
) -> Result<Vec<crate::BatchRenderOutput>, crate::BatchRenderError> {
use crate::{BatchRenderError, BatchRenderOutput};
if requests.is_empty() {
return Ok(Vec::new());
}
let first = &requests[0];
if first.render_config != self.render_config {
return Err(BatchRenderError::InvalidConfig(
"RenderSession render_config mismatch: session was constructed with a different \
RenderConfig than the first request carries. Session config cannot change after \
`new()`; construct a new session if you need a different resolution/camera."
.to_string(),
));
}
for r in &requests[1..] {
if r.object_dir != first.object_dir
|| r.object_rotation != first.object_rotation
|| r.render_config != first.render_config
{
return Err(BatchRenderError::InvalidConfig(
"Phase 1 RenderSession::render requires homogeneous requests \
(same object_dir, object_rotation, and render_config across the batch). \
Call render() once per group instead."
.to_string(),
));
}
}
let object_dir = std::fs::canonicalize(&first.object_dir).map_err(|e| {
BatchRenderError::InvalidConfig(format!(
"Cannot canonicalize object directory {}: {}",
first.object_dir.display(),
e
))
})?;
let mesh_path = object_dir.join(GOOGLE_16K_MESH_RELATIVE);
let texture_path = object_dir.join(GOOGLE_16K_TEXTURE_RELATIVE);
if !mesh_path.exists() {
return Err(BatchRenderError::InvalidConfig(format!(
"Mesh not found: {}",
mesh_path.display()
)));
}
if !texture_path.exists() {
return Err(BatchRenderError::InvalidConfig(format!(
"Texture not found: {}",
texture_path.display()
)));
}
let viewpoints: Vec<Transform> = requests.iter().map(|r| r.viewpoint).collect();
{
let world = self.app.world_mut();
let stale: Vec<Entity> = world
.query_filtered::<Entity, With<SessionScene>>()
.iter(world)
.collect();
for entity in stale {
world.entity_mut(entity).despawn_recursive();
}
if let Ok(mut guard) = self.shared_rgba.0.lock() {
*guard = None;
}
if let Ok(mut guard) = self.shared_depth.0.lock() {
*guard = None;
}
*world.resource_mut::<RenderState>() = RenderState::default();
let new_request = RenderRequest {
mesh_path: mesh_path.display().to_string(),
texture_path: texture_path.display().to_string(),
camera_transform: viewpoints[0],
object_rotation: first.object_rotation.clone(),
config: self.render_config.clone(),
};
world.insert_resource(new_request);
let asset_server = world.resource::<AssetServer>().clone();
let scene_handle: Handle<Scene> = asset_server.load(mesh_path.display().to_string());
let texture_handle: Handle<Image> =
asset_server.load(texture_path.display().to_string());
world.insert_resource(LoadedScene(scene_handle.clone()));
world.insert_resource(LoadedTexture(texture_handle));
world.spawn((
SceneRoot(scene_handle),
Transform::from_rotation(first.object_rotation.to_quat()),
RenderedObject,
SessionScene,
));
let camera_entity = world
.query_filtered::<Entity, With<RenderCamera>>()
.iter(world)
.next();
if let Some(cam) = camera_entity {
if let Some(mut transform) = world.entity_mut(cam).get_mut::<Transform>() {
*transform = viewpoints[0];
}
}
world.insert_resource(HeadlessBatchSequence::new(viewpoints.clone()));
}
let timeout = std::time::Duration::from_secs(RENDER_TIMEOUT_SECS);
let start = std::time::Instant::now();
loop {
if start.elapsed() > timeout {
return Err(BatchRenderError::TotalFailure(format!(
"RenderSession::render timed out after {}s",
RENDER_TIMEOUT_SECS
)));
}
self.app.update();
if self.app.world().resource::<HeadlessBatchSequence>().done {
break;
}
}
let mut sequence = self.app.world_mut().resource_mut::<HeadlessBatchSequence>();
if sequence.outputs.len() != requests.len() {
return Err(BatchRenderError::TotalFailure(format!(
"RenderSession produced {} outputs for {} requests",
sequence.outputs.len(),
requests.len()
)));
}
let outputs = std::mem::take(&mut sequence.outputs);
Ok(requests
.iter()
.cloned()
.zip(outputs)
.map(|(req, out)| BatchRenderOutput::from_render_output(req, out))
.collect())
}
}
#[derive(Component)]
struct PersistentScene;
pub struct PersistentRenderer {
app: App,
object_dir: PathBuf,
render_config: RenderConfig,
shared_rgba: SharedRgbaBuffer,
shared_depth: SharedDepthBuffer,
_not_send_sync: std::marker::PhantomData<*const ()>,
}
impl PersistentRenderer {
pub fn new(
object_dir: &Path,
render_config: &RenderConfig,
) -> Result<Self, crate::RenderError> {
let object_dir =
std::fs::canonicalize(object_dir).map_err(|e| crate::RenderError::FileNotFound {
path: object_dir.display().to_string(),
reason: e.to_string(),
})?;
let mesh_path = object_dir.join(GOOGLE_16K_MESH_RELATIVE);
let texture_path = object_dir.join(GOOGLE_16K_TEXTURE_RELATIVE);
if !mesh_path.exists() {
return Err(crate::RenderError::MeshNotFound(
mesh_path.display().to_string(),
));
}
if !texture_path.exists() {
return Err(crate::RenderError::TextureNotFound(
texture_path.display().to_string(),
));
}
let shared_rgba: SharedRgbaBuffer = SharedRgbaBuffer::default();
let shared_depth: SharedDepthBuffer = SharedDepthBuffer::default();
let mut app = App::new();
app.add_plugins(
DefaultPlugins
.set(WindowPlugin {
primary_window: None,
exit_condition: ExitCondition::DontExit,
..default()
})
.disable::<bevy::winit::WinitPlugin>()
.disable::<LogPlugin>()
.disable::<TerminalCtrlCHandlerPlugin>(),
)
.add_plugins(ObjPlugin)
.add_plugins(ImageCopyPlugin {
shared_rgba: shared_rgba.clone(),
})
.add_plugins(DepthReadbackPlugin {
shared_depth: shared_depth.clone(),
near: render_config.near_plane,
far: render_config.far_plane,
})
.insert_resource(SessionRenderConfig(render_config.clone()))
.insert_resource(shared_rgba.clone())
.init_resource::<RenderState>()
.add_systems(Startup, setup_session_persistent_scene)
.add_systems(
Update,
(
check_assets_loaded,
apply_materials,
tick_headless_batch_warmup,
request_headless_capture,
check_headless_capture_ready,
extract_and_continue_headless_batch,
)
.chain()
.run_if(bevy::ecs::schedule::common_conditions::resource_exists::<RenderRequest>),
);
app.finish();
app.cleanup();
app.update();
let initial_request = RenderRequest {
mesh_path: mesh_path.display().to_string(),
texture_path: texture_path.display().to_string(),
camera_transform: Transform::default(),
object_rotation: ObjectRotation::identity(),
config: render_config.clone(),
};
{
let world = app.world_mut();
let asset_server = world.resource::<AssetServer>().clone();
let scene_handle: Handle<Scene> = asset_server.load(mesh_path.display().to_string());
let texture_handle: Handle<Image> =
asset_server.load(texture_path.display().to_string());
world.insert_resource(LoadedScene(scene_handle.clone()));
world.insert_resource(LoadedTexture(texture_handle));
world.insert_resource(initial_request);
world.spawn((
SceneRoot(scene_handle),
Transform::from_rotation(ObjectRotation::identity().to_quat()),
RenderedObject,
PersistentScene,
));
world.insert_resource(HeadlessBatchSequence::new(vec![Transform::default()]));
}
let timeout = std::time::Duration::from_secs(RENDER_TIMEOUT_SECS);
let start = std::time::Instant::now();
loop {
if start.elapsed() > timeout {
return Err(crate::RenderError::RenderFailed(format!(
"PersistentRenderer::new warmup render timed out after {RENDER_TIMEOUT_SECS}s"
)));
}
app.update();
if app.world().resource::<HeadlessBatchSequence>().done {
break;
}
}
app.world_mut()
.resource_mut::<HeadlessBatchSequence>()
.outputs
.clear();
Ok(Self {
app,
object_dir,
render_config: render_config.clone(),
shared_rgba,
shared_depth,
_not_send_sync: std::marker::PhantomData,
})
}
pub fn render(
&mut self,
camera_transform: &Transform,
object_rotation: &ObjectRotation,
) -> Result<RenderOutput, crate::RenderError> {
let camera_transform = *camera_transform;
let object_rotation_owned = object_rotation.clone();
{
let world = self.app.world_mut();
let scene_entity = world
.query_filtered::<Entity, With<PersistentScene>>()
.iter(world)
.next();
if let Some(entity) = scene_entity {
if let Some(mut transform) = world.entity_mut(entity).get_mut::<Transform>() {
*transform = Transform::from_rotation(object_rotation_owned.to_quat());
}
}
let cam_entity = world
.query_filtered::<Entity, With<RenderCamera>>()
.iter(world)
.next();
if let Some(cam) = cam_entity {
if let Some(mut transform) = world.entity_mut(cam).get_mut::<Transform>() {
*transform = camera_transform;
}
}
{
let mut state = world.resource_mut::<RenderState>();
state.exit_requested = false;
state.screenshot_requested = false;
state.captured = false;
state.rgba_data = None;
state.depth_data = None;
state.frame_count = 0;
state.image_width = 0;
state.image_height = 0;
state.capture_ready = true;
}
if let Ok(mut guard) = self.shared_rgba.0.lock() {
*guard = None;
}
if let Ok(mut guard) = self.shared_depth.0.lock() {
*guard = None;
}
{
let mut req = world.resource_mut::<RenderRequest>();
req.camera_transform = camera_transform;
req.object_rotation = object_rotation_owned.clone();
}
let mut batch = HeadlessBatchSequence::new(vec![camera_transform]);
batch.warmup_frames_remaining = PERSISTENT_WARMUP_FRAMES;
world.insert_resource(batch);
}
let timeout = std::time::Duration::from_secs(RENDER_TIMEOUT_SECS);
let start = std::time::Instant::now();
loop {
if start.elapsed() > timeout {
return Err(crate::RenderError::RenderFailed(format!(
"PersistentRenderer::render timed out after {RENDER_TIMEOUT_SECS}s"
)));
}
self.app.update();
if self.app.world().resource::<HeadlessBatchSequence>().done {
break;
}
}
let mut sequence = self.app.world_mut().resource_mut::<HeadlessBatchSequence>();
let mut outputs = std::mem::take(&mut sequence.outputs);
if outputs.len() != 1 {
return Err(crate::RenderError::RenderFailed(format!(
"PersistentRenderer::render expected 1 output, got {}",
outputs.len()
)));
}
Ok(outputs.remove(0))
}
pub fn object_dir(&self) -> &Path {
&self.object_dir
}
pub fn render_config(&self) -> &RenderConfig {
&self.render_config
}
pub fn close(self) {
}
}
pub fn render_to_files(
object_dir: &Path,
camera_transform: &Transform,
object_rotation: &ObjectRotation,
config: &RenderConfig,
rgba_path: &Path,
depth_path: &Path,
) -> Result<(), RenderError> {
let mesh_path = object_dir.join(GOOGLE_16K_MESH_RELATIVE);
let texture_path = object_dir.join(GOOGLE_16K_TEXTURE_RELATIVE);
if !mesh_path.exists() {
return Err(RenderError::MeshNotFound(mesh_path.display().to_string()));
}
if !texture_path.exists() {
return Err(RenderError::TextureNotFound(
texture_path.display().to_string(),
));
}
let request = RenderRequest {
mesh_path: mesh_path.display().to_string(),
texture_path: texture_path.display().to_string(),
camera_transform: *camera_transform,
object_rotation: object_rotation.clone(),
config: config.clone(),
};
let shared_output: SharedOutput = SharedOutput(Arc::new(Mutex::new(None)));
let output_poll = shared_output.clone();
let rgba_path = rgba_path.to_path_buf();
let depth_path = depth_path.to_path_buf();
let shared_rgba: SharedRgbaBuffer = SharedRgbaBuffer::default();
let shared_depth: SharedDepthBuffer = SharedDepthBuffer::default();
std::thread::spawn(move || {
let timeout = std::time::Duration::from_secs(RENDER_TIMEOUT_SECS);
let start = std::time::Instant::now();
let poll_interval = std::time::Duration::from_millis(100);
loop {
if let Ok(guard) = output_poll.0.lock() {
if let Some(output) = guard.as_ref() {
if let Err(e) =
save_rgba_to_png(&output.rgba, output.width, output.height, &rgba_path)
{
eprintln!("Failed to save RGBA: {:?}", e);
std::process::exit(1);
}
if let Err(e) = save_depth_to_binary(&output.depth, &depth_path) {
eprintln!("Failed to save depth: {:?}", e);
std::process::exit(1);
}
std::process::exit(0);
}
}
if start.elapsed() > timeout {
eprintln!(
"Error: Render timeout after {} seconds",
RENDER_TIMEOUT_SECS
);
eprintln!("Debug info: This may indicate GPU issues, missing assets, or insufficient system resources.");
std::process::exit(1);
}
std::thread::sleep(poll_interval);
}
});
static BACKEND_INIT: OnceLock<()> = OnceLock::new();
BACKEND_INIT.get_or_init(|| {
let backend_config = BackendConfig::headless();
backend_config.apply_env();
});
build_headless_app(request, shared_output, shared_rgba, shared_depth).run();
Err(RenderError::RenderFailed(
"Render did not complete".to_string(),
))
}
fn save_rgba_to_png(rgba: &[u8], width: u32, height: u32, path: &Path) -> Result<(), String> {
use image::{ImageBuffer, Rgba};
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
}
let img: ImageBuffer<Rgba<u8>, Vec<u8>> =
ImageBuffer::from_raw(width, height, rgba.to_vec())
.ok_or_else(|| "Failed to create image buffer".to_string())?;
img.save(path).map_err(|e| e.to_string())
}
fn save_depth_to_binary(depth: &[f64], path: &Path) -> Result<(), String> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
}
let bytes: Vec<u8> = depth.iter().flat_map(|f| f.to_le_bytes()).collect();
std::fs::write(path, &bytes).map_err(|e| e.to_string())
}
#[cfg(test)]
mod smoke_tests {
use super::{headless_scene_setup_count, reset_headless_scene_setup_count};
use crate::{
BatchRenderConfig, BatchRenderRequest, ObjectRotation, RenderConfig, TargetingPolicy, Vec3,
ViewpointConfig,
};
use image::{ImageBuffer, Rgba};
use tempfile::TempDir;
fn write_synthetic_object() -> TempDir {
let temp_dir = TempDir::new().expect("create temp dir for synthetic object");
let object_dir = temp_dir.path().join("synthetic_cube").join("google_16k");
std::fs::create_dir_all(&object_dir).expect("create synthetic google_16k dir");
let obj = r#"o SyntheticCube
v -0.10 -0.10 0.10
v 0.10 -0.10 0.10
v 0.10 0.10 0.10
v -0.10 0.10 0.10
v -0.10 -0.10 -0.10
v 0.10 -0.10 -0.10
v 0.10 0.10 -0.10
v -0.10 0.10 -0.10
vt 0.0 0.0
vt 1.0 0.0
vt 1.0 1.0
vt 0.0 1.0
f 1/1 2/2 3/3
f 1/1 3/3 4/4
f 6/1 5/2 8/3
f 6/1 8/3 7/4
f 2/1 6/2 7/3
f 2/1 7/3 3/4
f 5/1 1/2 4/3
f 5/1 4/3 8/4
f 4/1 3/2 7/3
f 4/1 7/3 8/4
f 5/1 6/2 2/3
f 5/1 2/3 1/4
"#;
std::fs::write(object_dir.join("textured.obj"), obj).expect("write synthetic obj");
let texture = ImageBuffer::from_fn(2, 2, |x, y| match (x, y) {
(0, 0) => Rgba([255u8, 48, 48, 255]),
(1, 0) => Rgba([48u8, 255, 48, 255]),
(0, 1) => Rgba([48u8, 48, 255, 255]),
_ => Rgba([255u8, 255, 64, 255]),
});
texture
.save(object_dir.join("texture_map.png"))
.expect("write synthetic texture");
temp_dir
}
#[test]
#[ignore = "headless throughput smoke check is opt-in because it needs a local render backend"]
fn test_headless_batch_throughput_smoke() {
crate::initialize();
reset_headless_scene_setup_count();
let object_root = write_synthetic_object();
let object_dir = object_root.path().join("synthetic_cube");
let viewpoints = crate::generate_viewpoints(&ViewpointConfig::default());
let request_count = 5usize;
let config = RenderConfig::tbp_default();
let requests: Vec<_> = viewpoints
.iter()
.take(request_count)
.copied()
.map(|viewpoint| BatchRenderRequest {
object_dir: object_dir.clone(),
viewpoint,
object_rotation: ObjectRotation::identity(),
render_config: config.clone(),
target_point: Vec3::ZERO,
targeting_policy: TargetingPolicy::Origin,
})
.collect();
let start = std::time::Instant::now();
let outputs = crate::render_batch(requests, &BatchRenderConfig::default())
.expect("synthetic headless batch render should succeed");
let elapsed = start.elapsed();
assert_eq!(outputs.len(), request_count);
assert_eq!(
headless_scene_setup_count(),
1,
"homogeneous batch smoke check should reuse one headless app setup"
);
for (idx, output) in outputs.iter().enumerate() {
assert_eq!(output.width, config.width, "output {idx} width mismatch");
assert_eq!(output.height, config.height, "output {idx} height mismatch");
assert_eq!(
output.rgba.len(),
(config.width * config.height * 4) as usize,
"output {idx} rgba size mismatch"
);
assert_eq!(
output.depth.len(),
(config.width * config.height) as usize,
"output {idx} depth size mismatch"
);
assert!(
output
.rgba
.chunks_exact(4)
.any(|px| px[0] != 0 || px[1] != 0 || px[2] != 0),
"output {idx} should contain visible color"
);
}
assert!(
elapsed < std::time::Duration::from_secs(8),
"5 synthetic headless captures took {:.2}s, expected < 8.0s",
elapsed.as_secs_f64()
);
}
}