1
2
3//! The PIZ compression method is a wavelet compression,
4//! based on the PIZ image format, customized for OpenEXR.
5// inspired by https://github.com/AcademySoftwareFoundation/openexr/blob/master/OpenEXR/IlmImf/ImfPizCompressor.cpp
6
7mod huffman;
8mod wavelet;
9
10use crate::prelude::*;
11use crate::io::Data;
12use crate::meta::attribute::*;
13use crate::compression::{ByteVec, Bytes, mod_p};
14use crate::error::{usize_to_i32, usize_to_u16};
15use std::convert::TryFrom;
16
17
18const U16_RANGE: usize = (1_i32 << 16_i32) as usize;
19const BITMAP_SIZE: usize = (U16_RANGE as i32 >> 3_i32) as usize;
20
21#[derive(Debug)]
22struct ChannelData {
23 tmp_start_index: usize,
24 tmp_end_index: usize,
25
26 resolution: Vec2<usize>,
27 y_sampling: usize,
28 samples_per_pixel: usize,
29}
30
31
32pub fn decompress(
33 channels: &ChannelList,
34 compressed: ByteVec,
35 rectangle: IntegerBounds,
36 expected_byte_size: usize, // TODO remove expected byte size as it can be computed with `rectangle.size.area() * channels.bytes_per_pixel`
37 pedantic: bool
38) -> Result<ByteVec>
39{
40 let expected_u16_count = expected_byte_size / 2;
41 debug_assert_eq!(expected_byte_size, rectangle.size.area() * channels.bytes_per_pixel);
42 debug_assert!(!channels.list.is_empty());
43
44 if compressed.is_empty() {
45 return Ok(Vec::new());
46 }
47
48 debug_assert_ne!(expected_u16_count, 0);
49
50 let mut bitmap = vec![0_u8; BITMAP_SIZE]; // FIXME use bit_vec!
51
52 let mut remaining_input = compressed.as_slice();
53 let min_non_zero = u16::read(&mut remaining_input)? as usize;
54 let max_non_zero = u16::read(&mut remaining_input)? as usize;
55
56 if max_non_zero >= BITMAP_SIZE || min_non_zero >= BITMAP_SIZE {
57 return Err(Error::invalid("compression data"));
58 }
59
60 if min_non_zero <= max_non_zero {
61 u8::read_slice(&mut remaining_input, &mut bitmap[min_non_zero ..= max_non_zero])?;
62 }
63
64 let (lookup_table, max_value) = reverse_lookup_table_from_bitmap(&bitmap);
65
66 {
67 let length = i32::read(&mut remaining_input)?;
68 if pedantic && length as i64 != remaining_input.len() as i64 {
69 // TODO length might be smaller than remaining??
70 return Err(Error::invalid("compression data"));
71 }
72 }
73
74 let mut tmp_u16_buffer = huffman::decompress(remaining_input, expected_u16_count)?;
75
76 let mut channel_data: SmallVec<[ChannelData; 6]> = {
77 let mut tmp_read_index = 0;
78
79 let channel_data = channels.list.iter().map(|channel| {
80 let channel_data = ChannelData {
81 tmp_start_index: tmp_read_index,
82 tmp_end_index: tmp_read_index,
83 y_sampling: channel.sampling.y(),
84 resolution: channel.subsampled_resolution(rectangle.size),
85 samples_per_pixel: channel.sample_type.bytes_per_sample() / SampleType::F16.bytes_per_sample()
86 };
87
88 tmp_read_index += channel_data.resolution.area() * channel_data.samples_per_pixel;
89 channel_data
90 }).collect();
91
92 debug_assert_eq!(tmp_read_index, expected_u16_count);
93 channel_data
94 };
95
96 for channel in &channel_data {
97 let u16_count = channel.resolution.area() * channel.samples_per_pixel;
98 let u16s = &mut tmp_u16_buffer[channel.tmp_start_index .. channel.tmp_start_index + u16_count];
99
100 for offset in 0..channel.samples_per_pixel { // if channel is 32 bit, compress interleaved as two 16 bit values
101 wavelet::decode(
102 &mut u16s[offset..],
103 channel.resolution,
104 Vec2(channel.samples_per_pixel, channel.resolution.x() * channel.samples_per_pixel),
105 max_value
106 )?;
107 }
108 }
109
110 // Expand the pixel data to their original range
111 apply_lookup_table(&mut tmp_u16_buffer, &lookup_table);
112
113 // let out_buffer_size = (max_scan_line_size * scan_line_count) + 65536 + 8192; // TODO not use expected byte size?
114 let mut out = Vec::with_capacity(expected_byte_size);
115
116 for y in rectangle.position.y() .. rectangle.end().y() {
117 for channel in &mut channel_data {
118 if mod_p(y, usize_to_i32(channel.y_sampling)) != 0 {
119 continue;
120 }
121
122 let u16s_per_line = channel.resolution.x() * channel.samples_per_pixel;
123 let next_tmp_end_index = channel.tmp_end_index + u16s_per_line;
124 let values = &tmp_u16_buffer[channel.tmp_end_index .. next_tmp_end_index];
125 channel.tmp_end_index = next_tmp_end_index;
126
127 // TODO do not convert endianness for f16-only images
128 // see https://github.com/AcademySoftwareFoundation/openexr/blob/3bd93f85bcb74c77255f28cdbb913fdbfbb39dfe/OpenEXR/IlmImf/ImfTiledOutputFile.cpp#L750-L842
129 // We can support uncompressed data in the machine's native format
130 // if all image channels are of type HALF, and if the Xdr and the
131 // native representations of a half have the same size.
132 u16::write_slice(&mut out, values).expect("write to in-memory failed");
133 }
134 }
135
136 for (previous, current) in channel_data.iter().zip(channel_data.iter().skip(1)) {
137 debug_assert_eq!(previous.tmp_end_index, current.tmp_start_index);
138 }
139
140 debug_assert_eq!(channel_data.last().unwrap().tmp_end_index, tmp_u16_buffer.len());
141 debug_assert_eq!(out.len(), expected_byte_size);
142
143 // TODO optimize for when all channels are f16!
144 // we should be able to omit endianness conversions in that case
145 // see https://github.com/AcademySoftwareFoundation/openexr/blob/3bd93f85bcb74c77255f28cdbb913fdbfbb39dfe/OpenEXR/IlmImf/ImfTiledOutputFile.cpp#L750-L842
146 Ok(super::convert_little_endian_to_current(out, channels, rectangle))
147}
148
149
150
151pub fn compress(
152 channels: &ChannelList,
153 uncompressed: ByteVec,
154 rectangle: IntegerBounds
155) -> Result<ByteVec>
156{
157 if uncompressed.is_empty() {
158 return Ok(Vec::new());
159 }
160
161 // TODO do not convert endianness for f16-only images
162 // see https://github.com/AcademySoftwareFoundation/openexr/blob/3bd93f85bcb74c77255f28cdbb913fdbfbb39dfe/OpenEXR/IlmImf/ImfTiledOutputFile.cpp#L750-L842
163 let uncompressed = super::convert_current_to_little_endian(uncompressed, channels, rectangle);
164 let uncompressed = uncompressed.as_slice();// TODO no alloc
165
166 let mut tmp = vec![0_u16; uncompressed.len() / 2 ];
167 let mut channel_data: SmallVec<[ChannelData; 6]> = {
168 let mut tmp_end_index = 0;
169
170 let vec = channels.list.iter().map(|channel| {
171 let number_samples = channel.subsampled_resolution(rectangle.size);
172 let byte_size = channel.sample_type.bytes_per_sample() / SampleType::F16.bytes_per_sample();
173 let byte_count = byte_size * number_samples.area();
174
175 let channel = ChannelData {
176 tmp_end_index,
177 tmp_start_index: tmp_end_index,
178 y_sampling: channel.sampling.y(),
179 resolution: number_samples,
180 samples_per_pixel: byte_size,
181 };
182
183 tmp_end_index += byte_count;
184 channel
185 }).collect();
186
187 debug_assert_eq!(tmp_end_index, tmp.len());
188 vec
189 };
190
191 let mut remaining_uncompressed_bytes = uncompressed;
192 for y in rectangle.position.y() .. rectangle.end().y() {
193 for channel in &mut channel_data {
194 if mod_p(y, usize_to_i32(channel.y_sampling)) != 0 { continue; }
195 let u16s_per_line = channel.resolution.x() * channel.samples_per_pixel;
196 let next_tmp_end_index = channel.tmp_end_index + u16s_per_line;
197 let target = &mut tmp[channel.tmp_end_index .. next_tmp_end_index];
198 channel.tmp_end_index = next_tmp_end_index;
199
200 // TODO do not convert endianness for f16-only images
201 // see https://github.com/AcademySoftwareFoundation/openexr/blob/3bd93f85bcb74c77255f28cdbb913fdbfbb39dfe/OpenEXR/IlmImf/ImfTiledOutputFile.cpp#L750-L842
202 // We can support uncompressed data in the machine's native format
203 // if all image channels are of type HALF, and if the Xdr and the
204 // native representations of a half have the same size.
205 u16::read_slice(&mut remaining_uncompressed_bytes, target).expect("in-memory read failed");
206 }
207 }
208
209
210 let (min_non_zero, max_non_zero, bitmap) = bitmap_from_data(&tmp);
211 let (max_value, table) = forward_lookup_table_from_bitmap(&bitmap);
212 apply_lookup_table(&mut tmp, &table);
213
214 let mut piz_compressed = Vec::with_capacity(uncompressed.len() / 2);
215 u16::try_from(min_non_zero)?.write(&mut piz_compressed)?;
216 u16::try_from(max_non_zero)?.write(&mut piz_compressed)?;
217
218 if min_non_zero <= max_non_zero {
219 piz_compressed.extend_from_slice(&bitmap[min_non_zero ..= max_non_zero]);
220 }
221
222 for channel in channel_data {
223 for offset in 0 .. channel.samples_per_pixel {
224 wavelet::encode(
225 &mut tmp[channel.tmp_start_index + offset .. channel.tmp_end_index],
226 channel.resolution,
227 Vec2(channel.samples_per_pixel, channel.resolution.x() * channel.samples_per_pixel),
228 max_value
229 )?;
230 }
231 }
232
233 let huffman_compressed: Vec<u8> = huffman::compress(&tmp)?;
234 u8::write_i32_sized_slice(&mut piz_compressed, &huffman_compressed).expect("in-memory write failed");
235
236 Ok(piz_compressed)
237}
238
239
240pub fn bitmap_from_data(data: &[u16]) -> (usize, usize, Vec<u8>) {
241 let mut bitmap: Vec = vec![0_u8; BITMAP_SIZE];
242
243 for value: &u16 in data {
244 bitmap[*value as usize >> 3] |= 1 << (*value as u8 & 7);
245 }
246
247 bitmap[0] = bitmap[0] & !1; // zero is not explicitly stored in the bitmap; we assume that the data always contain zeroes
248
249 let min_index: Option = bitmap.iter().position(|&value: u8| value != 0);
250 let max_index: Option = min_index.map(|min: usize| // only if min was found
251 min + bitmap[min..].iter().rposition(|&value| value != 0).expect(msg:"[min] not found")
252 );
253
254 (min_index.unwrap_or(default:0), max_index.unwrap_or(default:0), bitmap)
255}
256
257pub fn forward_lookup_table_from_bitmap(bitmap: &[u8]) -> (u16, Vec<u16>) {
258 debug_assert_eq!(bitmap.len(), BITMAP_SIZE);
259
260 let mut table: Vec = vec![0_u16; U16_RANGE];
261 let mut count: usize = 0_usize;
262
263 for (index: usize, entry: &mut u16) in table.iter_mut().enumerate() {
264 if index == 0 || bitmap[index >> 3] as usize & (1 << (index & 7)) != 0 {
265 *entry = usize_to_u16(count).unwrap();
266 count += 1;
267 }
268 }
269
270 (usize_to_u16(count - 1).unwrap(), table)
271}
272
273fn reverse_lookup_table_from_bitmap(bitmap: Bytes<'_>) -> (Vec<u16>, u16) {
274 let mut table: Vec = Vec::with_capacity(U16_RANGE);
275
276 for index: usize in 0 .. U16_RANGE { // cannot use iter because filter removes capacity sizehint
277 if index == 0 || ((bitmap[index >> 3] as usize & (1 << (index & 7))) != 0) {
278 table.push(usize_to_u16(index).unwrap());
279 }
280 }
281
282 debug_assert!(!table.is_empty());
283 let max_value: u16 = usize_to_u16(table.len() - 1).unwrap();
284
285 // fill remaining up to u16 range
286 assert!(table.len() <= U16_RANGE);
287 table.resize(U16_RANGE, value:0);
288
289 (table, max_value)
290}
291
292fn apply_lookup_table(data: &mut [u16], table: &[u16]) {
293 for data: &mut u16 in data {
294 *data = table[*data as usize];
295 }
296}
297
298#[cfg(test)]
299mod test {
300 use crate::prelude::*;
301 use crate::compression::ByteVec;
302 use crate::compression::piz;
303 use crate::meta::attribute::*;
304
305 fn test_roundtrip_noise_with(channels: ChannelList, rectangle: IntegerBounds){
306 let pixel_bytes: ByteVec = (0 .. 37).map(|_| rand::random()).collect::<Vec<u8>>().into_iter()
307 .cycle().take(channels.bytes_per_pixel * rectangle.size.area())
308 .collect();
309
310 let compressed = piz::compress(&channels, pixel_bytes.clone(), rectangle).unwrap();
311 let decompressed = piz::decompress(&channels, compressed, rectangle, pixel_bytes.len(), true).unwrap();
312
313 assert_eq!(pixel_bytes, decompressed);
314 }
315
316
317 #[test]
318 fn roundtrip_any_sample_type(){
319 for &sample_type in &[SampleType::F16, SampleType::F32, SampleType::U32] {
320 let channel = ChannelDescription {
321 sample_type,
322
323 name: Default::default(),
324 quantize_linearly: false,
325 sampling: Vec2(1,1)
326 };
327
328 let channels = ChannelList::new(smallvec![ channel.clone(), channel ]);
329
330 let rectangle = IntegerBounds {
331 position: Vec2(-30, 100),
332 size: Vec2(1080, 720),
333 };
334
335 test_roundtrip_noise_with(channels, rectangle);
336 }
337 }
338
339 #[test]
340 fn roundtrip_two_channels(){
341 let channel = ChannelDescription {
342 sample_type: SampleType::F16,
343
344 name: Default::default(),
345 quantize_linearly: false,
346 sampling: Vec2(1,1)
347 };
348
349 let channel2 = ChannelDescription {
350 sample_type: SampleType::F32,
351
352 name: Default::default(),
353 quantize_linearly: false,
354 sampling: Vec2(1,1)
355 };
356
357 let channels = ChannelList::new(smallvec![ channel, channel2 ]);
358
359 let rectangle = IntegerBounds {
360 position: Vec2(-3, 1),
361 size: Vec2(223, 3132),
362 };
363
364 test_roundtrip_noise_with(channels, rectangle);
365 }
366
367
368
369 #[test]
370 fn roundtrip_seven_channels(){
371 let channels = ChannelList::new(smallvec![
372 ChannelDescription {
373 sample_type: SampleType::F32,
374
375 name: Default::default(),
376 quantize_linearly: false,
377 sampling: Vec2(1,1)
378 },
379
380 ChannelDescription {
381 sample_type: SampleType::F32,
382
383 name: Default::default(),
384 quantize_linearly: false,
385 sampling: Vec2(1,1)
386 },
387
388 ChannelDescription {
389 sample_type: SampleType::F32,
390
391 name: Default::default(),
392 quantize_linearly: false,
393 sampling: Vec2(1,1)
394 },
395
396 ChannelDescription {
397 sample_type: SampleType::F16,
398
399 name: Default::default(),
400 quantize_linearly: false,
401 sampling: Vec2(1,1)
402 },
403
404 ChannelDescription {
405 sample_type: SampleType::F32,
406
407 name: Default::default(),
408 quantize_linearly: false,
409 sampling: Vec2(1,1)
410 },
411
412 ChannelDescription {
413 sample_type: SampleType::F32,
414
415 name: Default::default(),
416 quantize_linearly: false,
417 sampling: Vec2(1,1)
418 },
419
420 ChannelDescription {
421 sample_type: SampleType::U32,
422
423 name: Default::default(),
424 quantize_linearly: false,
425 sampling: Vec2(1,1)
426 },
427 ]);
428
429 let rectangle = IntegerBounds {
430 position: Vec2(-3, 1),
431 size: Vec2(1323, 132),
432 };
433
434 test_roundtrip_noise_with(channels, rectangle);
435 }
436
437}