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

veloren_voxygen/render/renderer/
screenshot.rs
