Skip to main content

polyscope_structures/floating/
render_image.rs

1//! Render image floating quantities for depth-composited external renders.
2
3use glam::{Vec3, Vec4};
4use polyscope_core::quantity::{Quantity, QuantityKind};
5
6use super::ImageOrigin;
7
8/// A depth render image (geometry from an external renderer).
9///
10/// Stores per-pixel radial depth values and optional world-space normals.
11/// Used for depth-compositing external renders with the polyscope scene.
12pub struct FloatingDepthRenderImage {
13    name: String,
14    width: u32,
15    height: u32,
16    depths: Vec<f32>,           // Radial distance from camera
17    normals: Option<Vec<Vec3>>, // Optional world-space normals
18    origin: ImageOrigin,
19    enabled: bool,
20}
21
22impl FloatingDepthRenderImage {
23    /// Creates a new depth render image.
24    pub fn new(name: impl Into<String>, width: u32, height: u32, depths: Vec<f32>) -> Self {
25        Self {
26            name: name.into(),
27            width,
28            height,
29            depths,
30            normals: None,
31            origin: ImageOrigin::default(),
32            enabled: true,
33        }
34    }
35
36    /// Sets per-pixel world-space normals.
37    pub fn set_normals(&mut self, normals: Vec<Vec3>) -> &mut Self {
38        self.normals = Some(normals);
39        self
40    }
41
42    /// Returns the image width.
43    #[must_use]
44    pub fn width(&self) -> u32 {
45        self.width
46    }
47
48    /// Returns the image height.
49    #[must_use]
50    pub fn height(&self) -> u32 {
51        self.height
52    }
53
54    /// Returns the depth values.
55    #[must_use]
56    pub fn depths(&self) -> &[f32] {
57        &self.depths
58    }
59
60    /// Returns the normals, if set.
61    #[must_use]
62    pub fn normals(&self) -> Option<&[Vec3]> {
63        self.normals.as_deref()
64    }
65
66    /// Gets the image origin.
67    #[must_use]
68    pub fn origin(&self) -> ImageOrigin {
69        self.origin
70    }
71
72    /// Sets the image origin.
73    pub fn set_origin(&mut self, origin: ImageOrigin) -> &mut Self {
74        self.origin = origin;
75        self
76    }
77
78    /// Returns the depth at pixel (x, y), accounting for image origin.
79    #[must_use]
80    pub fn depth_at(&self, x: u32, y: u32) -> f32 {
81        let row = match self.origin {
82            ImageOrigin::UpperLeft => y,
83            ImageOrigin::LowerLeft => self.height - 1 - y,
84        };
85        self.depths[(row * self.width + x) as usize]
86    }
87
88    /// Returns whether a pixel has valid depth (not infinity/NaN).
89    #[must_use]
90    pub fn has_depth(&self, x: u32, y: u32) -> bool {
91        let d = self.depth_at(x, y);
92        d.is_finite() && d > 0.0
93    }
94}
95
96impl Quantity for FloatingDepthRenderImage {
97    fn as_any(&self) -> &dyn std::any::Any {
98        self
99    }
100    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
101        self
102    }
103    fn name(&self) -> &str {
104        &self.name
105    }
106    #[allow(clippy::unnecessary_literal_bound)]
107    fn structure_name(&self) -> &str {
108        "" // No parent structure
109    }
110    fn kind(&self) -> QuantityKind {
111        QuantityKind::Other
112    }
113    fn is_enabled(&self) -> bool {
114        self.enabled
115    }
116    fn set_enabled(&mut self, enabled: bool) {
117        self.enabled = enabled;
118    }
119    fn build_ui(&mut self, _ui: &dyn std::any::Any) {}
120    fn refresh(&mut self) {}
121    fn data_size(&self) -> usize {
122        self.depths.len()
123    }
124}
125
126/// A color render image (colored geometry from an external renderer).
127///
128/// Stores per-pixel depth, color, and optional normals for depth-compositing
129/// externally rendered content with polyscope's scene.
130pub struct FloatingColorRenderImage {
131    name: String,
132    width: u32,
133    height: u32,
134    depths: Vec<f32>,
135    colors: Vec<Vec4>, // Per-pixel RGBA
136    normals: Option<Vec<Vec3>>,
137    origin: ImageOrigin,
138    enabled: bool,
139}
140
141impl FloatingColorRenderImage {
142    /// Creates a new color render image.
143    pub fn new(
144        name: impl Into<String>,
145        width: u32,
146        height: u32,
147        depths: Vec<f32>,
148        colors: Vec<Vec3>,
149    ) -> Self {
150        let colors = colors.into_iter().map(|c| c.extend(1.0)).collect();
151        Self {
152            name: name.into(),
153            width,
154            height,
155            depths,
156            colors,
157            normals: None,
158            origin: ImageOrigin::default(),
159            enabled: true,
160        }
161    }
162
163    /// Sets per-pixel world-space normals.
164    pub fn set_normals(&mut self, normals: Vec<Vec3>) -> &mut Self {
165        self.normals = Some(normals);
166        self
167    }
168
169    /// Returns the image width.
170    #[must_use]
171    pub fn width(&self) -> u32 {
172        self.width
173    }
174
175    /// Returns the image height.
176    #[must_use]
177    pub fn height(&self) -> u32 {
178        self.height
179    }
180
181    /// Returns the depth values.
182    #[must_use]
183    pub fn depths(&self) -> &[f32] {
184        &self.depths
185    }
186
187    /// Returns the pixel colors.
188    #[must_use]
189    pub fn colors(&self) -> &[Vec4] {
190        &self.colors
191    }
192
193    /// Returns the normals, if set.
194    #[must_use]
195    pub fn normals(&self) -> Option<&[Vec3]> {
196        self.normals.as_deref()
197    }
198
199    /// Gets the image origin.
200    #[must_use]
201    pub fn origin(&self) -> ImageOrigin {
202        self.origin
203    }
204
205    /// Sets the image origin.
206    pub fn set_origin(&mut self, origin: ImageOrigin) -> &mut Self {
207        self.origin = origin;
208        self
209    }
210
211    /// Returns the depth at pixel (x, y).
212    #[must_use]
213    pub fn depth_at(&self, x: u32, y: u32) -> f32 {
214        let row = match self.origin {
215            ImageOrigin::UpperLeft => y,
216            ImageOrigin::LowerLeft => self.height - 1 - y,
217        };
218        self.depths[(row * self.width + x) as usize]
219    }
220
221    /// Returns the color at pixel (x, y).
222    #[must_use]
223    pub fn color_at(&self, x: u32, y: u32) -> Vec4 {
224        let row = match self.origin {
225            ImageOrigin::UpperLeft => y,
226            ImageOrigin::LowerLeft => self.height - 1 - y,
227        };
228        self.colors[(row * self.width + x) as usize]
229    }
230}
231
232impl Quantity for FloatingColorRenderImage {
233    fn as_any(&self) -> &dyn std::any::Any {
234        self
235    }
236    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
237        self
238    }
239    fn name(&self) -> &str {
240        &self.name
241    }
242    #[allow(clippy::unnecessary_literal_bound)]
243    fn structure_name(&self) -> &str {
244        "" // No parent structure
245    }
246    fn kind(&self) -> QuantityKind {
247        QuantityKind::Color
248    }
249    fn is_enabled(&self) -> bool {
250        self.enabled
251    }
252    fn set_enabled(&mut self, enabled: bool) {
253        self.enabled = enabled;
254    }
255    fn build_ui(&mut self, _ui: &dyn std::any::Any) {}
256    fn refresh(&mut self) {}
257    fn data_size(&self) -> usize {
258        self.colors.len()
259    }
260}
261
262/// A raw color render image (direct display, no shading).
263///
264/// Stores per-pixel colors for direct display without depth compositing
265/// or lighting. Suitable for pre-lit content or UI overlays.
266pub struct FloatingRawColorImage {
267    name: String,
268    width: u32,
269    height: u32,
270    colors: Vec<Vec4>,
271    origin: ImageOrigin,
272    enabled: bool,
273}
274
275impl FloatingRawColorImage {
276    /// Creates a new raw color render image.
277    pub fn new(name: impl Into<String>, width: u32, height: u32, colors: Vec<Vec3>) -> Self {
278        let colors = colors.into_iter().map(|c| c.extend(1.0)).collect();
279        Self {
280            name: name.into(),
281            width,
282            height,
283            colors,
284            origin: ImageOrigin::default(),
285            enabled: true,
286        }
287    }
288
289    /// Returns the image width.
290    #[must_use]
291    pub fn width(&self) -> u32 {
292        self.width
293    }
294
295    /// Returns the image height.
296    #[must_use]
297    pub fn height(&self) -> u32 {
298        self.height
299    }
300
301    /// Returns the pixel colors.
302    #[must_use]
303    pub fn colors(&self) -> &[Vec4] {
304        &self.colors
305    }
306
307    /// Gets the image origin.
308    #[must_use]
309    pub fn origin(&self) -> ImageOrigin {
310        self.origin
311    }
312
313    /// Sets the image origin.
314    pub fn set_origin(&mut self, origin: ImageOrigin) -> &mut Self {
315        self.origin = origin;
316        self
317    }
318
319    /// Returns the color at pixel (x, y).
320    #[must_use]
321    pub fn color_at(&self, x: u32, y: u32) -> Vec4 {
322        let row = match self.origin {
323            ImageOrigin::UpperLeft => y,
324            ImageOrigin::LowerLeft => self.height - 1 - y,
325        };
326        self.colors[(row * self.width + x) as usize]
327    }
328}
329
330impl Quantity for FloatingRawColorImage {
331    fn as_any(&self) -> &dyn std::any::Any {
332        self
333    }
334    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
335        self
336    }
337    fn name(&self) -> &str {
338        &self.name
339    }
340    #[allow(clippy::unnecessary_literal_bound)]
341    fn structure_name(&self) -> &str {
342        "" // No parent structure
343    }
344    fn kind(&self) -> QuantityKind {
345        QuantityKind::Color
346    }
347    fn is_enabled(&self) -> bool {
348        self.enabled
349    }
350    fn set_enabled(&mut self, enabled: bool) {
351        self.enabled = enabled;
352    }
353    fn build_ui(&mut self, _ui: &dyn std::any::Any) {}
354    fn refresh(&mut self) {}
355    fn data_size(&self) -> usize {
356        self.colors.len()
357    }
358}
359
360#[cfg(test)]
361mod tests {
362    use super::*;
363
364    #[test]
365    fn test_depth_render_image_creation() {
366        let depths = vec![1.0, 2.0, 3.0, 4.0];
367        let img = FloatingDepthRenderImage::new("depth", 2, 2, depths);
368
369        assert_eq!(img.name(), "depth");
370        assert_eq!(img.width(), 2);
371        assert_eq!(img.height(), 2);
372        assert_eq!(img.data_size(), 4);
373        assert!(img.normals().is_none());
374    }
375
376    #[test]
377    fn test_depth_render_image_pixel_access() {
378        let depths = vec![1.0, 2.0, 3.0, 4.0];
379        let img = FloatingDepthRenderImage::new("depth", 2, 2, depths);
380
381        assert_eq!(img.depth_at(0, 0), 1.0);
382        assert_eq!(img.depth_at(1, 0), 2.0);
383        assert_eq!(img.depth_at(0, 1), 3.0);
384        assert_eq!(img.depth_at(1, 1), 4.0);
385    }
386
387    #[test]
388    fn test_depth_render_image_has_depth() {
389        let depths = vec![1.0, f32::INFINITY, 0.0, -1.0];
390        let img = FloatingDepthRenderImage::new("depth", 2, 2, depths);
391
392        assert!(img.has_depth(0, 0)); // 1.0 — valid
393        assert!(!img.has_depth(1, 0)); // inf — invalid
394        assert!(!img.has_depth(0, 1)); // 0.0 — invalid (not > 0)
395        assert!(!img.has_depth(1, 1)); // -1.0 — invalid
396    }
397
398    #[test]
399    fn test_depth_render_image_with_normals() {
400        let depths = vec![1.0; 4];
401        let normals = vec![Vec3::Z; 4];
402        let mut img = FloatingDepthRenderImage::new("depth", 2, 2, depths);
403        img.set_normals(normals);
404
405        assert!(img.normals().is_some());
406        assert_eq!(img.normals().unwrap().len(), 4);
407    }
408
409    #[test]
410    fn test_color_render_image_creation() {
411        let depths = vec![1.0, 2.0, 3.0, 4.0];
412        let colors = vec![Vec3::X, Vec3::Y, Vec3::Z, Vec3::ONE];
413        let img = FloatingColorRenderImage::new("colored", 2, 2, depths, colors);
414
415        assert_eq!(img.name(), "colored");
416        assert_eq!(img.width(), 2);
417        assert_eq!(img.height(), 2);
418        assert_eq!(img.data_size(), 4);
419        assert_eq!(img.kind(), QuantityKind::Color);
420    }
421
422    #[test]
423    fn test_color_render_image_pixel_access() {
424        let depths = vec![1.0, 2.0, 3.0, 4.0];
425        let colors = vec![Vec3::X, Vec3::Y, Vec3::Z, Vec3::ONE];
426        let img = FloatingColorRenderImage::new("colored", 2, 2, depths, colors);
427
428        assert_eq!(img.depth_at(0, 0), 1.0);
429        assert_eq!(img.color_at(0, 0), Vec4::new(1.0, 0.0, 0.0, 1.0));
430        assert_eq!(img.color_at(1, 1), Vec4::new(1.0, 1.0, 1.0, 1.0));
431    }
432
433    #[test]
434    fn test_raw_color_image_creation() {
435        let colors = vec![Vec3::X, Vec3::Y, Vec3::Z, Vec3::ONE];
436        let img = FloatingRawColorImage::new("raw", 2, 2, colors);
437
438        assert_eq!(img.name(), "raw");
439        assert_eq!(img.data_size(), 4);
440        assert_eq!(img.kind(), QuantityKind::Color);
441    }
442
443    #[test]
444    fn test_raw_color_image_pixel_access() {
445        let colors = vec![Vec3::X, Vec3::Y, Vec3::Z, Vec3::ONE];
446        let img = FloatingRawColorImage::new("raw", 2, 2, colors);
447
448        assert_eq!(img.color_at(0, 0), Vec4::new(1.0, 0.0, 0.0, 1.0));
449        assert_eq!(img.color_at(1, 0), Vec4::new(0.0, 1.0, 0.0, 1.0));
450        assert_eq!(img.color_at(0, 1), Vec4::new(0.0, 0.0, 1.0, 1.0));
451        assert_eq!(img.color_at(1, 1), Vec4::new(1.0, 1.0, 1.0, 1.0));
452    }
453}