From 54bd86f7ee78dadae8e9c80a60696b8d4053aee9 Mon Sep 17 00:00:00 2001 From: "Micah Chambers (eos)" Date: Tue, 30 Jul 2024 21:50:14 -0700 Subject: [PATCH] add fp16 cleanup fix predict_f16 which had the wrong word size fix missing prediction add example move ug --- Cargo.toml | 1 + src/bytecast.rs | 3 +++ src/decoder/image.rs | 3 ++- src/decoder/mod.rs | 30 ++++++++++++++++++++++++++++++ tests/decode_fp16_images.rs | 34 ++++++++++++++++++++++++++++++++++ 5 files changed, 70 insertions(+), 1 deletion(-) create mode 100644 tests/decode_fp16_images.rs diff --git a/Cargo.toml b/Cargo.toml index a71fd27f..51c203c6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,7 @@ categories = ["multimedia::images", "multimedia::encoding"] exclude = ["tests/images/*", "tests/fuzz_images/*"] [dependencies] +half = { version = "2.4.1" } weezl = "0.1.0" jpeg = { package = "jpeg-decoder", version = "0.3.0", default-features = false } flate2 = "1.0.20" diff --git a/src/bytecast.rs b/src/bytecast.rs index 6e9d762a..88d562c7 100644 --- a/src/bytecast.rs +++ b/src/bytecast.rs @@ -12,6 +12,8 @@ //! TODO: Would like to use std-lib here. use std::{mem, slice}; +use half::f16; + macro_rules! integral_slice_as_bytes{($int:ty, $const:ident $(,$mut:ident)*) => { pub(crate) fn $const(slice: &[$int]) -> &[u8] { assert!(mem::align_of::<$int>() <= mem::size_of::<$int>()); @@ -31,4 +33,5 @@ integral_slice_as_bytes!(i32, i32_as_ne_bytes, i32_as_ne_mut_bytes); integral_slice_as_bytes!(u64, u64_as_ne_bytes, u64_as_ne_mut_bytes); integral_slice_as_bytes!(i64, i64_as_ne_bytes, i64_as_ne_mut_bytes); integral_slice_as_bytes!(f32, f32_as_ne_bytes, f32_as_ne_mut_bytes); +integral_slice_as_bytes!(f16, f16_as_ne_bytes, f16_as_ne_mut_bytes); integral_slice_as_bytes!(f64, f64_as_ne_bytes, f64_as_ne_mut_bytes); diff --git a/src/decoder/image.rs b/src/decoder/image.rs index 532a8fb9..b9e9976f 100644 --- a/src/decoder/image.rs +++ b/src/decoder/image.rs @@ -1,7 +1,7 @@ use super::ifd::{Directory, Value}; use super::stream::{ByteOrder, DeflateReader, LZWReader, PackBitsReader}; use super::tag_reader::TagReader; -use super::{predict_f32, predict_f64, Limits}; +use super::{predict_f16, predict_f32, predict_f64, Limits}; use super::{stream::SmartReader, ChunkType}; use crate::tags::{ CompressionMethod, PhotometricInterpretation, PlanarConfiguration, Predictor, SampleFormat, Tag, @@ -674,6 +674,7 @@ impl Image { let row = &mut row[..data_row_bytes]; match color_type.bit_depth() { + 16 => predict_f16(&mut encoded, row, samples), 32 => predict_f32(&mut encoded, row, samples), 64 => predict_f64(&mut encoded, row, samples), _ => unreachable!(), diff --git a/src/decoder/mod.rs b/src/decoder/mod.rs index 2b8b0dd2..9cb92a43 100644 --- a/src/decoder/mod.rs +++ b/src/decoder/mod.rs @@ -8,6 +8,7 @@ use crate::tags::{ use crate::{ bytecast, ColorType, TiffError, TiffFormatError, TiffResult, TiffUnsupportedError, UsageError, }; +use half::f16; use self::ifd::Directory; use self::image::Image; @@ -29,6 +30,8 @@ pub enum DecodingResult { U32(Vec), /// A vector of 64 bit unsigned ints U64(Vec), + /// A vector of 16 bit IEEE floats (held in u16) + F16(Vec), /// A vector of 32 bit IEEE floats F32(Vec), /// A vector of 64 bit IEEE floats @@ -92,6 +95,14 @@ impl DecodingResult { } } + fn new_f16(size: usize, limits: &Limits) -> TiffResult { + if size > limits.decoding_buffer_size / std::mem::size_of::() { + Err(TiffError::LimitsExceeded) + } else { + Ok(DecodingResult::F16(vec![f16::ZERO; size])) + } + } + fn new_i8(size: usize, limits: &Limits) -> TiffResult { if size > limits.decoding_buffer_size / std::mem::size_of::() { Err(TiffError::LimitsExceeded) @@ -130,6 +141,7 @@ impl DecodingResult { DecodingResult::U16(ref mut buf) => DecodingBuffer::U16(&mut buf[start..]), DecodingResult::U32(ref mut buf) => DecodingBuffer::U32(&mut buf[start..]), DecodingResult::U64(ref mut buf) => DecodingBuffer::U64(&mut buf[start..]), + DecodingResult::F16(ref mut buf) => DecodingBuffer::F16(&mut buf[start..]), DecodingResult::F32(ref mut buf) => DecodingBuffer::F32(&mut buf[start..]), DecodingResult::F64(ref mut buf) => DecodingBuffer::F64(&mut buf[start..]), DecodingResult::I8(ref mut buf) => DecodingBuffer::I8(&mut buf[start..]), @@ -150,6 +162,8 @@ pub enum DecodingBuffer<'a> { U32(&'a mut [u32]), /// A slice of 64 bit unsigned ints U64(&'a mut [u64]), + /// A slice of 16 bit IEEE floats + F16(&'a mut [f16]), /// A slice of 32 bit IEEE floats F32(&'a mut [f32]), /// A slice of 64 bit IEEE floats @@ -175,6 +189,7 @@ impl<'a> DecodingBuffer<'a> { DecodingBuffer::I32(buf) => bytecast::i32_as_ne_mut_bytes(buf), DecodingBuffer::U64(buf) => bytecast::u64_as_ne_mut_bytes(buf), DecodingBuffer::I64(buf) => bytecast::i64_as_ne_mut_bytes(buf), + DecodingBuffer::F16(buf) => bytecast::f16_as_ne_mut_bytes(buf), DecodingBuffer::F32(buf) => bytecast::f32_as_ne_mut_bytes(buf), DecodingBuffer::F64(buf) => bytecast::f64_as_ne_mut_bytes(buf), } @@ -303,6 +318,19 @@ fn predict_f32(input: &mut [u8], output: &mut [u8], samples: usize) { } } +fn predict_f16(input: &mut [u8], output: &mut [u8], samples: usize) { + for i in samples..input.len() { + input[i] = input[i].wrapping_add(input[i - samples]); + } + + for (i, chunk) in output.chunks_mut(2).enumerate() { + chunk.copy_from_slice(&u16::to_ne_bytes(u16::from_be_bytes([ + input[i], + input[input.len() / 2 + i], + ]))); + } +} + fn predict_f64(input: &mut [u8], output: &mut [u8], samples: usize) { for i in samples..input.len() { input[i] = input[i].wrapping_add(input[i - samples]); @@ -340,6 +368,7 @@ fn fix_endianness_and_predict( Predictor::FloatingPoint => { let mut buffer_copy = buf.to_vec(); match bit_depth { + 16 => predict_f16(&mut buffer_copy, buf, samples), 32 => predict_f32(&mut buffer_copy, buf, samples), 64 => predict_f64(&mut buffer_copy, buf, samples), _ => unreachable!("Caller should have validated arguments. Please file a bug."), @@ -1004,6 +1033,7 @@ impl Decoder { )), }, SampleFormat::IEEEFP => match max_sample_bits { + 16 => DecodingResult::new_f16(buffer_size, &self.limits), 32 => DecodingResult::new_f32(buffer_size, &self.limits), 64 => DecodingResult::new_f64(buffer_size, &self.limits), n => Err(TiffError::UnsupportedError( diff --git a/tests/decode_fp16_images.rs b/tests/decode_fp16_images.rs new file mode 100644 index 00000000..88c7655a --- /dev/null +++ b/tests/decode_fp16_images.rs @@ -0,0 +1,34 @@ +extern crate tiff; + +use tiff::decoder::{ifd, Decoder, DecodingResult}; +use tiff::tags::Tag; +use tiff::ColorType; + +use std::fs::File; +use std::path::PathBuf; + +const TEST_IMAGE_DIR: &str = "./tests/images/"; + +#[test] +fn test_ieee_fp16() { + let filenames = ["example_fp16.tif"]; + + for filename in filenames.iter() { + let path = PathBuf::from(TEST_IMAGE_DIR).join(filename); + let img_file = File::open(path).expect("Cannot find test image!"); + let mut decoder = Decoder::new(img_file).expect("Cannot create decoder"); + assert_eq!( + decoder.dimensions().expect("Cannot get dimensions"), + (513, 513) + ); + assert_eq!( + decoder.colortype().expect("Cannot get colortype"), + ColorType::Gray(16) + ); + if let DecodingResult::F16(img_res) = decoder.read_image().unwrap() { + //assert_eq!(image_data, img_res); + } else { + panic!("Wrong data type"); + } + } +}