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
13pub 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 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 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 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
151pub 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() }
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(ref mut v) = 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 let data = std::fs::read(&path).unwrap();
221 assert!(data.len() > 16); 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); 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(ref mut v) = 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 let read_back = writer.read_file().unwrap();
269 if let (NDDataBuffer::U8(ref orig), NDDataBuffer::U8(ref 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(ref mut v) = 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(ref orig), NDDataBuffer::U16(ref 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}