diff options
Diffstat (limited to 'vendor/exr/src/compression/piz/mod.rs')
| -rw-r--r-- | vendor/exr/src/compression/piz/mod.rs | 437 |
1 files changed, 437 insertions, 0 deletions
diff --git a/vendor/exr/src/compression/piz/mod.rs b/vendor/exr/src/compression/piz/mod.rs new file mode 100644 index 0000000..1d77663 --- /dev/null +++ b/vendor/exr/src/compression/piz/mod.rs @@ -0,0 +1,437 @@ + + +//! The PIZ compression method is a wavelet compression, +//! based on the PIZ image format, customized for OpenEXR. +// inspired by https://github.com/AcademySoftwareFoundation/openexr/blob/master/OpenEXR/IlmImf/ImfPizCompressor.cpp + +mod huffman; +mod wavelet; + +use crate::prelude::*; +use crate::io::Data; +use crate::meta::attribute::*; +use crate::compression::{ByteVec, Bytes, mod_p}; +use crate::error::{usize_to_i32, usize_to_u16}; +use std::convert::TryFrom; + + +const U16_RANGE: usize = (1_i32 << 16_i32) as usize; +const BITMAP_SIZE: usize = (U16_RANGE as i32 >> 3_i32) as usize; + +#[derive(Debug)] +struct ChannelData { + tmp_start_index: usize, + tmp_end_index: usize, + + resolution: Vec2<usize>, + y_sampling: usize, + samples_per_pixel: usize, +} + + +pub fn decompress( + channels: &ChannelList, + compressed: ByteVec, + rectangle: IntegerBounds, + expected_byte_size: usize, // TODO remove expected byte size as it can be computed with `rectangle.size.area() * channels.bytes_per_pixel` + pedantic: bool +) -> Result<ByteVec> +{ + let expected_u16_count = expected_byte_size / 2; + debug_assert_eq!(expected_byte_size, rectangle.size.area() * channels.bytes_per_pixel); + debug_assert!(!channels.list.is_empty()); + + if compressed.is_empty() { + return Ok(Vec::new()); + } + + debug_assert_ne!(expected_u16_count, 0); + + let mut bitmap = vec![0_u8; BITMAP_SIZE]; // FIXME use bit_vec! + + let mut remaining_input = compressed.as_slice(); + let min_non_zero = u16::read(&mut remaining_input)? as usize; + let max_non_zero = u16::read(&mut remaining_input)? as usize; + + if max_non_zero >= BITMAP_SIZE || min_non_zero >= BITMAP_SIZE { + return Err(Error::invalid("compression data")); + } + + if min_non_zero <= max_non_zero { + u8::read_slice(&mut remaining_input, &mut bitmap[min_non_zero ..= max_non_zero])?; + } + + let (lookup_table, max_value) = reverse_lookup_table_from_bitmap(&bitmap); + + { + let length = i32::read(&mut remaining_input)?; + if pedantic && length as i64 != remaining_input.len() as i64 { + // TODO length might be smaller than remaining?? + return Err(Error::invalid("compression data")); + } + } + + let mut tmp_u16_buffer = huffman::decompress(remaining_input, expected_u16_count)?; + + let mut channel_data: SmallVec<[ChannelData; 6]> = { + let mut tmp_read_index = 0; + + let channel_data = channels.list.iter().map(|channel| { + let channel_data = ChannelData { + tmp_start_index: tmp_read_index, + tmp_end_index: tmp_read_index, + y_sampling: channel.sampling.y(), + resolution: channel.subsampled_resolution(rectangle.size), + samples_per_pixel: channel.sample_type.bytes_per_sample() / SampleType::F16.bytes_per_sample() + }; + + tmp_read_index += channel_data.resolution.area() * channel_data.samples_per_pixel; + channel_data + }).collect(); + + debug_assert_eq!(tmp_read_index, expected_u16_count); + channel_data + }; + + for channel in &channel_data { + let u16_count = channel.resolution.area() * channel.samples_per_pixel; + let u16s = &mut tmp_u16_buffer[channel.tmp_start_index .. channel.tmp_start_index + u16_count]; + + for offset in 0..channel.samples_per_pixel { // if channel is 32 bit, compress interleaved as two 16 bit values + wavelet::decode( + &mut u16s[offset..], + channel.resolution, + Vec2(channel.samples_per_pixel, channel.resolution.x() * channel.samples_per_pixel), + max_value + )?; + } + } + + // Expand the pixel data to their original range + apply_lookup_table(&mut tmp_u16_buffer, &lookup_table); + + // let out_buffer_size = (max_scan_line_size * scan_line_count) + 65536 + 8192; // TODO not use expected byte size? + let mut out = Vec::with_capacity(expected_byte_size); + + for y in rectangle.position.y() .. rectangle.end().y() { + for channel in &mut channel_data { + if mod_p(y, usize_to_i32(channel.y_sampling)) != 0 { + continue; + } + + let u16s_per_line = channel.resolution.x() * channel.samples_per_pixel; + let next_tmp_end_index = channel.tmp_end_index + u16s_per_line; + let values = &tmp_u16_buffer[channel.tmp_end_index .. next_tmp_end_index]; + channel.tmp_end_index = next_tmp_end_index; + + // TODO do not convert endianness for f16-only images + // see https://github.com/AcademySoftwareFoundation/openexr/blob/3bd93f85bcb74c77255f28cdbb913fdbfbb39dfe/OpenEXR/IlmImf/ImfTiledOutputFile.cpp#L750-L842 + // We can support uncompressed data in the machine's native format + // if all image channels are of type HALF, and if the Xdr and the + // native representations of a half have the same size. + u16::write_slice(&mut out, values).expect("write to in-memory failed"); + } + } + + for (previous, current) in channel_data.iter().zip(channel_data.iter().skip(1)) { + debug_assert_eq!(previous.tmp_end_index, current.tmp_start_index); + } + + debug_assert_eq!(channel_data.last().unwrap().tmp_end_index, tmp_u16_buffer.len()); + debug_assert_eq!(out.len(), expected_byte_size); + + // TODO optimize for when all channels are f16! + // we should be able to omit endianness conversions in that case + // see https://github.com/AcademySoftwareFoundation/openexr/blob/3bd93f85bcb74c77255f28cdbb913fdbfbb39dfe/OpenEXR/IlmImf/ImfTiledOutputFile.cpp#L750-L842 + Ok(super::convert_little_endian_to_current(out, channels, rectangle)) +} + + + +pub fn compress( + channels: &ChannelList, + uncompressed: ByteVec, + rectangle: IntegerBounds +) -> Result<ByteVec> +{ + if uncompressed.is_empty() { + return Ok(Vec::new()); + } + + // TODO do not convert endianness for f16-only images + // see https://github.com/AcademySoftwareFoundation/openexr/blob/3bd93f85bcb74c77255f28cdbb913fdbfbb39dfe/OpenEXR/IlmImf/ImfTiledOutputFile.cpp#L750-L842 + let uncompressed = super::convert_current_to_little_endian(uncompressed, channels, rectangle); + let uncompressed = uncompressed.as_slice();// TODO no alloc + + let mut tmp = vec![0_u16; uncompressed.len() / 2 ]; + let mut channel_data: SmallVec<[ChannelData; 6]> = { + let mut tmp_end_index = 0; + + let vec = channels.list.iter().map(|channel| { + let number_samples = channel.subsampled_resolution(rectangle.size); + let byte_size = channel.sample_type.bytes_per_sample() / SampleType::F16.bytes_per_sample(); + let byte_count = byte_size * number_samples.area(); + + let channel = ChannelData { + tmp_end_index, + tmp_start_index: tmp_end_index, + y_sampling: channel.sampling.y(), + resolution: number_samples, + samples_per_pixel: byte_size, + }; + + tmp_end_index += byte_count; + channel + }).collect(); + + debug_assert_eq!(tmp_end_index, tmp.len()); + vec + }; + + let mut remaining_uncompressed_bytes = uncompressed; + for y in rectangle.position.y() .. rectangle.end().y() { + for channel in &mut channel_data { + if mod_p(y, usize_to_i32(channel.y_sampling)) != 0 { continue; } + let u16s_per_line = channel.resolution.x() * channel.samples_per_pixel; + let next_tmp_end_index = channel.tmp_end_index + u16s_per_line; + let target = &mut tmp[channel.tmp_end_index .. next_tmp_end_index]; + channel.tmp_end_index = next_tmp_end_index; + + // TODO do not convert endianness for f16-only images + // see https://github.com/AcademySoftwareFoundation/openexr/blob/3bd93f85bcb74c77255f28cdbb913fdbfbb39dfe/OpenEXR/IlmImf/ImfTiledOutputFile.cpp#L750-L842 + // We can support uncompressed data in the machine's native format + // if all image channels are of type HALF, and if the Xdr and the + // native representations of a half have the same size. + u16::read_slice(&mut remaining_uncompressed_bytes, target).expect("in-memory read failed"); + } + } + + + let (min_non_zero, max_non_zero, bitmap) = bitmap_from_data(&tmp); + let (max_value, table) = forward_lookup_table_from_bitmap(&bitmap); + apply_lookup_table(&mut tmp, &table); + + let mut piz_compressed = Vec::with_capacity(uncompressed.len() / 2); + u16::try_from(min_non_zero)?.write(&mut piz_compressed)?; + u16::try_from(max_non_zero)?.write(&mut piz_compressed)?; + + if min_non_zero <= max_non_zero { + piz_compressed.extend_from_slice(&bitmap[min_non_zero ..= max_non_zero]); + } + + for channel in channel_data { + for offset in 0 .. channel.samples_per_pixel { + wavelet::encode( + &mut tmp[channel.tmp_start_index + offset .. channel.tmp_end_index], + channel.resolution, + Vec2(channel.samples_per_pixel, channel.resolution.x() * channel.samples_per_pixel), + max_value + )?; + } + } + + let huffman_compressed: Vec<u8> = huffman::compress(&tmp)?; + u8::write_i32_sized_slice(&mut piz_compressed, &huffman_compressed).expect("in-memory write failed"); + + Ok(piz_compressed) +} + + +pub fn bitmap_from_data(data: &[u16]) -> (usize, usize, Vec<u8>) { + let mut bitmap = vec![0_u8; BITMAP_SIZE]; + + for value in data { + bitmap[*value as usize >> 3] |= 1 << (*value as u8 & 7); + } + + bitmap[0] = bitmap[0] & !1; // zero is not explicitly stored in the bitmap; we assume that the data always contain zeroes + + let min_index = bitmap.iter().position(|&value| value != 0); + let max_index = min_index.map(|min| // only if min was found + min + bitmap[min..].iter().rposition(|&value| value != 0).expect("[min] not found") + ); + + (min_index.unwrap_or(0), max_index.unwrap_or(0), bitmap) +} + +pub fn forward_lookup_table_from_bitmap(bitmap: &[u8]) -> (u16, Vec<u16>) { + debug_assert_eq!(bitmap.len(), BITMAP_SIZE); + + let mut table = vec![0_u16; U16_RANGE]; + let mut count = 0_usize; + + for (index, entry) in table.iter_mut().enumerate() { + if index == 0 || bitmap[index >> 3] as usize & (1 << (index & 7)) != 0 { + *entry = usize_to_u16(count).unwrap(); + count += 1; + } + } + + (usize_to_u16(count - 1).unwrap(), table) +} + +fn reverse_lookup_table_from_bitmap(bitmap: Bytes<'_>) -> (Vec<u16>, u16) { + let mut table = Vec::with_capacity(U16_RANGE); + + for index in 0 .. U16_RANGE { // cannot use iter because filter removes capacity sizehint + if index == 0 || ((bitmap[index >> 3] as usize & (1 << (index & 7))) != 0) { + table.push(usize_to_u16(index).unwrap()); + } + } + + debug_assert!(!table.is_empty()); + let max_value = usize_to_u16(table.len() - 1).unwrap(); + + // fill remaining up to u16 range + assert!(table.len() <= U16_RANGE); + table.resize(U16_RANGE, 0); + + (table, max_value) +} + +fn apply_lookup_table(data: &mut [u16], table: &[u16]) { + for data in data { + *data = table[*data as usize]; + } +} + +#[cfg(test)] +mod test { + use crate::prelude::*; + use crate::compression::ByteVec; + use crate::compression::piz; + use crate::meta::attribute::*; + + fn test_roundtrip_noise_with(channels: ChannelList, rectangle: IntegerBounds){ + let pixel_bytes: ByteVec = (0 .. 37).map(|_| rand::random()).collect::<Vec<u8>>().into_iter() + .cycle().take(channels.bytes_per_pixel * rectangle.size.area()) + .collect(); + + let compressed = piz::compress(&channels, pixel_bytes.clone(), rectangle).unwrap(); + let decompressed = piz::decompress(&channels, compressed, rectangle, pixel_bytes.len(), true).unwrap(); + + assert_eq!(pixel_bytes, decompressed); + } + + + #[test] + fn roundtrip_any_sample_type(){ + for &sample_type in &[SampleType::F16, SampleType::F32, SampleType::U32] { + let channel = ChannelDescription { + sample_type, + + name: Default::default(), + quantize_linearly: false, + sampling: Vec2(1,1) + }; + + let channels = ChannelList::new(smallvec![ channel.clone(), channel ]); + + let rectangle = IntegerBounds { + position: Vec2(-30, 100), + size: Vec2(1080, 720), + }; + + test_roundtrip_noise_with(channels, rectangle); + } + } + + #[test] + fn roundtrip_two_channels(){ + let channel = ChannelDescription { + sample_type: SampleType::F16, + + name: Default::default(), + quantize_linearly: false, + sampling: Vec2(1,1) + }; + + let channel2 = ChannelDescription { + sample_type: SampleType::F32, + + name: Default::default(), + quantize_linearly: false, + sampling: Vec2(1,1) + }; + + let channels = ChannelList::new(smallvec![ channel, channel2 ]); + + let rectangle = IntegerBounds { + position: Vec2(-3, 1), + size: Vec2(223, 3132), + }; + + test_roundtrip_noise_with(channels, rectangle); + } + + + + #[test] + fn roundtrip_seven_channels(){ + let channels = ChannelList::new(smallvec![ + ChannelDescription { + sample_type: SampleType::F32, + + name: Default::default(), + quantize_linearly: false, + sampling: Vec2(1,1) + }, + + ChannelDescription { + sample_type: SampleType::F32, + + name: Default::default(), + quantize_linearly: false, + sampling: Vec2(1,1) + }, + + ChannelDescription { + sample_type: SampleType::F32, + + name: Default::default(), + quantize_linearly: false, + sampling: Vec2(1,1) + }, + + ChannelDescription { + sample_type: SampleType::F16, + + name: Default::default(), + quantize_linearly: false, + sampling: Vec2(1,1) + }, + + ChannelDescription { + sample_type: SampleType::F32, + + name: Default::default(), + quantize_linearly: false, + sampling: Vec2(1,1) + }, + + ChannelDescription { + sample_type: SampleType::F32, + + name: Default::default(), + quantize_linearly: false, + sampling: Vec2(1,1) + }, + + ChannelDescription { + sample_type: SampleType::U32, + + name: Default::default(), + quantize_linearly: false, + sampling: Vec2(1,1) + }, + ]); + + let rectangle = IntegerBounds { + position: Vec2(-3, 1), + size: Vec2(1323, 132), + }; + + test_roundtrip_noise_with(channels, rectangle); + } + +}
\ No newline at end of file |