polyscope_render/
screenshot.rs

1//! Screenshot functionality for capturing rendered frames.
2
3use image::{ImageBuffer, Rgba};
4use std::path::Path;
5
6/// Options for taking screenshots.
7#[derive(Debug, Clone, Default)]
8pub struct ScreenshotOptions {
9    /// Whether to use transparent background (PNG only).
10    pub transparent_background: bool,
11}
12
13/// Saves raw BGRA pixel data to an image file.
14///
15/// # Arguments
16/// * `filename` - Output filename (supports .png, .jpg, .jpeg)
17/// * `data` - Raw BGRA pixel data (4 bytes per pixel, as from wgpu `Bgra8UnormSrgb` format)
18/// * `width` - Image width in pixels
19/// * `height` - Image height in pixels
20///
21/// # Errors
22/// Returns an error if the file cannot be written or format is unsupported.
23pub fn save_image(
24    filename: &str,
25    data: &[u8],
26    width: u32,
27    height: u32,
28) -> Result<(), ScreenshotError> {
29    let path = Path::new(filename);
30    let extension = path
31        .extension()
32        .and_then(|e| e.to_str())
33        .map(str::to_lowercase)
34        .unwrap_or_default();
35
36    // Convert BGRA to RGBA (wgpu surface format is Bgra8UnormSrgb)
37    let mut rgba_data = data.to_vec();
38    for chunk in rgba_data.chunks_exact_mut(4) {
39        chunk.swap(0, 2); // Swap B and R
40    }
41
42    // Create image buffer from converted RGBA data
43    // Note: wgpu uses top-left origin, so no vertical flip needed
44    let img: ImageBuffer<Rgba<u8>, Vec<u8>> =
45        ImageBuffer::from_raw(width, height, rgba_data).ok_or(ScreenshotError::InvalidImageData)?;
46
47    match extension.as_str() {
48        "png" => {
49            img.save_with_format(path, image::ImageFormat::Png)?;
50        }
51        "jpg" | "jpeg" => {
52            // Convert to RGB for JPEG (no alpha)
53            let rgb_img = image::DynamicImage::ImageRgba8(img).to_rgb8();
54            rgb_img.save_with_format(path, image::ImageFormat::Jpeg)?;
55        }
56        _ => {
57            return Err(ScreenshotError::UnsupportedFormat(extension));
58        }
59    }
60
61    Ok(())
62}
63
64/// Saves raw BGRA pixel data to a PNG buffer in memory.
65///
66/// # Arguments
67/// * `data` - Raw BGRA pixel data (4 bytes per pixel, as from wgpu `Bgra8UnormSrgb` format)
68/// * `width` - Image width in pixels
69/// * `height` - Image height in pixels
70///
71/// # Returns
72/// PNG-encoded image data as a byte vector.
73pub fn save_to_buffer(data: &[u8], width: u32, height: u32) -> Result<Vec<u8>, ScreenshotError> {
74    // Convert BGRA to RGBA
75    let mut rgba_data = data.to_vec();
76    for chunk in rgba_data.chunks_exact_mut(4) {
77        chunk.swap(0, 2); // Swap B and R
78    }
79
80    // Note: wgpu uses top-left origin, so no vertical flip needed
81    let img: ImageBuffer<Rgba<u8>, Vec<u8>> =
82        ImageBuffer::from_raw(width, height, rgba_data).ok_or(ScreenshotError::InvalidImageData)?;
83
84    let mut buffer = std::io::Cursor::new(Vec::new());
85    img.write_to(&mut buffer, image::ImageFormat::Png)?;
86
87    Ok(buffer.into_inner())
88}
89
90/// Error type for screenshot operations.
91#[derive(Debug, thiserror::Error)]
92pub enum ScreenshotError {
93    #[error("Failed to save image: {0}")]
94    IoError(#[from] std::io::Error),
95
96    #[error("Image encoding error: {0}")]
97    ImageError(#[from] image::ImageError),
98
99    #[error("Unsupported image format: {0}")]
100    UnsupportedFormat(String),
101
102    #[error("Invalid image data")]
103    InvalidImageData,
104
105    #[error("GPU buffer mapping failed")]
106    BufferMapFailed,
107}
polyscope_render/screenshot.rs

polyscope_render/
screenshot.rs
