Skip to main content

ad_plugins_rs/
file_tiff.rs

1use std::path::{Path, PathBuf};
2
3use ad_core_rs::attributes::{NDAttrSource, NDAttrValue, NDAttribute};
4use ad_core_rs::color::{NDColorMode, convert_rgb_layout};
5use ad_core_rs::error::{ADError, ADResult};
6use ad_core_rs::ndarray::{NDArray, NDDataBuffer, NDDataType, NDDimension};
7use ad_core_rs::ndarray_pool::NDArrayPool;
8use ad_core_rs::plugin::file_base::{NDFileMode, NDFileWriter};
9use ad_core_rs::plugin::file_controller::FilePluginController;
10use ad_core_rs::plugin::runtime::{
11    NDPluginProcess, ParamChangeResult, PluginParamSnapshot, ProcessResult,
12};
13
14use tiff::ColorType;
15use tiff::decoder::Decoder;
16use tiff::encoder::TiffEncoder;
17use tiff::encoder::colortype;
18
19/// TIFF file writer using the `tiff` crate for proper encoding/decoding.
20pub struct TiffWriter {
21    current_path: Option<PathBuf>,
22}
23
24impl TiffWriter {
25    pub fn new() -> Self {
26        Self { current_path: None }
27    }
28
29    fn array_color_mode(array: &NDArray) -> NDColorMode {
30        array
31            .attributes
32            .get("ColorMode")
33            .and_then(|attr| attr.value.as_i64())
34            .map(|v| NDColorMode::from_i32(v as i32))
35            .unwrap_or_else(|| match array.dims.as_slice() {
36                [a, _, _] if a.size == 3 => NDColorMode::RGB1,
37                [_, b, _] if b.size == 3 => NDColorMode::RGB2,
38                [_, _, c] if c.size == 3 => NDColorMode::RGB3,
39                _ => NDColorMode::Mono,
40            })
41    }
42
43    fn normalize_for_write(array: &NDArray) -> ADResult<(NDArray, u32, u32, bool)> {
44        match array.dims.as_slice() {
45            [x] => {
46                let mut normalized = NDArray::new(
47                    vec![NDDimension::new(x.size), NDDimension::new(1)],
48                    array.data.data_type(),
49                );
50                normalized.data = array.data.clone();
51                normalized.unique_id = array.unique_id;
52                normalized.timestamp = array.timestamp;
53                normalized.attributes = array.attributes.clone();
54                normalized.codec = array.codec.clone();
55                Ok((normalized, x.size as u32, 1, false))
56            }
57            [x, y] => Ok((array.clone(), x.size as u32, y.size as u32, false)),
58            [_, _, _] => {
59                let color_mode = Self::array_color_mode(array);
60                let rgb1 = match color_mode {
61                    NDColorMode::RGB1 => array.clone(),
62                    NDColorMode::RGB2 | NDColorMode::RGB3 => {
63                        convert_rgb_layout(array, color_mode, NDColorMode::RGB1)?
64                    }
65                    other => {
66                        return Err(ADError::UnsupportedConversion(format!(
67                            "unsupported TIFF color mode: {:?}",
68                            other
69                        )));
70                    }
71                };
72                Ok((
73                    rgb1.clone(),
74                    rgb1.dims[1].size as u32,
75                    rgb1.dims[2].size as u32,
76                    true,
77                ))
78            }
79            _ => Err(ADError::InvalidDimensions(
80                "unsupported TIFF array dimensions".into(),
81            )),
82        }
83    }
84
85    fn attach_color_mode(array: &mut NDArray, color_mode: NDColorMode) {
86        array.attributes.add(NDAttribute {
87            name: "ColorMode".into(),
88            description: "Color mode".into(),
89            source: NDAttrSource::Driver,
90            value: NDAttrValue::Int32(color_mode as i32),
91        });
92    }
93}
94
95impl NDFileWriter for TiffWriter {
96    fn open_file(&mut self, path: &Path, _mode: NDFileMode, _array: &NDArray) -> ADResult<()> {
97        self.current_path = Some(path.to_path_buf());
98        Ok(())
99    }
100
101    fn write_file(&mut self, array: &NDArray) -> ADResult<()> {
102        let path = self
103            .current_path
104            .as_ref()
105            .ok_or_else(|| ADError::UnsupportedConversion("no file open".into()))?;
106        let (array, width, height, is_rgb) = Self::normalize_for_write(array)?;
107
108        let file = std::fs::File::create(path)?;
109        let mut encoder = TiffEncoder::new(file)
110            .map_err(|e| ADError::UnsupportedConversion(format!("TIFF encoder error: {}", e)))?;
111
112        match &array.data {
113            NDDataBuffer::U8(v) => {
114                if is_rgb {
115                    encoder.write_image::<colortype::RGB8>(width, height, v)
116                } else {
117                    encoder.write_image::<colortype::Gray8>(width, height, v)
118                }
119            }
120            NDDataBuffer::I8(v) => {
121                if is_rgb {
122                    return Err(ADError::UnsupportedConversion(
123                        "TIFF crate does not support signed RGB8".into(),
124                    ));
125                }
126                encoder.write_image::<colortype::GrayI8>(width, height, v)
127            }
128            NDDataBuffer::U16(v) => {
129                if is_rgb {
130                    encoder.write_image::<colortype::RGB16>(width, height, v)
131                } else {
132                    encoder.write_image::<colortype::Gray16>(width, height, v)
133                }
134            }
135            NDDataBuffer::I16(v) => {
136                if is_rgb {
137                    return Err(ADError::UnsupportedConversion(
138                        "TIFF crate does not support signed RGB16".into(),
139                    ));
140                }
141                encoder.write_image::<colortype::GrayI16>(width, height, v)
142            }
143            NDDataBuffer::U32(v) => {
144                if is_rgb {
145                    encoder.write_image::<colortype::RGB32>(width, height, v)
146                } else {
147                    encoder.write_image::<colortype::Gray32>(width, height, v)
148                }
149            }
150            NDDataBuffer::I32(v) => {
151                if is_rgb {
152                    return Err(ADError::UnsupportedConversion(
153                        "TIFF crate does not support signed RGB32".into(),
154                    ));
155                }
156                encoder.write_image::<colortype::GrayI32>(width, height, v)
157            }
158            NDDataBuffer::I64(v) => {
159                if is_rgb {
160                    return Err(ADError::UnsupportedConversion(
161                        "TIFF crate does not support signed RGB64".into(),
162                    ));
163                }
164                encoder.write_image::<colortype::GrayI64>(width, height, v)
165            }
166            NDDataBuffer::U64(v) => {
167                if is_rgb {
168                    encoder.write_image::<colortype::RGB64>(width, height, v)
169                } else {
170                    encoder.write_image::<colortype::Gray64>(width, height, v)
171                }
172            }
173            NDDataBuffer::F32(v) => {
174                if is_rgb {
175                    encoder.write_image::<colortype::RGB32Float>(width, height, v)
176                } else {
177                    encoder.write_image::<colortype::Gray32Float>(width, height, v)
178                }
179            }
180            NDDataBuffer::F64(v) => {
181                if is_rgb {
182                    encoder.write_image::<colortype::RGB64Float>(width, height, v)
183                } else {
184                    encoder.write_image::<colortype::Gray64Float>(width, height, v)
185                }
186            }
187        }
188        .map_err(|e| ADError::UnsupportedConversion(format!("TIFF write error: {}", e)))?;
189
190        Ok(())
191    }
192
193    fn read_file(&mut self) -> ADResult<NDArray> {
194        let path = self
195            .current_path
196            .as_ref()
197            .ok_or_else(|| ADError::UnsupportedConversion("no file open".into()))?;
198
199        let file = std::fs::File::open(path)?;
200        let mut decoder = Decoder::new(file)
201            .map_err(|e| ADError::UnsupportedConversion(format!("TIFF decode error: {}", e)))?;
202
203        let (width, height) = decoder
204            .dimensions()
205            .map_err(|e| ADError::UnsupportedConversion(format!("TIFF dimensions error: {}", e)))?;
206        let color_type = decoder
207            .colortype()
208            .map_err(|e| ADError::UnsupportedConversion(format!("TIFF colortype error: {}", e)))?;
209
210        let result = decoder
211            .read_image()
212            .map_err(|e| ADError::UnsupportedConversion(format!("TIFF read error: {}", e)))?;
213
214        let (dims, color_mode) = match color_type {
215            ColorType::Gray(_) => (
216                vec![
217                    NDDimension::new(width as usize),
218                    NDDimension::new(height as usize),
219                ],
220                NDColorMode::Mono,
221            ),
222            ColorType::RGB(_) => (
223                vec![
224                    NDDimension::new(3),
225                    NDDimension::new(width as usize),
226                    NDDimension::new(height as usize),
227                ],
228                NDColorMode::RGB1,
229            ),
230            other => {
231                return Err(ADError::UnsupportedConversion(format!(
232                    "unsupported TIFF color type: {:?}",
233                    other
234                )));
235            }
236        };
237
238        let mut array = match result {
239            tiff::decoder::DecodingResult::U8(data) => {
240                let mut arr = NDArray::new(dims.clone(), NDDataType::UInt8);
241                arr.data = NDDataBuffer::U8(data);
242                arr
243            }
244            tiff::decoder::DecodingResult::U16(data) => {
245                let mut arr = NDArray::new(dims.clone(), NDDataType::UInt16);
246                arr.data = NDDataBuffer::U16(data);
247                arr
248            }
249            tiff::decoder::DecodingResult::U32(data) => {
250                let mut arr = NDArray::new(dims.clone(), NDDataType::UInt32);
251                arr.data = NDDataBuffer::U32(data);
252                arr
253            }
254            tiff::decoder::DecodingResult::U64(data) => {
255                let mut arr = NDArray::new(dims.clone(), NDDataType::UInt64);
256                arr.data = NDDataBuffer::U64(data);
257                arr
258            }
259            tiff::decoder::DecodingResult::I8(data) => {
260                let mut arr = NDArray::new(dims.clone(), NDDataType::Int8);
261                arr.data = NDDataBuffer::I8(data);
262                arr
263            }
264            tiff::decoder::DecodingResult::I16(data) => {
265                let mut arr = NDArray::new(dims.clone(), NDDataType::Int16);
266                arr.data = NDDataBuffer::I16(data);
267                arr
268            }
269            tiff::decoder::DecodingResult::I32(data) => {
270                let mut arr = NDArray::new(dims.clone(), NDDataType::Int32);
271                arr.data = NDDataBuffer::I32(data);
272                arr
273            }
274            tiff::decoder::DecodingResult::I64(data) => {
275                let mut arr = NDArray::new(dims.clone(), NDDataType::Int64);
276                arr.data = NDDataBuffer::I64(data);
277                arr
278            }
279            tiff::decoder::DecodingResult::F32(data) => {
280                let mut arr = NDArray::new(dims.clone(), NDDataType::Float32);
281                arr.data = NDDataBuffer::F32(data);
282                arr
283            }
284            tiff::decoder::DecodingResult::F64(data) => {
285                let mut arr = NDArray::new(dims.clone(), NDDataType::Float64);
286                arr.data = NDDataBuffer::F64(data);
287                arr
288            }
289        };
290        Self::attach_color_mode(&mut array, color_mode);
291        Ok(array)
292    }
293
294    fn close_file(&mut self) -> ADResult<()> {
295        self.current_path = None;
296        Ok(())
297    }
298
299    fn supports_multiple_arrays(&self) -> bool {
300        false
301    }
302}
303
304/// TIFF file processor wrapping FilePluginController<TiffWriter>.
305pub struct TiffFileProcessor {
306    pub ctrl: FilePluginController<TiffWriter>,
307}
308
309impl TiffFileProcessor {
310    pub fn new() -> Self {
311        Self {
312            ctrl: FilePluginController::new(TiffWriter::new()),
313        }
314    }
315}
316
317impl Default for TiffFileProcessor {
318    fn default() -> Self {
319        Self::new()
320    }
321}
322
323impl NDPluginProcess for TiffFileProcessor {
324    fn process_array(&mut self, array: &NDArray, _pool: &NDArrayPool) -> ProcessResult {
325        self.ctrl.process_array(array)
326    }
327
328    fn plugin_type(&self) -> &str {
329        "NDFileTIFF"
330    }
331
332    fn register_params(
333        &mut self,
334        base: &mut asyn_rs::port::PortDriverBase,
335    ) -> asyn_rs::error::AsynResult<()> {
336        self.ctrl.register_params(base)
337    }
338
339    fn on_param_change(
340        &mut self,
341        reason: usize,
342        params: &PluginParamSnapshot,
343    ) -> ParamChangeResult {
344        self.ctrl.on_param_change(reason, params)
345    }
346}
347
348#[cfg(test)]
349mod tests {
350    use super::*;
351    use ad_core_rs::ndarray::NDDataBuffer;
352    use ad_core_rs::params::ndarray_driver::NDArrayDriverParams;
353    use ad_core_rs::plugin::runtime::{ParamChangeValue, ParamUpdate, PluginParamSnapshot};
354    use asyn_rs::port::{PortDriverBase, PortFlags};
355    use std::sync::atomic::{AtomicU32, Ordering};
356
357    static TEST_COUNTER: AtomicU32 = AtomicU32::new(0);
358
359    fn temp_path(prefix: &str) -> PathBuf {
360        let n = TEST_COUNTER.fetch_add(1, Ordering::Relaxed);
361        std::env::temp_dir().join(format!("adcore_test_{}_{}.tif", prefix, n))
362    }
363
364    #[test]
365    fn test_write_u8_mono() {
366        let path = temp_path("tiff_u8");
367        let mut writer = TiffWriter::new();
368
369        let mut arr = NDArray::new(
370            vec![NDDimension::new(4), NDDimension::new(4)],
371            NDDataType::UInt8,
372        );
373        if let NDDataBuffer::U8(v) = &mut arr.data {
374            for i in 0..16 {
375                v[i] = i as u8;
376            }
377        }
378
379        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
380        writer.write_file(&arr).unwrap();
381        writer.close_file().unwrap();
382
383        let data = std::fs::read(&path).unwrap();
384        assert!(data.len() > 16);
385        assert!(
386            &data[0..2] == &[0x49, 0x49] || &data[0..2] == &[0x4D, 0x4D],
387            "Expected TIFF magic bytes"
388        );
389
390        std::fs::remove_file(&path).ok();
391    }
392
393    #[test]
394    fn test_write_u16() {
395        let path = temp_path("tiff_u16");
396        let mut writer = TiffWriter::new();
397
398        let arr = NDArray::new(
399            vec![NDDimension::new(4), NDDimension::new(4)],
400            NDDataType::UInt16,
401        );
402
403        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
404        writer.write_file(&arr).unwrap();
405        writer.close_file().unwrap();
406
407        let data = std::fs::read(&path).unwrap();
408        assert!(data.len() > 32);
409
410        std::fs::remove_file(&path).ok();
411    }
412
413    #[test]
414    fn test_roundtrip_u8() {
415        let path = temp_path("tiff_rt_u8");
416        let mut writer = TiffWriter::new();
417
418        let mut arr = NDArray::new(
419            vec![NDDimension::new(4), NDDimension::new(4)],
420            NDDataType::UInt8,
421        );
422        if let NDDataBuffer::U8(v) = &mut arr.data {
423            for i in 0..16 {
424                v[i] = (i * 10) as u8;
425            }
426        }
427
428        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
429        writer.write_file(&arr).unwrap();
430
431        let read_back = writer.read_file().unwrap();
432        if let (NDDataBuffer::U8(orig), NDDataBuffer::U8(read)) = (&arr.data, &read_back.data) {
433            assert_eq!(orig, read);
434        } else {
435            panic!("data type mismatch on roundtrip");
436        }
437
438        writer.close_file().unwrap();
439        std::fs::remove_file(&path).ok();
440    }
441
442    #[test]
443    fn test_roundtrip_u16() {
444        let path = temp_path("tiff_rt_u16");
445        let mut writer = TiffWriter::new();
446
447        let mut arr = NDArray::new(
448            vec![NDDimension::new(4), NDDimension::new(4)],
449            NDDataType::UInt16,
450        );
451        if let NDDataBuffer::U16(v) = &mut arr.data {
452            for i in 0..16 {
453                v[i] = (i * 1000) as u16;
454            }
455        }
456
457        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
458        writer.write_file(&arr).unwrap();
459
460        let read_back = writer.read_file().unwrap();
461        if let (NDDataBuffer::U16(orig), NDDataBuffer::U16(read)) = (&arr.data, &read_back.data) {
462            assert_eq!(orig, read);
463        } else {
464            panic!("data type mismatch on roundtrip");
465        }
466
467        writer.close_file().unwrap();
468        std::fs::remove_file(&path).ok();
469    }
470
471    #[test]
472    fn test_on_param_change_read_file_emits_array_and_resets_busy() {
473        let path = temp_path("tiff_read_param");
474        let mut writer = TiffWriter::new();
475
476        let mut arr = NDArray::new(
477            vec![NDDimension::new(4), NDDimension::new(3)],
478            NDDataType::UInt8,
479        );
480        arr.unique_id = 77;
481        if let NDDataBuffer::U8(v) = &mut arr.data {
482            for (i, item) in v.iter_mut().enumerate() {
483                *item = i as u8;
484            }
485        }
486
487        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
488        writer.write_file(&arr).unwrap();
489        writer.close_file().unwrap();
490
491        let mut base = PortDriverBase::new("TIFFTEST", 1, PortFlags::default());
492        let _nd_params = NDArrayDriverParams::create(&mut base).unwrap();
493
494        let mut proc = TiffFileProcessor::new();
495        proc.register_params(&mut base).unwrap();
496
497        let reason_path = base.find_param("FILE_PATH").unwrap();
498        let reason_name = base.find_param("FILE_NAME").unwrap();
499        let reason_template = base.find_param("FILE_TEMPLATE").unwrap();
500        let reason_read = base.find_param("READ_FILE").unwrap();
501
502        let _ = proc.on_param_change(
503            reason_path,
504            &PluginParamSnapshot {
505                enable_callbacks: true,
506                reason: reason_path,
507                addr: 0,
508                value: ParamChangeValue::Octet(
509                    path.parent().unwrap().to_str().unwrap().to_string(),
510                ),
511            },
512        );
513        let _ = proc.on_param_change(
514            reason_name,
515            &PluginParamSnapshot {
516                enable_callbacks: true,
517                reason: reason_name,
518                addr: 0,
519                value: ParamChangeValue::Octet(
520                    path.file_name().unwrap().to_str().unwrap().to_string(),
521                ),
522            },
523        );
524        let _ = proc.on_param_change(
525            reason_template,
526            &PluginParamSnapshot {
527                enable_callbacks: true,
528                reason: reason_template,
529                addr: 0,
530                value: ParamChangeValue::Octet("%s%s".into()),
531            },
532        );
533
534        let result = proc.on_param_change(
535            reason_read,
536            &PluginParamSnapshot {
537                enable_callbacks: true,
538                reason: reason_read,
539                addr: 0,
540                value: ParamChangeValue::Int32(1),
541            },
542        );
543
544        assert_eq!(result.output_arrays.len(), 1);
545        assert!(result.param_updates.iter().any(|u| matches!(
546            u,
547            ParamUpdate::Int32 { reason, value: 0, .. } if *reason == reason_read
548        )));
549        match &result.output_arrays[0].data {
550            NDDataBuffer::U8(v) => assert_eq!(v.len(), 12),
551            other => panic!("unexpected data buffer: {other:?}"),
552        }
553
554        std::fs::remove_file(&path).ok();
555    }
556
557    #[test]
558    fn test_single_mode_requires_auto_save_for_automatic_write() {
559        let path = temp_path("tiff_autosave_single");
560        let full_name = path.to_string_lossy().to_string();
561        let file_path = path.parent().unwrap().to_str().unwrap().to_string();
562        let file_name = path.file_name().unwrap().to_str().unwrap().to_string();
563
564        let mut proc = TiffFileProcessor::new();
565        proc.ctrl.file_base.file_path = file_path.clone() + "/";
566        proc.ctrl.file_base.file_name = file_name;
567        proc.ctrl.file_base.file_template = "%s%s".into();
568        proc.ctrl.file_base.set_mode(NDFileMode::Single);
569
570        let mut arr = NDArray::new(
571            vec![NDDimension::new(4), NDDimension::new(4)],
572            NDDataType::UInt8,
573        );
574        if let NDDataBuffer::U8(v) = &mut arr.data {
575            for (i, item) in v.iter_mut().enumerate() {
576                *item = i as u8;
577            }
578        }
579
580        proc.ctrl.auto_save = false;
581        let _ = proc.process_array(&arr, &NDArrayPool::new(1024));
582        assert!(!std::path::Path::new(&full_name).exists());
583
584        proc.ctrl.auto_save = true;
585        let _ = proc.process_array(&arr, &NDArrayPool::new(1024));
586        assert!(std::path::Path::new(&full_name).exists());
587
588        std::fs::remove_file(&path).ok();
589    }
590}