bevy_sensor/
batch.rs

1//! Batch rendering API for multiple viewpoints and objects.
2//!
3//! This module provides efficient batch rendering that eliminates subprocess spawning and
4//! Bevy app initialization overhead. A single Bevy app instance is kept alive and reused
5//! to render multiple viewpoints, achieving 10-100x speedup for typical batches.
6//!
7//! # Example
8//!
9//! ```ignore
10//! use bevy_sensor::{
11//!     create_batch_renderer, queue_render_request, render_next_in_batch,
12//!     batch::BatchRenderRequest, BatchRenderConfig, RenderConfig, ObjectRotation,
13//! };
14//! use std::path::PathBuf;
15//!
16//! // Create a persistent renderer (initializes once)
17//! let config = BatchRenderConfig::default();
18//! let mut renderer = create_batch_renderer(&config)?;
19//!
20//! // Queue multiple renders
21//! for rotation in rotations {
22//!     for viewpoint in viewpoints {
23//!         queue_render_request(&mut renderer, BatchRenderRequest {
24//!             object_dir: "/tmp/ycb/003_cracker_box".into(),
25//!             viewpoint,
26//!             object_rotation: rotation.clone(),
27//!             render_config: RenderConfig::tbp_default(),
28//!         })?;
29//!     }
30//! }
31//!
32//! // Execute and collect results
33//! let mut results = Vec::new();
34//! loop {
35//!     match render_next_in_batch(&mut renderer, 500)? {
36//!         Some(output) => results.push(output),
37//!         None => break,
38//!     }
39//! }
40//! ```
41
42use crate::{CameraIntrinsics, ObjectRotation, RenderConfig, RenderOutput};
43use bevy::prelude::Transform;
44use std::collections::VecDeque;
45use std::path::PathBuf;
46
47/// Configuration for batch rendering.
48#[derive(Clone, Debug)]
49pub struct BatchRenderConfig {
50    /// Maximum number of renders to queue before automatic cleanup
51    pub max_batch_size: usize,
52    /// Timeout in milliseconds per individual render
53    pub frame_timeout_ms: u32,
54    /// Enable depth buffer readback
55    pub enable_depth_readback: bool,
56    /// Enable asset caching for repeated objects
57    pub enable_asset_caching: bool,
58    /// Number of renders before triggering resource cleanup
59    pub resource_cleanup_interval: u32,
60}
61
62impl Default for BatchRenderConfig {
63    fn default() -> Self {
64        Self {
65            max_batch_size: 256,
66            frame_timeout_ms: 500,
67            enable_depth_readback: true,
68            enable_asset_caching: true,
69            resource_cleanup_interval: 32,
70        }
71    }
72}
73
74/// A single render request in a batch.
75#[derive(Clone, Debug)]
76pub struct BatchRenderRequest {
77    /// Path to YCB object directory (e.g., "/tmp/ycb/003_cracker_box")
78    pub object_dir: PathBuf,
79    /// Camera transform (position and orientation)
80    pub viewpoint: Transform,
81    /// Object rotation to apply
82    pub object_rotation: ObjectRotation,
83    /// Render configuration (resolution, lighting, etc.)
84    pub render_config: RenderConfig,
85}
86
87/// Status of a single render in a batch.
88#[derive(Clone, Debug, Copy, PartialEq, Eq)]
89pub enum RenderStatus {
90    /// Render completed successfully with RGBA and depth
91    Success,
92    /// Render completed but depth extraction failed
93    PartialFailure,
94    /// Render failed completely
95    Failed,
96}
97
98/// Output from a single render in a batch.
99#[derive(Clone, Debug)]
100pub struct BatchRenderOutput {
101    /// Original request for this render
102    pub request: BatchRenderRequest,
103    /// RGBA pixel data (width * height * 4 bytes, row-major)
104    pub rgba: Vec<u8>,
105    /// Depth data in meters (width * height f64s)
106    pub depth: Vec<f64>,
107    /// Image width in pixels
108    pub width: u32,
109    /// Image height in pixels
110    pub height: u32,
111    /// Camera intrinsics used
112    pub intrinsics: CameraIntrinsics,
113    /// Status of this render
114    pub status: RenderStatus,
115    /// Error message if status is Failed or PartialFailure
116    pub error_message: Option<String>,
117}
118
119impl BatchRenderOutput {
120    /// Convert to neocortx-compatible RGB format: Vec<Vec<[u8; 3]>>
121    pub fn to_rgb_image(&self) -> Vec<Vec<[u8; 3]>> {
122        let mut image = Vec::with_capacity(self.height as usize);
123        for y in 0..self.height {
124            let mut row = Vec::with_capacity(self.width as usize);
125            for x in 0..self.width {
126                let idx = ((y * self.width + x) * 4) as usize;
127                if idx + 2 < self.rgba.len() {
128                    row.push([self.rgba[idx], self.rgba[idx + 1], self.rgba[idx + 2]]);
129                } else {
130                    row.push([0, 0, 0]);
131                }
132            }
133            image.push(row);
134        }
135        image
136    }
137
138    /// Convert depth to neocortx-compatible format: Vec<Vec<f64>>
139    pub fn to_depth_image(&self) -> Vec<Vec<f64>> {
140        let mut image = Vec::with_capacity(self.height as usize);
141        for y in 0..self.height {
142            let mut row = Vec::with_capacity(self.width as usize);
143            for x in 0..self.width {
144                let idx = (y * self.width + x) as usize;
145                if idx < self.depth.len() {
146                    row.push(self.depth[idx]);
147                } else {
148                    row.push(0.0);
149                }
150            }
151            image.push(row);
152        }
153        image
154    }
155
156    /// Convert from RenderOutput, copying all fields
157    pub fn from_render_output(request: BatchRenderRequest, output: RenderOutput) -> Self {
158        Self {
159            request,
160            rgba: output.rgba,
161            depth: output.depth,
162            width: output.width,
163            height: output.height,
164            intrinsics: output.intrinsics,
165            status: RenderStatus::Success,
166            error_message: None,
167        }
168    }
169}
170
171/// Error types for batch rendering.
172#[derive(Debug, Clone)]
173pub enum BatchRenderError {
174    /// Some renders succeeded, others failed
175    PartialFailure { successful: usize, failed: usize },
176    /// All renders failed
177    TotalFailure(String),
178    /// Invalid configuration
179    InvalidConfig(String),
180    /// Queue is full
181    QueueFull,
182    /// No renders queued
183    EmptyQueue,
184}
185
186impl std::fmt::Display for BatchRenderError {
187    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
188        match self {
189            BatchRenderError::PartialFailure { successful, failed } => {
190                write!(
191                    f,
192                    "Batch render partial failure: {} succeeded, {} failed",
193                    successful, failed
194                )
195            }
196            BatchRenderError::TotalFailure(msg) => write!(f, "Batch render total failure: {}", msg),
197            BatchRenderError::InvalidConfig(msg) => write!(f, "Invalid batch config: {}", msg),
198            BatchRenderError::QueueFull => write!(f, "Batch queue is full"),
199            BatchRenderError::EmptyQueue => write!(f, "No renders queued"),
200        }
201    }
202}
203
204impl std::error::Error for BatchRenderError {}
205
206/// State machine for batch rendering lifecycle.
207#[derive(Clone, Copy, Debug, PartialEq, Eq)]
208pub enum BatchState {
209    /// Idle, waiting for requests to queue
210    Idle,
211    /// Loading object assets (mesh, texture)
212    LoadingAssets,
213    /// Rendering frame to GPU buffer
214    RenderingFrame,
215    /// Extracting RGBA and depth from GPU
216    ExtractingResults,
217    /// Cleaning up resources
218    Cleanup,
219    /// Shutting down
220    Shutdown,
221}
222
223/// Manages a persistent Bevy app for batch rendering.
224pub struct BatchRenderer {
225    /// Queued render requests
226    pub pending_requests: VecDeque<BatchRenderRequest>,
227    /// Completed results
228    pub completed_results: Vec<BatchRenderOutput>,
229    /// Current request being processed
230    pub current_request: Option<BatchRenderRequest>,
231    /// Current render output being built
232    pub current_output: Option<BatchRenderOutput>,
233    /// Frame counter for timeout management
234    pub frame_count: u32,
235    /// Current state
236    pub state: BatchState,
237    /// Configuration
238    pub config: BatchRenderConfig,
239    /// Total renders processed
240    pub renders_processed: usize,
241}
242
243impl BatchRenderer {
244    /// Create a new batch renderer with default configuration.
245    pub fn new(config: BatchRenderConfig) -> Self {
246        Self {
247            pending_requests: VecDeque::new(),
248            completed_results: Vec::new(),
249            current_request: None,
250            current_output: None,
251            frame_count: 0,
252            state: BatchState::Idle,
253            config,
254            renders_processed: 0,
255        }
256    }
257
258    /// Queue a render request for batch processing.
259    pub fn queue_request(&mut self, request: BatchRenderRequest) -> Result<(), BatchRenderError> {
260        if self.pending_requests.len() >= self.config.max_batch_size {
261            return Err(BatchRenderError::QueueFull);
262        }
263        self.pending_requests.push_back(request);
264        Ok(())
265    }
266
267    /// Get the number of pending requests.
268    pub fn pending_count(&self) -> usize {
269        self.pending_requests.len()
270    }
271
272    /// Get the number of completed results.
273    pub fn completed_count(&self) -> usize {
274        self.completed_results.len()
275    }
276
277    /// Get all completed results and clear the internal list.
278    pub fn take_completed(&mut self) -> Vec<BatchRenderOutput> {
279        std::mem::take(&mut self.completed_results)
280    }
281
282    /// Check if all work is done (no pending requests and not currently rendering).
283    pub fn is_finished(&self) -> bool {
284        self.pending_requests.is_empty() && self.current_request.is_none()
285    }
286}
287
288#[cfg(test)]
289mod tests {
290    use super::*;
291
292    #[test]
293    fn test_batch_config_defaults() {
294        let config = BatchRenderConfig::default();
295        assert_eq!(config.max_batch_size, 256);
296        assert_eq!(config.frame_timeout_ms, 500);
297        assert!(config.enable_depth_readback);
298        assert!(config.enable_asset_caching);
299    }
300
301    #[test]
302    fn test_batch_renderer_creation() {
303        let config = BatchRenderConfig::default();
304        let renderer = BatchRenderer::new(config);
305        assert_eq!(renderer.state, BatchState::Idle);
306        assert_eq!(renderer.pending_count(), 0);
307        assert_eq!(renderer.completed_count(), 0);
308        assert!(renderer.is_finished());
309    }
310
311    #[test]
312    fn test_queue_request() {
313        let mut renderer = BatchRenderer::new(BatchRenderConfig::default());
314        let request = BatchRenderRequest {
315            object_dir: "/tmp/test".into(),
316            viewpoint: Transform::default(),
317            object_rotation: ObjectRotation::identity(),
318            render_config: RenderConfig::tbp_default(),
319        };
320        assert!(renderer.queue_request(request).is_ok());
321        assert_eq!(renderer.pending_count(), 1);
322    }
323
324    #[test]
325    fn test_queue_full() {
326        let config = BatchRenderConfig {
327            max_batch_size: 1,
328            ..BatchRenderConfig::default()
329        };
330        let mut renderer = BatchRenderer::new(config);
331
332        let request = BatchRenderRequest {
333            object_dir: "/tmp/test".into(),
334            viewpoint: Transform::default(),
335            object_rotation: ObjectRotation::identity(),
336            render_config: RenderConfig::tbp_default(),
337        };
338
339        assert!(renderer.queue_request(request.clone()).is_ok());
340        assert!(matches!(
341            renderer.queue_request(request),
342            Err(BatchRenderError::QueueFull)
343        ));
344    }
345
346    #[test]
347    fn test_batch_render_output_rgb_conversion() {
348        let request = BatchRenderRequest {
349            object_dir: "/tmp/test".into(),
350            viewpoint: Transform::default(),
351            object_rotation: ObjectRotation::identity(),
352            render_config: RenderConfig::tbp_default(),
353        };
354
355        // Create minimal output: 2x2 image
356        let mut rgba = vec![0u8; 2 * 2 * 4];
357        // Pixel (0,0) = red
358        rgba[0] = 255;
359        rgba[1] = 0;
360        rgba[2] = 0;
361        rgba[3] = 255;
362
363        let output = BatchRenderOutput {
364            request,
365            rgba,
366            depth: vec![1.0; 4],
367            width: 2,
368            height: 2,
369            intrinsics: RenderConfig::tbp_default().intrinsics(),
370            status: RenderStatus::Success,
371            error_message: None,
372        };
373
374        let rgb = output.to_rgb_image();
375        assert_eq!(rgb.len(), 2); // 2 rows
376        assert_eq!(rgb[0].len(), 2); // 2 cols
377        assert_eq!(rgb[0][0], [255, 0, 0]); // Red
378    }
379}
bevy_sensor/batch.rs

bevy_sensor/
batch.rs
