veloren_voxygen/render/renderer/
screenshot.rs

1use super::super::pipelines::blit;
2use common_base::prof_span;
3use crossbeam_channel;
4use tracing::error;
5
6pub type ScreenshotFn = Box<dyn FnOnce(Result<image::RgbImage, String>) + Send>;
7
8pub struct TakeScreenshot {
9    bind_group: blit::BindGroup,
10    view: wgpu::TextureView,
11    texture: wgpu::Texture,
12    buffer: wgpu::Buffer,
13    screenshot_fn: ScreenshotFn,
14    // Dimensions used for copying from the screenshot texture to a buffer
15    width: u32,
16    height: u32,
17    bytes_per_pixel: u32,
18    // Texture format
19    tex_format: wgpu::TextureFormat,
20}
21
22impl TakeScreenshot {
23    pub fn new(
24        device: &wgpu::Device,
25        blit_layout: &blit::BlitLayout,
26        sampler: &wgpu::Sampler,
27        // Used to determine the resolution and texture format
28        surface_config: &wgpu::SurfaceConfiguration,
29        // Function that is given the image after downloading it from the GPU
30        // This is executed in a background thread
31        screenshot_fn: ScreenshotFn,
32    ) -> Self {
33        let texture = device.create_texture(&wgpu::TextureDescriptor {
34            label: Some("screenshot tex"),
35            size: wgpu::Extent3d {
36                width: surface_config.width,
37                height: surface_config.height,
38                depth_or_array_layers: 1,
39            },
40            mip_level_count: 1,
41            sample_count: 1,
42            dimension: wgpu::TextureDimension::D2,
43            format: surface_config.format,
44            usage: wgpu::TextureUsages::COPY_SRC
45                | wgpu::TextureUsages::TEXTURE_BINDING
46                | wgpu::TextureUsages::RENDER_ATTACHMENT,
47            view_formats: &[],
48        });
49
50        let view = texture.create_view(&wgpu::TextureViewDescriptor {
51            label: Some("screenshot tex view"),
52            format: Some(surface_config.format),
53            dimension: Some(wgpu::TextureViewDimension::D2),
54            usage: None,
55            aspect: wgpu::TextureAspect::All,
56            base_mip_level: 0,
57            mip_level_count: None,
58            base_array_layer: 0,
59            array_layer_count: None,
60        });
61
62        let bind_group = blit_layout.bind(device, &view, sampler);
63
64        let bytes_per_pixel = surface_config.format.block_copy_size(None).unwrap();
65        let padded_bytes_per_row = padded_bytes_per_row(surface_config.width, bytes_per_pixel);
66
67        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
68            label: Some("screenshot download buffer"),
69            size: (padded_bytes_per_row * surface_config.height) as u64,
70            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
71            mapped_at_creation: false,
72        });
73
74        Self {
75            bind_group,
76            texture,
77            view,
78            buffer,
79            screenshot_fn,
80            width: surface_config.width,
81            height: surface_config.height,
82            bytes_per_pixel,
83            tex_format: surface_config.format,
84        }
85    }
86
87    /// Get the texture view for the screenshot
88    /// This can then be used as a render attachment
89    pub fn texture_view(&self) -> &wgpu::TextureView { &self.view }
90
91    /// Get the bind group used for blitting the screenshot to the current
92    /// swapchain image
93    pub fn bind_group(&self) -> &wgpu::BindGroup { &self.bind_group.bind_group }
94
95    /// Call this after rendering to the screenshot texture
96    ///
97    /// Issues a command to copy from the texture to a buffer and returns a
98    /// closure that needs to be called after submitting the encoder
99    /// to the queue. When called, the closure will spawn a new thread for
100    /// async mapping of the buffer and downloading of the screenshot.
101    pub fn copy_to_buffer(self, encoder: &mut wgpu::CommandEncoder) -> impl FnOnce() + use<> {
102        // Calculate padded bytes per row
103        let padded_bytes_per_row = padded_bytes_per_row(self.width, self.bytes_per_pixel);
104        // Copy image to a buffer
105        encoder.copy_texture_to_buffer(
106            wgpu::TexelCopyTextureInfo {
107                texture: &self.texture,
108                mip_level: 0,
109                origin: wgpu::Origin3d::ZERO,
110                aspect: wgpu::TextureAspect::All,
111            },
112            wgpu::TexelCopyBufferInfo {
113                buffer: &self.buffer,
114                layout: wgpu::TexelCopyBufferLayout {
115                    offset: 0,
116                    bytes_per_row: Some(padded_bytes_per_row),
117                    rows_per_image: None,
118                },
119            },
120            wgpu::Extent3d {
121                width: self.width,
122                height: self.height,
123                depth_or_array_layers: 1,
124            },
125        );
126
127        move || {
128            // Send buffer to another thread for async
129            // mapping, downloading, and passing to the given handler function
130            // (which probably saves it to the disk)
131            std::thread::Builder::new()
132                .name("screenshot".into())
133                .spawn(move || {
134                    self.download_and_handle_internal();
135                })
136                .expect("Failed to spawn screenshot thread");
137        }
138    }
139
140    /// Don't call this from the main loop, it will block for a while
141    fn download_and_handle_internal(self) {
142        prof_span!("download_and_handle_internal");
143        // Calculate padded bytes per row
144        let padded_bytes_per_row = padded_bytes_per_row(self.width, self.bytes_per_pixel);
145
146        // Map buffer
147        let buffer = std::sync::Arc::new(self.buffer);
148        let buffer2 = std::sync::Arc::clone(&buffer);
149        let buffer_slice = buffer.slice(..);
150        let (map_result_sender, map_result_receiver) = crossbeam_channel::bounded(1);
151        buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
152            map_result_sender
153                .send(result)
154                .expect("seems like the receiver broke, which should not happen");
155        });
156        let result = match map_result_receiver.recv() {
157            Ok(result) => result,
158            Err(e) => {
159                error!(
160                    ?e,
161                    "map_async never send the result for the screenshot mapping"
162                );
163                return;
164            },
165        };
166        let padded_buffer;
167        let buffer_slice = buffer2.slice(..);
168        let rows = match result {
169            Ok(()) => {
170                // Copy to a Vec
171                padded_buffer = buffer_slice.get_mapped_range();
172                padded_buffer
173                    .chunks(padded_bytes_per_row as usize)
174                    .map(|padded_chunk| {
175                        &padded_chunk[..self.width as usize * self.bytes_per_pixel as usize]
176                    })
177            },
178            // Error
179            Err(err) => {
180                error!(
181                    ?err,
182                    "Failed to map buffer for downloading a screenshot from the GPU"
183                );
184                return;
185            },
186        };
187
188        // Note: we don't use bytes_per_pixel here since we expect only certain formats
189        // below.
190        let bytes_per_rgb = 3;
191        let mut pixel_bytes =
192            Vec::with_capacity(self.width as usize * self.height as usize * bytes_per_rgb);
193        // Construct image
194        let image = match self.tex_format {
195            wgpu::TextureFormat::Bgra8UnormSrgb => {
196                prof_span!("copy image");
197                rows.for_each(|row| {
198                    let (pixels, rest) = row.as_chunks();
199                    assert!(
200                        rest.is_empty(),
201                        "Always valid because each pixel uses four bytes"
202                    );
203                    // Swap blue and red components and drop alpha to get a RGB texture.
204                    for &[b, g, r, _a] in pixels {
205                        pixel_bytes.extend_from_slice(&[r, g, b])
206                    }
207                });
208
209                Ok(pixel_bytes)
210            },
211            wgpu::TextureFormat::Rgba8UnormSrgb => {
212                prof_span!("copy image");
213                rows.for_each(|row| {
214                    let (pixels, rest) = row.as_chunks();
215                    assert!(
216                        rest.is_empty(),
217                        "Always valid because each pixel uses four bytes"
218                    );
219                    // Drop alpha to get a RGB texture.
220                    for &[r, g, b, _a] in pixels {
221                        pixel_bytes.extend_from_slice(&[r, g, b])
222                    }
223                });
224
225                Ok(pixel_bytes)
226            },
227            format => Err(format!(
228                "Unhandled format for screenshot texture: {:?}",
229                format,
230            )),
231        }
232        .map(|pixel_bytes| {
233            image::RgbImage::from_vec(self.width, self.height, pixel_bytes).expect(
234                "Failed to create ImageBuffer! Buffer was not large enough. This should not occur",
235            )
236        });
237        // Call supplied handler
238        (self.screenshot_fn)(image);
239    }
240}
241
242// Graphics API requires a specific alignment for buffer copies
243fn padded_bytes_per_row(width: u32, bytes_per_pixel: u32) -> u32 {
244    let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
245    let unpadded_bytes_per_row = width * bytes_per_pixel;
246    let padding = (align - unpadded_bytes_per_row % align) % align;
247    unpadded_bytes_per_row + padding
248}
veloren_voxygen/render/renderer/screenshot.rs

veloren_voxygen/render/renderer/
screenshot.rs
