Skip to main content

ad_plugins/
file_tiff.rs

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3
4use ad_core::error::{ADError, ADResult};
5use ad_core::ndarray::{NDArray, NDDataBuffer, NDDataType, NDDimension};
6use ad_core::ndarray_pool::NDArrayPool;
7use ad_core::plugin::file_base::{NDFileMode, NDFileWriter, NDPluginFileBase};
8use ad_core::plugin::runtime::{NDPluginProcess, ProcessResult};
9
10use tiff::encoder::colortype;
11use tiff::encoder::TiffEncoder;
12
13/// TIFF file writer using the `tiff` crate for proper encoding/decoding.
14pub struct TiffWriter {
15    current_path: Option<PathBuf>,
16}
17
18impl TiffWriter {
19    pub fn new() -> Self {
20        Self { current_path: None }
21    }
22}
23
24impl NDFileWriter for TiffWriter {
25    fn open_file(&mut self, path: &Path, _mode: NDFileMode, _array: &NDArray) -> ADResult<()> {
26        self.current_path = Some(path.to_path_buf());
27        Ok(())
28    }
29
30    fn write_file(&mut self, array: &NDArray) -> ADResult<()> {
31        let path = self.current_path.as_ref()
32            .ok_or_else(|| ADError::UnsupportedConversion("no file open".into()))?;
33
34        let info = array.info();
35        let width = info.x_size as u32;
36        let height = info.y_size as u32;
37
38        let file = std::fs::File::create(path)?;
39        let mut encoder = TiffEncoder::new(file)
40            .map_err(|e| ADError::UnsupportedConversion(format!("TIFF encoder error: {}", e)))?;
41
42        match &array.data {
43            NDDataBuffer::U8(v) => {
44                if info.color_size == 3 {
45                    encoder.write_image::<colortype::RGB8>(width, height, v)
46                } else {
47                    encoder.write_image::<colortype::Gray8>(width, height, v)
48                }
49            }
50            NDDataBuffer::U16(v) => {
51                if info.color_size == 3 {
52                    encoder.write_image::<colortype::RGB16>(width, height, v)
53                } else {
54                    encoder.write_image::<colortype::Gray16>(width, height, v)
55                }
56            }
57            NDDataBuffer::I32(v) => {
58                // No signed Gray32 in tiff crate; reinterpret as u32
59                let u32_data: Vec<u32> = v.iter().map(|&x| x as u32).collect();
60                encoder.write_image::<colortype::Gray32>(width, height, &u32_data)
61            }
62            NDDataBuffer::U32(v) => {
63                encoder.write_image::<colortype::Gray32>(width, height, v)
64            }
65            NDDataBuffer::F32(_) => {
66                // Write as raw gray bytes via U8 fallback
67                let raw = array.data.as_u8_slice();
68                let raw_width = (info.x_size * 4) as u32;
69                encoder.write_image::<colortype::Gray8>(raw_width, height, raw)
70            }
71            NDDataBuffer::F64(_) => {
72                let raw = array.data.as_u8_slice();
73                let raw_width = (info.x_size * 8) as u32;
74                encoder.write_image::<colortype::Gray8>(raw_width, height, raw)
75            }
76            _ => {
77                // Fallback: write as raw bytes with Gray8
78                let raw = array.data.as_u8_slice();
79                encoder.write_image::<colortype::Gray8>(
80                    raw.len() as u32 / height.max(1),
81                    height,
82                    raw,
83                )
84            }
85        }
86        .map_err(|e| ADError::UnsupportedConversion(format!("TIFF write error: {}", e)))?;
87
88        Ok(())
89    }
90
91    fn read_file(&mut self) -> ADResult<NDArray> {
92        use tiff::decoder::Decoder;
93
94        let path = self.current_path.as_ref()
95            .ok_or_else(|| ADError::UnsupportedConversion("no file open".into()))?;
96
97        let file = std::fs::File::open(path)?;
98        let mut decoder = Decoder::new(file)
99            .map_err(|e| ADError::UnsupportedConversion(format!("TIFF decode error: {}", e)))?;
100
101        let (width, height) = decoder.dimensions()
102            .map_err(|e| ADError::UnsupportedConversion(format!("TIFF dimensions error: {}", e)))?;
103
104        let result = decoder.read_image()
105            .map_err(|e| ADError::UnsupportedConversion(format!("TIFF read error: {}", e)))?;
106
107        match result {
108            tiff::decoder::DecodingResult::U8(data) => {
109                let dims = vec![
110                    NDDimension::new(width as usize),
111                    NDDimension::new(height as usize),
112                ];
113                let mut arr = NDArray::new(dims, NDDataType::UInt8);
114                arr.data = NDDataBuffer::U8(data);
115                Ok(arr)
116            }
117            tiff::decoder::DecodingResult::U16(data) => {
118                let dims = vec![
119                    NDDimension::new(width as usize),
120                    NDDimension::new(height as usize),
121                ];
122                let mut arr = NDArray::new(dims, NDDataType::UInt16);
123                arr.data = NDDataBuffer::U16(data);
124                Ok(arr)
125            }
126            tiff::decoder::DecodingResult::U32(data) => {
127                let dims = vec![
128                    NDDimension::new(width as usize),
129                    NDDimension::new(height as usize),
130                ];
131                let mut arr = NDArray::new(dims, NDDataType::UInt32);
132                arr.data = NDDataBuffer::U32(data);
133                Ok(arr)
134            }
135            _ => Err(ADError::UnsupportedConversion(
136                "unsupported TIFF pixel format".into(),
137            )),
138        }
139    }
140
141    fn close_file(&mut self) -> ADResult<()> {
142        self.current_path = None;
143        Ok(())
144    }
145
146    fn supports_multiple_arrays(&self) -> bool {
147        false
148    }
149}
150
151/// TIFF file processor wrapping NDPluginFileBase + TiffWriter.
152pub struct TiffFileProcessor {
153    file_base: NDPluginFileBase,
154    writer: TiffWriter,
155}
156
157impl TiffFileProcessor {
158    pub fn new() -> Self {
159        Self {
160            file_base: NDPluginFileBase::new(),
161            writer: TiffWriter::new(),
162        }
163    }
164
165    pub fn file_base_mut(&mut self) -> &mut NDPluginFileBase {
166        &mut self.file_base
167    }
168}
169
170impl Default for TiffFileProcessor {
171    fn default() -> Self {
172        Self::new()
173    }
174}
175
176impl NDPluginProcess for TiffFileProcessor {
177    fn process_array(&mut self, array: &NDArray, _pool: &NDArrayPool) -> ProcessResult {
178        let _ = self
179            .file_base
180            .process_array(Arc::new(array.clone()), &mut self.writer);
181        ProcessResult::empty() // file plugins are sinks
182    }
183
184    fn plugin_type(&self) -> &str {
185        "NDFileTIFF"
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192    use ad_core::ndarray::NDDataBuffer;
193    use std::sync::atomic::{AtomicU32, Ordering};
194
195    static TEST_COUNTER: AtomicU32 = AtomicU32::new(0);
196
197    fn temp_path(prefix: &str) -> PathBuf {
198        let n = TEST_COUNTER.fetch_add(1, Ordering::Relaxed);
199        std::env::temp_dir().join(format!("adcore_test_{}_{}.tif", prefix, n))
200    }
201
202    #[test]
203    fn test_write_u8_mono() {
204        let path = temp_path("tiff_u8");
205        let mut writer = TiffWriter::new();
206
207        let mut arr = NDArray::new(
208            vec![NDDimension::new(4), NDDimension::new(4)],
209            NDDataType::UInt8,
210        );
211        if let NDDataBuffer::U8(v) = &mut arr.data {
212            for i in 0..16 { v[i] = i as u8; }
213        }
214
215        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
216        writer.write_file(&arr).unwrap();
217        writer.close_file().unwrap();
218
219        // Verify file exists and has content
220        let data = std::fs::read(&path).unwrap();
221        assert!(data.len() > 16); // header + data
222        // Check TIFF magic (little-endian II or big-endian MM)
223        assert!(
224            &data[0..2] == &[0x49, 0x49] || &data[0..2] == &[0x4D, 0x4D],
225            "Expected TIFF magic bytes"
226        );
227
228        std::fs::remove_file(&path).ok();
229    }
230
231    #[test]
232    fn test_write_u16() {
233        let path = temp_path("tiff_u16");
234        let mut writer = TiffWriter::new();
235
236        let arr = NDArray::new(
237            vec![NDDimension::new(4), NDDimension::new(4)],
238            NDDataType::UInt16,
239        );
240
241        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
242        writer.write_file(&arr).unwrap();
243        writer.close_file().unwrap();
244
245        let data = std::fs::read(&path).unwrap();
246        assert!(data.len() > 32); // 16 elements * 2 bytes + header
247
248        std::fs::remove_file(&path).ok();
249    }
250
251    #[test]
252    fn test_roundtrip_u8() {
253        let path = temp_path("tiff_rt_u8");
254        let mut writer = TiffWriter::new();
255
256        let mut arr = NDArray::new(
257            vec![NDDimension::new(4), NDDimension::new(4)],
258            NDDataType::UInt8,
259        );
260        if let NDDataBuffer::U8(v) = &mut arr.data {
261            for i in 0..16 { v[i] = (i * 10) as u8; }
262        }
263
264        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
265        writer.write_file(&arr).unwrap();
266
267        // Read it back
268        let read_back = writer.read_file().unwrap();
269        if let (NDDataBuffer::U8(orig), NDDataBuffer::U8(read)) =
270            (&arr.data, &read_back.data)
271        {
272            assert_eq!(orig, read);
273        } else {
274            panic!("data type mismatch on roundtrip");
275        }
276
277        writer.close_file().unwrap();
278        std::fs::remove_file(&path).ok();
279    }
280
281    #[test]
282    fn test_roundtrip_u16() {
283        let path = temp_path("tiff_rt_u16");
284        let mut writer = TiffWriter::new();
285
286        let mut arr = NDArray::new(
287            vec![NDDimension::new(4), NDDimension::new(4)],
288            NDDataType::UInt16,
289        );
290        if let NDDataBuffer::U16(v) = &mut arr.data {
291            for i in 0..16 { v[i] = (i * 1000) as u16; }
292        }
293
294        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
295        writer.write_file(&arr).unwrap();
296
297        let read_back = writer.read_file().unwrap();
298        if let (NDDataBuffer::U16(orig), NDDataBuffer::U16(read)) =
299            (&arr.data, &read_back.data)
300        {
301            assert_eq!(orig, read);
302        } else {
303            panic!("data type mismatch on roundtrip");
304        }
305
306        writer.close_file().unwrap();
307        std::fs::remove_file(&path).ok();
308    }
309}