Skip to main content

bevy_sensor/
batch.rs

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