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(); let padded_buffer = buffer_slice.get_mapped_range();
if self.write_to_disk { if self.write_to_disk {
buffered_texture.write_png( buffered_texture
.write_png(
&padded_buffer, &padded_buffer,
format!("frame_{}.png", current_frame).as_str(), 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 // 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, 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")] #[cfg(feature = "headless")]
impl BufferedTextureHead { impl BufferedTextureHead {
pub fn map_async(&self, device: &wgpu::Device) -> wgpu::BufferSlice { pub fn map_async(&self, device: &wgpu::Device) -> wgpu::BufferSlice {
@ -98,32 +107,32 @@ impl BufferedTextureHead {
self.output_buffer.unmap(); 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}; use std::{fs::File, io::Write};
let mut png_encoder = png::Encoder::new( 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.width as u32,
self.buffer_dimensions.height as u32, self.buffer_dimensions.height as u32,
); );
png_encoder.set_depth(png::BitDepth::Eight); png_encoder.set_depth(png::BitDepth::Eight);
png_encoder.set_color(png::ColorType::Rgba); png_encoder.set_color(png::ColorType::Rgba);
let mut png_writer = png_encoder let mut png_writer = png_encoder.write_header()?.into_stream_writer_with_size(
.write_header() self.buffer_dimensions.unpadded_bytes_per_row.get() as usize,
.unwrap() // TODO: Remove unwrap )?;
.into_stream_writer_with_size(
self.buffer_dimensions.unpadded_bytes_per_row.get() as usize
)
.unwrap(); // TODO: Remove unwrap
// from the padded_buffer we write just the unpadded bytes into the image // from the padded_buffer we write just the unpadded bytes into the image
for chunk in for chunk in
padded_buffer.chunks(self.buffer_dimensions.padded_bytes_per_row.get() as usize) padded_buffer.chunks(self.buffer_dimensions.padded_bytes_per_row.get() as usize)
{ {
png_writer png_writer
.write_all(&chunk[..self.buffer_dimensions.unpadded_bytes_per_row.get() as usize]) .write_all(&chunk[..self.buffer_dimensions.unpadded_bytes_per_row.get() as usize])?
.unwrap(); // TODO: Remove unwrap
} }
png_writer.finish().unwrap(); // TODO: Remove unwrap png_writer.finish()?;
Ok(())
} }
pub fn copy_texture(&self) -> wgpu::ImageCopyTexture<'_> { pub fn copy_texture(&self) -> wgpu::ImageCopyTexture<'_> {