egui/crates/egui_kittest/src/texture_to_image.rs

84 lines
3.0 KiB
Rust

use egui_wgpu::wgpu;
use egui_wgpu::wgpu::{Device, Extent3d, Queue, Texture};
use image::RgbaImage;
use std::iter;
use std::mem::size_of;
use std::sync::mpsc::channel;
pub(crate) fn texture_to_image(device: &Device, queue: &Queue, texture: &Texture) -> RgbaImage {
let buffer_dimensions =
BufferDimensions::new(texture.width() as usize, texture.height() as usize);
let output_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Texture to bytes output buffer"),
size: (buffer_dimensions.padded_bytes_per_row * buffer_dimensions.height) as u64,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
label: Some("Texture to bytes encoder"),
});
// Copy the data from the texture to the buffer
encoder.copy_texture_to_buffer(
texture.as_image_copy(),
wgpu::ImageCopyBuffer {
buffer: &output_buffer,
layout: wgpu::ImageDataLayout {
offset: 0,
bytes_per_row: Some(buffer_dimensions.padded_bytes_per_row as u32),
rows_per_image: None,
},
},
Extent3d {
width: texture.width(),
height: texture.height(),
depth_or_array_layers: 1,
},
);
let submission_index = queue.submit(iter::once(encoder.finish()));
// Note that we're not calling `.await` here.
let buffer_slice = output_buffer.slice(..);
// Sets the buffer up for mapping, sending over the result of the mapping back to us when it is finished.
let (sender, receiver) = channel();
buffer_slice.map_async(wgpu::MapMode::Read, move |v| drop(sender.send(v)));
// Poll the device in a blocking manner so that our future resolves.
device.poll(wgpu::Maintain::WaitForSubmissionIndex(submission_index));
receiver.recv().unwrap().unwrap();
let buffer_slice = output_buffer.slice(..);
let data = buffer_slice.get_mapped_range();
let data = data
.chunks_exact(buffer_dimensions.padded_bytes_per_row)
.flat_map(|row| row.iter().take(buffer_dimensions.unpadded_bytes_per_row))
.copied()
.collect::<Vec<_>>();
RgbaImage::from_raw(texture.width(), texture.height(), data).expect("Failed to create image")
}
struct BufferDimensions {
height: usize,
unpadded_bytes_per_row: usize,
padded_bytes_per_row: usize,
}
impl BufferDimensions {
fn new(width: usize, height: usize) -> Self {
let bytes_per_pixel = size_of::<u32>();
let unpadded_bytes_per_row = width * bytes_per_pixel;
let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT as usize;
let padded_bytes_per_row_padding = (align - unpadded_bytes_per_row % align) % align;
let padded_bytes_per_row = unpadded_bytes_per_row + padded_bytes_per_row_padding;
Self {
height,
unpadded_bytes_per_row,
padded_bytes_per_row,
}
}
}