Remove unwrap from write_png (#261)

* Remove unwrap from write_png

---------

Co-authored-by: Max Ammann <maximilian.ammann@trailofbits.com>
This commit is contained in:
Tristram Gräbener 2023-02-27 14:38:15 +01:00 committed by GitHub
parent 1fc7c4f7a9
commit 5cd3e503cc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 27 additions and 16 deletions

View File

@ -55,10 +55,12 @@ impl System for WriteSurfaceBufferSystem {
let padded_buffer = buffer_slice.get_mapped_range();
if self.write_to_disk {
buffered_texture.write_png(
&padded_buffer,
format!("frame_{}.png", current_frame).as_str(),
);
buffered_texture
.write_png(
&padded_buffer,
format!("frame_{}.png", current_frame).as_str(),
)
.expect("Could save frame to disk");
}
// With the current interface, we have to make sure all mapped views are

View File

@ -80,6 +80,15 @@ pub struct BufferedTextureHead {
buffer_dimensions: BufferDimensions,
}
#[cfg(feature = "headless")]
#[derive(thiserror::Error, Debug)]
pub enum WriteImageError {
#[error("error while rendering to image")]
WriteImage(#[from] png::EncodingError),
#[error("could not create file to save as an image")]
CreateImageFileFailed(#[from] std::io::Error),
}
#[cfg(feature = "headless")]
impl BufferedTextureHead {
pub fn map_async(&self, device: &wgpu::Device) -> wgpu::BufferSlice {
@ -98,32 +107,32 @@ impl BufferedTextureHead {
self.output_buffer.unmap();
}
pub fn write_png<'a>(&self, padded_buffer: &wgpu::BufferView<'a>, png_output_path: &str) {
pub fn write_png<'a>(
&self,
padded_buffer: &wgpu::BufferView<'a>,
png_output_path: &str,
) -> Result<(), WriteImageError> {
use std::{fs::File, io::Write};
let mut png_encoder = png::Encoder::new(
File::create(png_output_path).unwrap(), // TODO: Remove unwrap
File::create(png_output_path)?,
self.buffer_dimensions.width as u32,
self.buffer_dimensions.height as u32,
);
png_encoder.set_depth(png::BitDepth::Eight);
png_encoder.set_color(png::ColorType::Rgba);
let mut png_writer = png_encoder
.write_header()
.unwrap() // TODO: Remove unwrap
.into_stream_writer_with_size(
self.buffer_dimensions.unpadded_bytes_per_row.get() as usize
)
.unwrap(); // TODO: Remove unwrap
let mut png_writer = png_encoder.write_header()?.into_stream_writer_with_size(
self.buffer_dimensions.unpadded_bytes_per_row.get() as usize,
)?;
// from the padded_buffer we write just the unpadded bytes into the image
for chunk in
padded_buffer.chunks(self.buffer_dimensions.padded_bytes_per_row.get() as usize)
{
png_writer
.write_all(&chunk[..self.buffer_dimensions.unpadded_bytes_per_row.get() as usize])
.unwrap(); // TODO: Remove unwrap
.write_all(&chunk[..self.buffer_dimensions.unpadded_bytes_per_row.get() as usize])?
}
png_writer.finish().unwrap(); // TODO: Remove unwrap
png_writer.finish()?;
Ok(())
}
pub fn copy_texture(&self) -> wgpu::ImageCopyTexture<'_> {