Progress on image tile decoding (#2)

* Progress on image tile decoding

* Update decoder code

* update gitignore

* Use Bytes in API instead of Vec<u8>

* rename

Author: kylebarron

.gitignore

@@ -1,4 +1,5 @@
 *.tif
+*.buf
 
 # Generated by Cargo
 # will have compiled files and executables

Cargo.toml

@@ -6,11 +6,14 @@ edition = "2021"
 [dependencies]
 byteorder = "1"
 bytes = "1.7.0"
+flate2 = "1.0.20"
+jpeg = { package = "jpeg-decoder", version = "0.3.0", default-features = false }
 ndarray = "*"
 num_enum = "*"
 object_store = "0.11"
 thiserror = "1"
 tiff = "0.9"
+weezl = "0.1.0"
 
 [dev-dependencies]
 tokio = { version = "1.9", features = ["macros", "fs", "rt-multi-thread"] }

show_image.py

@@ -0,0 +1,23 @@
+import numpy as np
+from rasterio.plot import show, reshape_as_raster
+
+with open("img.buf", "rb") as f:
+    buffer = f.read()
+
+arr = np.frombuffer(buffer, np.uint8)
+arr = arr.reshape(512, 512, 3)
+arr = reshape_as_raster(arr)
+show(arr, adjust=True)
+
+###########
+
+with open("img_from_tiff.buf", "rb") as f:
+    buffer = f.read()
+
+arr = np.frombuffer(buffer, np.uint8)
+arr.reshape()
+
+# We first need to reshape into "image" axes, then we need to reshape as "raster" axes.
+arr = arr.reshape(12410, 9680, 3)
+arr2 = reshape_as_raster(arr)
+show(arr2, adjust=True)

src/cog.rs

@@ -59,14 +59,40 @@ impl COGReader {
 
 #[cfg(test)]
 mod test {
+    use std::io::BufReader;
+
     use super::*;
     use object_store::local::LocalFileSystem;
+    use tiff::decoder::{DecodingResult, Limits};
 
     #[tokio::test]
     async fn tmp() {
         let folder = "/Users/kyle/github/developmentseed/aiocogeo-rs/";
         let path = Path::parse("m_4007307_sw_18_060_20220803.tif").unwrap();
         let store = Arc::new(LocalFileSystem::new_with_prefix(folder).unwrap());
-        let _reader = COGReader::try_open(store, path).await.unwrap();
+        let reader = COGReader::try_open(store.clone(), path.clone())
+            .await
+            .unwrap();
+        let cursor = ObjectStoreCursor::new(store.clone(), path.clone());
+        let ifd = &reader.ifds.as_ref()[4];
+        dbg!(ifd.compression);
+        let tile = ifd.get_tile(0, 0, &cursor).await.unwrap();
+        std::fs::write("img.buf", tile).unwrap();
+        // dbg!(tile.len());
+    }
+
+    #[test]
+    fn tmp_tiff_example() {
+        let path =
+            "/Users/kyle/github/developmentseed/aiocogeo-rs/m_4007307_sw_18_060_20220803.tif";
+        let reader = std::fs::File::open(path).unwrap();
+        let mut decoder = tiff::decoder::Decoder::new(BufReader::new(reader))
+            .unwrap()
+            .with_limits(Limits::unlimited());
+        let result = decoder.read_image().unwrap();
+        match result {
+            DecodingResult::U8(content) => std::fs::write("img_from_tiff.buf", content).unwrap(),
+            _ => todo!(),
+        }
     }
 }

src/compression.rs

@@ -1,4 +1,5 @@
 use std::io::Cursor;
+use std::ops::Range;
 use std::sync::Arc;
 
 use byteorder::{BigEndian, LittleEndian, ReadBytesExt};
@@ -94,6 +95,17 @@ impl ObjectStoreCursor {
         }
     }
 
+    pub(crate) fn store(&self) -> &Arc<dyn ObjectStore> {
+        &self.store
+    }
+
+    pub(crate) async fn get_range(
+        &self,
+        range: Range<usize>,
+    ) -> Result<Bytes, object_store::Error> {
+        Ok(self.store.get_range(&self.path, range).await?)
+    }
+
     /// Advance cursor position by a set amount
     pub(crate) fn advance(&mut self, amount: usize) {
         self.offset += amount;

src/cursor.rs

@@ -0,0 +1,102 @@
+use std::io::{Cursor, Read};
+
+use bytes::Bytes;
+use flate2::bufread::ZlibDecoder;
+use tiff::tags::{CompressionMethod, PhotometricInterpretation};
+use tiff::{TiffError, TiffUnsupportedError};
+
+use crate::error::Result;
+
+// https://github.com/image-rs/image-tiff/blob/3bfb43e83e31b0da476832067ada68a82b378b7b/src/decoder/image.rs#L370
+pub(crate) fn decode_tile(
+    buf: Bytes,
+    photometric_interpretation: PhotometricInterpretation,
+    compression_method: CompressionMethod,
+    // compressed_length: u64,
+    jpeg_tables: Option<&Vec<u8>>,
+) -> Result<Bytes> {
+    match compression_method {
+        CompressionMethod::None => Ok(buf),
+        CompressionMethod::LZW => decode_lzw(buf),
+        CompressionMethod::Deflate | CompressionMethod::OldDeflate => decode_deflate(buf),
+        CompressionMethod::ModernJPEG => {
+            decode_modern_jpeg(buf, photometric_interpretation, jpeg_tables)
+        }
+        method => Err(TiffError::UnsupportedError(
+            TiffUnsupportedError::UnsupportedCompressionMethod(method),
+        )
+        .into()),
+    }
+}
+
+fn decode_lzw(buf: Bytes) -> Result<Bytes> {
+    // https://github.com/image-rs/image-tiff/blob/90ae5b8e54356a35e266fb24e969aafbcb26e990/src/decoder/stream.rs#L147
+    let mut decoder = weezl::decode::Decoder::with_tiff_size_switch(weezl::BitOrder::Msb, 8);
+    let decoded = decoder.decode(&buf).expect("failed to decode LZW data");
+    Ok(decoded.into())
+}
+
+fn decode_deflate(buf: Bytes) -> Result<Bytes> {
+    let mut decoder = ZlibDecoder::new(Cursor::new(buf));
+    let mut buf = Vec::new();
+    decoder.read_to_end(&mut buf)?;
+    Ok(buf.into())
+}
+
+// https://github.com/image-rs/image-tiff/blob/3bfb43e83e31b0da476832067ada68a82b378b7b/src/decoder/image.rs#L389-L450
+fn decode_modern_jpeg(
+    buf: Bytes,
+    photometric_interpretation: PhotometricInterpretation,
+    jpeg_tables: Option<&Vec<u8>>,
+) -> Result<Bytes> {
+    // Construct new jpeg_reader wrapping a SmartReader.
+    //
+    // JPEG compression in TIFF allows saving quantization and/or huffman tables in one central
+    // location. These `jpeg_tables` are simply prepended to the remaining jpeg image data. Because
+    // these `jpeg_tables` start with a `SOI` (HEX: `0xFFD8`) or __start of image__ marker which is
+    // also at the beginning of the remaining JPEG image data and would confuse the JPEG renderer,
+    // one of these has to be taken off. In this case the first two bytes of the remaining JPEG
+    // data is removed because it follows `jpeg_tables`. Similary, `jpeg_tables` ends with a `EOI`
+    // (HEX: `0xFFD9`) or __end of image__ marker, this has to be removed as well (last two bytes
+    // of `jpeg_tables`).
+    let reader = Cursor::new(buf);
+
+    let jpeg_reader = match jpeg_tables {
+        Some(jpeg_tables) => {
+            let mut reader = reader;
+            reader.read_exact(&mut [0; 2])?;
+
+            Box::new(Cursor::new(&jpeg_tables[..jpeg_tables.len() - 2]).chain(reader))
+                as Box<dyn Read>
+        }
+        None => Box::new(reader),
+    };
+
+    let mut decoder = jpeg::Decoder::new(jpeg_reader);
+
+    match photometric_interpretation {
+        PhotometricInterpretation::RGB => decoder.set_color_transform(jpeg::ColorTransform::RGB),
+        PhotometricInterpretation::WhiteIsZero => {
+            decoder.set_color_transform(jpeg::ColorTransform::None)
+        }
+        PhotometricInterpretation::BlackIsZero => {
+            decoder.set_color_transform(jpeg::ColorTransform::None)
+        }
+        PhotometricInterpretation::TransparencyMask => {
+            decoder.set_color_transform(jpeg::ColorTransform::None)
+        }
+        PhotometricInterpretation::CMYK => decoder.set_color_transform(jpeg::ColorTransform::CMYK),
+        PhotometricInterpretation::YCbCr => {
+            decoder.set_color_transform(jpeg::ColorTransform::YCbCr)
+        }
+        photometric_interpretation => {
+            return Err(TiffError::UnsupportedError(
+                TiffUnsupportedError::UnsupportedInterpretation(photometric_interpretation),
+            )
+            .into());
+        }
+    }
+
+    let data = decoder.decode()?;
+    Ok(data.into())
+}

src/decoder.rs

@@ -8,6 +8,18 @@ pub enum AiocogeoError {
     /// General error.
     #[error("General error: {0}")]
     General(String),
+
+    #[error(transparent)]
+    IOError(#[from] std::io::Error),
+
+    #[error(transparent)]
+    JPEGDecodingError(#[from] jpeg::Error),
+
+    #[error(transparent)]
+    ObjectStore(#[from] object_store::Error),
+
+    #[error(transparent)]
+    TIFFError(#[from] tiff::TiffError),
 }
 
 /// Crate-specific result type.