Skip to content

Commit 4b3f2b2

Browse files
kylebarronkylebarron
authored
Decoder registry (#20)
1 parent 97c4e02 commit 4b3f2b2

File tree

8 files changed

+227
-39
lines changed

8 files changed

+227
-39
lines changed

‎python/python/async_tiff/_decoder.pyi‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
from typing import Protocol
2+
from collections.abc import Buffer
3+
4+
from .enums import CompressionMethod
5+
6+
class Decoder(Protocol):
7+
@staticmethod
8+
def __call__(buffer: Buffer) -> Buffer: ...
9+
10+
class DecoderRegistry:
11+
def __init__(self) -> None: ...
12+
def add(self, compression: CompressionMethod | int, decoder: Decoder) -> None: ...

‎python/src/decoder.rs‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
use async_tiff::decoder::{Decoder, DecoderRegistry};
2+
use async_tiff::error::AiocogeoError;
3+
use bytes::Bytes;
4+
use pyo3::exceptions::PyTypeError;
5+
use pyo3::intern;
6+
use pyo3::prelude::*;
7+
use pyo3::types::{PyDict, PyTuple};
8+
use pyo3_bytes::PyBytes;
9+
10+
use crate::enums::PyCompressionMethod;
11+
12+
#[pyclass(name = "DecoderRegistry")]
13+
pub(crate) struct PyDecoderRegistry(DecoderRegistry);
14+
15+
#[pymethods]
16+
impl PyDecoderRegistry {
17+
#[new]
18+
fn new() -> Self {
19+
Self(DecoderRegistry::default())
20+
}
21+
22+
fn add(&mut self, compression: PyCompressionMethod, decoder: PyDecoder) {
23+
self.0
24+
.as_mut()
25+
.insert(compression.into(), Box::new(decoder));
26+
}
27+
}
28+
29+
#[derive(Debug)]
30+
struct PyDecoder(PyObject);
31+
32+
impl PyDecoder {
33+
fn call(&self, py: Python, buffer: Bytes) -> PyResult<PyBytes> {
34+
let kwargs = PyDict::new(py);
35+
kwargs.set_item(intern!(py, "buffer"), PyBytes::new(buffer))?;
36+
let result = self.0.call(py, PyTuple::empty(py), Some(&kwargs))?;
37+
result.extract(py)
38+
}
39+
}
40+
41+
impl<'py> FromPyObject<'py> for PyDecoder {
42+
fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
43+
if !ob.hasattr(intern!(ob.py(), "__call__"))? {
44+
return Err(PyTypeError::new_err(
45+
"Expected callable object for custom decoder.",
46+
));
47+
}
48+
Ok(Self(ob.clone().unbind()))
49+
}
50+
}
51+
52+
impl Decoder for PyDecoder {
53+
fn decode_tile(
54+
&self,
55+
buffer: bytes::Bytes,
56+
_photometric_interpretation: tiff::tags::PhotometricInterpretation,
57+
_jpeg_tables: Option<&[u8]>,
58+
) -> async_tiff::error::Result<bytes::Bytes> {
59+
let decoded_buffer = Python::with_gil(|py| self.call(py, buffer))
60+
.map_err(|err| AiocogeoError::General(err.to_string()))?;
61+
Ok(decoded_buffer.into_inner())
62+
}
63+
}

‎python/src/enums.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,18 @@ impl From<CompressionMethod> for PyCompressionMethod {
1414
}
1515
}
1616

