Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove some allocations from HDR Decoder #1965

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading