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
19pub 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
304pub 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}