17+
impl From<PyCompressionMethod> for CompressionMethod {
18+
fn from(value: PyCompressionMethod) -> Self {
19+
value.0
20+
}
21+
}
22+
23+
impl<'py> FromPyObject<'py> for PyCompressionMethod {
24+
fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
25+
Ok(Self(CompressionMethod::from_u16_exhaustive(ob.extract()?)))
26+
}
27+
}
28+
1729
impl<'py> IntoPyObject<'py> for PyCompressionMethod {
1830
type Target = PyAny;
1931
type Output = Bound<'py, PyAny>;

‎python/src/lib.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
#![deny(clippy::undocumented_unsafe_blocks)]
22

3+
mod decoder;
34
mod enums;
45
mod geo;
56
mod ifd;
67
mod tiff;
78

89
use pyo3::prelude::*;
910

11+
use crate::decoder::PyDecoderRegistry;
1012
use crate::geo::PyGeoKeyDirectory;
1113
use crate::ifd::PyImageFileDirectory;
1214
use crate::tiff::PyTIFF;
@@ -43,6 +45,7 @@ fn _async_tiff(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
4345
check_debug_build(py)?;
4446

4547
m.add_wrapped(wrap_pyfunction!(___version))?;
48+
m.add_class::<PyDecoderRegistry>()?;
4649
m.add_class::<PyGeoKeyDirectory>()?;
4750
m.add_class::<PyImageFileDirectory>()?;
4851
m.add_class::<PyTIFF>()?;

‎src/cog.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ mod test {
4949
use std::io::BufReader;
5050
use std::sync::Arc;
5151

52+
use crate::decoder::DecoderRegistry;
5253
use crate::ObjectReader;
5354

5455
use super::*;
@@ -66,7 +67,11 @@ mod test {
6667
let cog_reader = COGReader::try_open(Box::new(reader.clone())).await.unwrap();
6768

6869
let ifd = &cog_reader.ifds.as_ref()[1];
69-
let tile = ifd.get_tile(0, 0, Box::new(reader)).await.unwrap();
70+
let decoder_registry = DecoderRegistry::default();
71+
let tile = ifd
72+
.get_tile(0, 0, Box::new(reader), &decoder_registry)
73+
.await
74+
.unwrap();
7075
std::fs::write("img.buf", tile).unwrap();
7176
}
7277

‎src/decoder.rs‎

Lines changed: 123 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
use std::collections::HashMap;
2+
use std::fmt::Debug;
13
use std::io::{Cursor, Read};
24

35
use bytes::Bytes;
@@ -7,47 +9,138 @@ use tiff::{TiffError, TiffUnsupportedError};
79

810
use crate::error::Result;
911

12+
/// A registry of decoders.
13+
#[derive(Debug)]
14+
pub struct DecoderRegistry(HashMap<CompressionMethod, Box<dyn Decoder>>);
15+
16+
impl DecoderRegistry {
17+
/// Create a new decoder registry with no decoders registered
18+
pub fn new() -> Self {
19+
Self(HashMap::new())
20+
}
21+
}
22+
23+
impl AsRef<HashMap<CompressionMethod, Box<dyn Decoder>>> for DecoderRegistry {
24+
fn as_ref(&self) -> &HashMap<CompressionMethod, Box<dyn Decoder>> {
25+
&self.0
26+
}
27+
}
28+
29+
impl AsMut<HashMap<CompressionMethod, Box<dyn Decoder>>> for DecoderRegistry {
30+
fn as_mut(&mut self) -> &mut HashMap<CompressionMethod, Box<dyn Decoder>> {
31+
&mut self.0
32+
}
33+
}
34+
35+
impl Default for DecoderRegistry {
36+
fn default() -> Self {
37+
let mut registry = HashMap::with_capacity(5);
38+
registry.insert(CompressionMethod::None, Box::new(UncompressedDecoder) as _);
39+
registry.insert(CompressionMethod::Deflate, Box::new(DeflateDecoder) as _);
40+
registry.insert(CompressionMethod::OldDeflate, Box::new(DeflateDecoder) as _);
41+
registry.insert(CompressionMethod::LZW, Box::new(LZWDecoder) as _);
42+
registry.insert(CompressionMethod::ModernJPEG, Box::new(JPEGDecoder) as _);
43+
Self(registry)
44+
}
45+
}
46+
47+
/// A trait to decode a TIFF tile.
48+
pub trait Decoder: Debug + Send + Sync {
49+
fn decode_tile(
50+
&self,
51+
buffer: Bytes,
52+
photometric_interpretation: PhotometricInterpretation,
53+
jpeg_tables: Option<&[u8]>,
54+
) -> Result<Bytes>;
55+
}
56+
57+
#[derive(Debug, Clone)]
58+
pub struct DeflateDecoder;
59+
60+
impl Decoder for DeflateDecoder {
61+
fn decode_tile(
62+
&self,
63+
buffer: Bytes,
64+
_photometric_interpretation: PhotometricInterpretation,
65+
_jpeg_tables: Option<&[u8]>,
66+
) -> Result<Bytes> {
67+
let mut decoder = ZlibDecoder::new(Cursor::new(buffer));
68+
let mut buf = Vec::new();
69+
decoder.read_to_end(&mut buf)?;
70+
Ok(buf.into())
71+
}
72+
}
73+
74+
#[derive(Debug, Clone)]
75+
pub struct JPEGDecoder;
76+
77+
impl Decoder for JPEGDecoder {
78+
fn decode_tile(
79+
&self,
80+
buffer: Bytes,
81+
photometric_interpretation: PhotometricInterpretation,
82+
jpeg_tables: Option<&[u8]>,
83+
) -> Result<Bytes> {
84+
decode_modern_jpeg(buffer, photometric_interpretation, jpeg_tables)
85+
}
86+
}
87+
88+
#[derive(Debug, Clone)]
89+
pub struct LZWDecoder;
90+
91+
impl Decoder for LZWDecoder {
92+
fn decode_tile(
93+
&self,
94+
buffer: Bytes,
95+
_photometric_interpretation: PhotometricInterpretation,
96+
_jpeg_tables: Option<&[u8]>,
97+
) -> Result<Bytes> {
98+
// https://github.com/image-rs/image-tiff/blob/90ae5b8e54356a35e266fb24e969aafbcb26e990/src/decoder/stream.rs#L147
99+
let mut decoder = weezl::decode::Decoder::with_tiff_size_switch(weezl::BitOrder::Msb, 8);
100+
let decoded = decoder.decode(&buffer).expect("failed to decode LZW data");
101+
Ok(decoded.into())
102+
}
103+
}
104+
105+
#[derive(Debug, Clone)]
106+
pub struct UncompressedDecoder;
107+
108+
impl Decoder for UncompressedDecoder {
109+
fn decode_tile(
110+
&self,
111+
buffer: Bytes,
112+
_photometric_interpretation: PhotometricInterpretation,
113+
_jpeg_tables: Option<&[u8]>,
114+
) -> Result<Bytes> {
115+
Ok(buffer)
116+
}
117+
}
118+
10119
// https://github.com/image-rs/image-tiff/blob/3bfb43e83e31b0da476832067ada68a82b378b7b/src/decoder/image.rs#L370
11120
pub(crate) fn decode_tile(
12121
buf: Bytes,
13122
photometric_interpretation: PhotometricInterpretation,
14123
compression_method: CompressionMethod,
15124
// compressed_length: u64,
16-
jpeg_tables: Option<&Vec<u8>>,
125+
jpeg_tables: Option<&[u8]>,
126+
decoder_registry: &DecoderRegistry,
17127
) -> Result<Bytes> {
18-
match compression_method {
19-
CompressionMethod::None => Ok(buf),
20-
CompressionMethod::LZW => decode_lzw(buf),
21-
CompressionMethod::Deflate | CompressionMethod::OldDeflate => decode_deflate(buf),
22-
CompressionMethod::ModernJPEG => {
23-
decode_modern_jpeg(buf, photometric_interpretation, jpeg_tables)
24-
}
25-
method => Err(TiffError::UnsupportedError(
26-
TiffUnsupportedError::UnsupportedCompressionMethod(method),
27-
)
28-
.into()),
29-
}
30-
}
31-
32-
fn decode_lzw(buf: Bytes) -> Result<Bytes> {
33-
// https://github.com/image-rs/image-tiff/blob/90ae5b8e54356a35e266fb24e969aafbcb26e990/src/decoder/stream.rs#L147
34-
let mut decoder = weezl::decode::Decoder::with_tiff_size_switch(weezl::BitOrder::Msb, 8);
35-
let decoded = decoder.decode(&buf).expect("failed to decode LZW data");
36-
Ok(decoded.into())
37-
}
128+
let decoder =
129+
decoder_registry
130+
.0
131+
.get(&compression_method)
132+
.ok_or(TiffError::UnsupportedError(
133+
TiffUnsupportedError::UnsupportedCompressionMethod(compression_method),
134+
))?;
38135

39-
fn decode_deflate(buf: Bytes) -> Result<Bytes> {
40-
let mut decoder = ZlibDecoder::new(Cursor::new(buf));
41-
let mut buf = Vec::new();
42-
decoder.read_to_end(&mut buf)?;
43-
Ok(buf.into())
136+
decoder.decode_tile(buf, photometric_interpretation, jpeg_tables)
44137
}
45138

46139
// https://github.com/image-rs/image-tiff/blob/3bfb43e83e31b0da476832067ada68a82b378b7b/src/decoder/image.rs#L389-L450
47140
fn decode_modern_jpeg(
48141
buf: Bytes,
49142
photometric_interpretation: PhotometricInterpretation,
50-
jpeg_tables: Option<&Vec<u8>>,
143+
jpeg_tables: Option<&[u8]>,
51144
) -> Result<Bytes> {
52145
// Construct new jpeg_reader wrapping a SmartReader.
53146
//
@@ -76,13 +169,9 @@ fn decode_modern_jpeg(
76169

77170
match photometric_interpretation {
78171
PhotometricInterpretation::RGB => decoder.set_color_transform(jpeg::ColorTransform::RGB),
79-
PhotometricInterpretation::WhiteIsZero => {
80-
decoder.set_color_transform(jpeg::ColorTransform::None)
81-
}
82-
PhotometricInterpretation::BlackIsZero => {
83-
decoder.set_color_transform(jpeg::ColorTransform::None)
84-
}
85-
PhotometricInterpretation::TransparencyMask => {
172+
PhotometricInterpretation::WhiteIsZero
173+
| PhotometricInterpretation::BlackIsZero
174+
| PhotometricInterpretation::TransparencyMask => {
86175
decoder.set_color_transform(jpeg::ColorTransform::None)
87176
}
88177
PhotometricInterpretation::CMYK => decoder.set_color_transform(jpeg::ColorTransform::CMYK),

‎src/ifd.rs‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ use tiff::tags::{
1313
use tiff::TiffError;
1414

1515
use crate::async_reader::AsyncCursor;
16-
use crate::decoder::decode_tile;
16+
use crate::decoder::{decode_tile, DecoderRegistry};
1717
use crate::error::{AiocogeoError, Result};
1818
use crate::geo::{AffineTransform, GeoKeyDirectory, GeoKeyTag};
1919
use crate::AsyncFileReader;
@@ -681,14 +681,16 @@ impl ImageFileDirectory {
681681
x: usize,
682682
y: usize,
683683
mut reader: Box<dyn AsyncFileReader>,
684+
decoder_registry: &DecoderRegistry,
684685
) -> Result<Bytes> {
685686
let range = self.get_tile_byte_range(x, y);
686687
let buf = reader.get_bytes(range).await?;
687688
decode_tile(
688689
buf,
689690
self.photometric_interpretation,
690691
self.compression,
691-
self.jpeg_tables.as_ref(),
692+
self.jpeg_tables.as_deref(),
693+
decoder_registry,
692694
)
693695
}
694696

@@ -697,6 +699,7 @@ impl ImageFileDirectory {
697699
x: &[usize],
698700
y: &[usize],
699701
mut reader: Box<dyn AsyncFileReader>,
702+
decoder_registry: &DecoderRegistry,
700703
) -> Result<Vec<Bytes>> {
701704
assert_eq!(x.len(), y.len(), "x and y should have same len");
702705

@@ -717,7 +720,8 @@ impl ImageFileDirectory {
717720
buf,
718721
self.photometric_interpretation,
719722
self.compression,
720-
self.jpeg_tables.as_ref(),
723+
self.jpeg_tables.as_deref(),
724+
decoder_registry,
721725
)?;
722726
decoded_tiles.push(decoded);
723727
}

‎src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
mod async_reader;
44
mod cog;
5-
mod decoder;
5+
pub mod decoder;
66
pub mod error;
77
pub mod geo;
88
mod ifd;

0 commit comments

Comments
 (0)