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