Skip to content

Commit

Permalink
Remove some allocations from HDR image
Browse files Browse the repository at this point in the history
  • Loading branch information
JulianKnodt committed Jul 19, 2023
1 parent 2a5c5bf commit 3a0e1fa
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 28 deletions.
64 changes: 39 additions & 25 deletions src/codecs/hdr/decoder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -156,14 +156,7 @@ impl<R: BufRead> HdrAdapter<R> {
fn read_image_data(&mut self, buf: &mut [u8]) -> ImageResult<()> {
assert_eq!(u64::try_from(buf.len()), Ok(self.total_bytes()));
match self.inner.take() {
Some(decoder) => {
let img: Vec<Rgb<u8>> = decoder.read_image_ldr()?;
for (i, Rgb(data)) in img.into_iter().enumerate() {
buf[(i * 3)..][..3].copy_from_slice(&data);
}

Ok(())
}
Some(decoder) => decoder.read_image_ldr_buf(buf),
None => Err(ImageError::Parameter(ParameterError::from_kind(
ParameterErrorKind::NoMoreData,
))),
Expand Down Expand Up @@ -444,29 +437,46 @@ impl<R: BufRead> HdrDecoder<R> {

/// Consumes decoder and returns a vector of transformed pixels
pub fn read_image_transform<T: Send, F: Send + Sync + Fn(Rgbe8Pixel) -> T>(
self,
f: F,
output_slice: &mut [T],
) -> ImageResult<()> {
// a bit hacky, but this function was in the public API and
// thus needs to preserve it's signature.
self.read_image_transform_flat(move |v| [f(v)], output_slice)
}

/// Consumes decoder and returns a vector of transformed pixels
fn read_image_transform_flat<
T: Send,
F: Send + Sync + Fn(Rgbe8Pixel) -> [T; N],
const N: usize,
>(
mut self,
f: F,
output_slice: &mut [T],
) -> ImageResult<()> {
assert_eq!(
output_slice.len(),
self.width as usize * self.height as usize
self.width as usize * self.height as usize * N
);

// Don't read anything if image is empty
if self.width == 0 || self.height == 0 {
return Ok(());
}

let chunks_iter = output_slice.chunks_mut(self.width as usize);
let chunks_iter = output_slice.chunks_mut(self.width as usize * N);

let mut buf = vec![Default::default(); self.width as usize];
for chunk in chunks_iter {
// read_scanline overwrites the entire buffer or returns an Err,
// so not resetting the buffer here is ok.
read_scanline(&mut self.r, &mut buf[..])?;
for (dst, &pix) in chunk.iter_mut().zip(buf.iter()) {
*dst = f(pix);
for (dst, &pix) in chunk.chunks_mut(N).zip(buf.iter()) {
for (d, s) in dst.iter_mut().zip(IntoIterator::into_iter(f(pix))) {
*d = s;
}
}
}
Ok(())
Expand All @@ -475,17 +485,22 @@ impl<R: BufRead> HdrDecoder<R> {
/// Consumes decoder and returns a vector of `Rgb<u8>` pixels.
/// scale = 1, gamma = 2.2
pub fn read_image_ldr(self) -> ImageResult<Vec<Rgb<u8>>> {
let mut ret = vec![Rgb([0, 0, 0]); self.width as usize * self.height as usize];
self.read_image_transform(|pix| pix.to_ldr(), &mut ret[..])?;
Ok(ret)
let sz = (self.width * self.height) as usize;
let mut buf = vec![Rgb([0; 3]); sz];
self.read_image_transform(|pix| pix.to_ldr(), &mut buf[..])?;
Ok(buf)
}
fn read_image_ldr_buf(self, buf: &mut [u8]) -> ImageResult<()> {
let sz = (self.width * self.height * 3) as usize;
assert!(buf.len() >= sz);
self.read_image_transform_flat(|pix| pix.to_ldr().0, &mut buf[..sz])
}

/// Consumes decoder and returns a vector of `Rgb<f32>` pixels.
///
pub fn read_image_hdr(self) -> ImageResult<Vec<Rgb<f32>>> {
let mut ret = vec![Rgb([0.0, 0.0, 0.0]); self.width as usize * self.height as usize];
self.read_image_transform(|pix| pix.to_hdr(), &mut ret[..])?;
Ok(ret)
let sz = (self.width * self.height) as usize;
let mut buf = vec![Rgb([0.0f32; 3]); sz];
self.read_image_transform(|pix| pix.to_hdr(), &mut buf[..])?;
Ok(buf)
}
}

Expand Down Expand Up @@ -971,15 +986,14 @@ fn split_at_first_test() {
// or return None to indicate end of file
fn read_line_u8<R: BufRead>(r: &mut R) -> ::std::io::Result<Option<Vec<u8>>> {
let mut ret = Vec::with_capacity(16);
match r.read_until(b'\n', &mut ret) {
Ok(0) => Ok(None),
Ok(_) => {
if let Some(&b'\n') = ret[..].last() {
match r.read_until(b'\n', &mut ret)? {
0 => Ok(None),
_ => {
if let Some(&b'\n') = ret.last() {
let _ = ret.pop();
}
Ok(Some(ret))
}
Err(err) => Err(err),
}
}

Expand Down
5 changes: 2 additions & 3 deletions src/codecs/openexr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -467,12 +467,11 @@ mod test {
.clone()
.join("overexposed gradient - data window equals display window.exr");

let hdr: Vec<Rgb<f32>> = crate::codecs::hdr::HdrDecoder::new(std::io::BufReader::new(
let hdr_decoder = crate::codecs::hdr::HdrDecoder::new(std::io::BufReader::new(
std::fs::File::open(&reference_path).unwrap(),
))
.unwrap()
.read_image_hdr()
.unwrap();
let hdr = hdr_decoder.read_image_hdr().unwrap();

let exr_pixels: Rgb32FImage = read_as_rgb_image_from_file(exr_path).unwrap();
assert_eq!(
Expand Down

0 comments on commit 3a0e1fa

Please sign in to comment.