ad_plugins_rs/
file_magick.rs

1use std::path::{Path, PathBuf};
2
3use ad_core_rs::color::{NDColorMode, convert_rgb_layout};
4use ad_core_rs::error::{ADError, ADResult};
5use ad_core_rs::ndarray::{NDArray, NDDataBuffer, NDDataType, NDDimension};
6use ad_core_rs::ndarray_pool::NDArrayPool;
7use ad_core_rs::plugin::file_base::{NDFileMode, NDFileWriter};
8use ad_core_rs::plugin::file_controller::FilePluginController;
9use ad_core_rs::plugin::runtime::{
10    NDPluginProcess, ParamChangeResult, PluginParamSnapshot, ProcessResult,
11};
12
13use image::{DynamicImage, ImageFormat};
14
15/// NDFileMagick: file writer using the `image` crate.
16///
17/// Format is determined by the file extension (PNG, BMP, GIF, TIFF, etc.).
18/// Supports UInt8 and UInt16 data in mono and RGB color modes.
19pub struct MagickWriter {
20    current_path: Option<PathBuf>,
21    quality: u8,
22    bit_depth: u32,
23}
24
25impl MagickWriter {
26    pub fn new() -> Self {
27        Self {
28            current_path: None,
29            quality: 100,
30            bit_depth: 8,
31        }
32    }
33
34    pub fn set_quality(&mut self, q: u8) {
35        self.quality = q;
36    }
37
38    pub fn set_bit_depth(&mut self, depth: u32) {
39        self.bit_depth = depth;
40    }
41
42    fn color_mode(array: &NDArray) -> NDColorMode {
43        array
44            .attributes
45            .get("ColorMode")
46            .and_then(|attr| attr.value.as_i64())
47            .map(|v| NDColorMode::from_i32(v as i32))
48            .unwrap_or_else(|| match array.dims.as_slice() {
49                [a, _, _] if a.size == 3 => NDColorMode::RGB1,
50                [_, b, _] if b.size == 3 => NDColorMode::RGB2,
51                [_, _, c] if c.size == 3 => NDColorMode::RGB3,
52                _ => NDColorMode::Mono,
53            })
54    }
55
56    /// Convert NDArray to DynamicImage for encoding.
57    fn array_to_image(array: &NDArray) -> ADResult<DynamicImage> {
58        let info = array.info();
59        let width = info.x_size as u32;
60        let height = info.y_size as u32;
61        let color = Self::color_mode(array);
62        let is_rgb = matches!(
63            color,
64            NDColorMode::RGB1 | NDColorMode::RGB2 | NDColorMode::RGB3
65        );
66
67        // Convert to RGB1 layout if needed (image crate expects interleaved RGB)
68        let src = if is_rgb && color != NDColorMode::RGB1 {
69            &convert_rgb_layout(array, color, NDColorMode::RGB1)?
70        } else {
71            array
72        };
73
74        match &src.data {
75            NDDataBuffer::U8(v) => {
76                if is_rgb {
77                    image::RgbImage::from_raw(width, height, v.clone())
78                        .map(DynamicImage::ImageRgb8)
79                        .ok_or_else(|| {
80                            ADError::UnsupportedConversion("RGB8 buffer size mismatch".into())
81                        })
82                } else {
83                    image::GrayImage::from_raw(width, height, v.clone())
84                        .map(DynamicImage::ImageLuma8)
85                        .ok_or_else(|| {
86                            ADError::UnsupportedConversion("Gray8 buffer size mismatch".into())
87                        })
88                }
89            }
90            NDDataBuffer::I8(v) => {
91                let u8_data: Vec<u8> = v.iter().map(|&b| b as u8).collect();
92                if is_rgb {
93                    image::RgbImage::from_raw(width, height, u8_data)
94                        .map(DynamicImage::ImageRgb8)
95                        .ok_or_else(|| {
96                            ADError::UnsupportedConversion("RGB8 buffer size mismatch".into())
97                        })
98                } else {
99                    image::GrayImage::from_raw(width, height, u8_data)
100                        .map(DynamicImage::ImageLuma8)
101                        .ok_or_else(|| {
102                            ADError::UnsupportedConversion("Gray8 buffer size mismatch".into())
103                        })
104                }
105            }
106            NDDataBuffer::U16(v) => {
107                if is_rgb {
108                    image::ImageBuffer::<image::Rgb<u16>, Vec<u16>>::from_raw(
109                        width,
110                        height,
111                        v.clone(),
112                    )
113                    .map(DynamicImage::ImageRgb16)
114                    .ok_or_else(|| {
115                        ADError::UnsupportedConversion("RGB16 buffer size mismatch".into())
116                    })
117                } else {
118                    image::ImageBuffer::<image::Luma<u16>, Vec<u16>>::from_raw(
119                        width,
120                        height,
121                        v.clone(),
122                    )
123                    .map(DynamicImage::ImageLuma16)
124                    .ok_or_else(|| {
125                        ADError::UnsupportedConversion("Gray16 buffer size mismatch".into())
126                    })
127                }
128            }
129            NDDataBuffer::I16(v) => {
130                let u16_data: Vec<u16> = v.iter().map(|&b| b as u16).collect();
131                if is_rgb {
132                    image::ImageBuffer::<image::Rgb<u16>, Vec<u16>>::from_raw(
133                        width, height, u16_data,
134                    )
135                    .map(DynamicImage::ImageRgb16)
136                    .ok_or_else(|| {
137                        ADError::UnsupportedConversion("RGB16 buffer size mismatch".into())
138                    })
139                } else {
140                    image::ImageBuffer::<image::Luma<u16>, Vec<u16>>::from_raw(
141                        width, height, u16_data,
142                    )
143                    .map(DynamicImage::ImageLuma16)
144                    .ok_or_else(|| {
145                        ADError::UnsupportedConversion("Gray16 buffer size mismatch".into())
146                    })
147                }
148            }
149            NDDataBuffer::F32(v) => {
150                let u16_data: Vec<u16> = v
151                    .iter()
152                    .map(|&f| (f.clamp(0.0, 1.0) * 65535.0) as u16)
153                    .collect();
154                if is_rgb {
155                    image::ImageBuffer::<image::Rgb<u16>, Vec<u16>>::from_raw(
156                        width, height, u16_data,
157                    )
158                    .map(DynamicImage::ImageRgb16)
159                    .ok_or_else(|| {
160                        ADError::UnsupportedConversion("RGB16 buffer size mismatch".into())
161                    })
162                } else {
163                    image::ImageBuffer::<image::Luma<u16>, Vec<u16>>::from_raw(
164                        width, height, u16_data,
165                    )
166                    .map(DynamicImage::ImageLuma16)
167                    .ok_or_else(|| {
168                        ADError::UnsupportedConversion("Gray16 buffer size mismatch".into())
169                    })
170                }
171            }
172            _ => Err(ADError::UnsupportedConversion(format!(
173                "NDFileMagick: unsupported data type {:?}, use UInt8, Int8, UInt16, Int16, or Float32",
174                src.data.data_type()
175            ))),
176        }
177    }
178}
179
180impl NDFileWriter for MagickWriter {
181    fn open_file(&mut self, path: &Path, _mode: NDFileMode, _array: &NDArray) -> ADResult<()> {
182        self.current_path = Some(path.to_path_buf());
183        Ok(())
184    }
185
186    fn write_file(&mut self, array: &NDArray) -> ADResult<()> {
187        let path = self
188            .current_path
189            .as_ref()
190            .ok_or_else(|| ADError::UnsupportedConversion("no file open".into()))?;
191
192        let img = Self::array_to_image(array)?;
193
194        // Determine format from extension, default to PNG
195        let format = ImageFormat::from_path(path).unwrap_or(ImageFormat::Png);
196
197        // For JPEG, use quality setting
198        if format == ImageFormat::Jpeg {
199            let mut buf = Vec::new();
200            let encoder =
201                image::codecs::jpeg::JpegEncoder::new_with_quality(&mut buf, self.quality);
202            img.write_with_encoder(encoder)
203                .map_err(|e| ADError::UnsupportedConversion(format!("Magick encode error: {e}")))?;
204            std::fs::write(path, &buf)?;
205        } else {
206            img.save(path)
207                .map_err(|e| ADError::UnsupportedConversion(format!("Magick save error: {e}")))?;
208        }
209
210        Ok(())
211    }
212
213    fn read_file(&mut self) -> ADResult<NDArray> {
214        let path = self
215            .current_path
216            .as_ref()
217            .ok_or_else(|| ADError::UnsupportedConversion("no file open".into()))?;
218
219        let img = image::open(path)
220            .map_err(|e| ADError::UnsupportedConversion(format!("Magick read error: {e}")))?;
221
222        let width = img.width() as usize;
223        let height = img.height() as usize;
224
225        match img {
226            DynamicImage::ImageLuma8(buf) => {
227                let mut arr = NDArray::new(
228                    vec![NDDimension::new(width), NDDimension::new(height)],
229                    NDDataType::UInt8,
230                );
231                arr.data = NDDataBuffer::U8(buf.into_raw());
232                Ok(arr)
233            }
234            DynamicImage::ImageRgb8(buf) => {
235                let mut arr = NDArray::new(
236                    vec![
237                        NDDimension::new(3),
238                        NDDimension::new(width),
239                        NDDimension::new(height),
240                    ],
241                    NDDataType::UInt8,
242                );
243                arr.data = NDDataBuffer::U8(buf.into_raw());
244                Ok(arr)
245            }
246            DynamicImage::ImageLuma16(buf) => {
247                let mut arr = NDArray::new(
248                    vec![NDDimension::new(width), NDDimension::new(height)],
249                    NDDataType::UInt16,
250                );
251                arr.data = NDDataBuffer::U16(buf.into_raw());
252                Ok(arr)
253            }
254            DynamicImage::ImageRgb16(buf) => {
255                let mut arr = NDArray::new(
256                    vec![
257                        NDDimension::new(3),
258                        NDDimension::new(width),
259                        NDDimension::new(height),
260                    ],
261                    NDDataType::UInt16,
262                );
263                arr.data = NDDataBuffer::U16(buf.into_raw());
264                Ok(arr)
265            }
266            other => {
267                // Convert anything else to RGB8
268                let rgb = other.to_rgb8();
269                let mut arr = NDArray::new(
270                    vec![
271                        NDDimension::new(3),
272                        NDDimension::new(width),
273                        NDDimension::new(height),
274                    ],
275                    NDDataType::UInt8,
276                );
277                arr.data = NDDataBuffer::U8(rgb.into_raw());
278                Ok(arr)
279            }
280        }
281    }
282
283    fn close_file(&mut self) -> ADResult<()> {
284        self.current_path = None;
285        Ok(())
286    }
287
288    fn supports_multiple_arrays(&self) -> bool {
289        false
290    }
291}
292
293/// Magick file processor wrapping FilePluginController<MagickWriter>.
294pub struct MagickFileProcessor {
295    ctrl: FilePluginController<MagickWriter>,
296    quality_idx: Option<usize>,
297    bit_depth_idx: Option<usize>,
298    compress_type_idx: Option<usize>,
299}
300
301impl MagickFileProcessor {
302    pub fn new() -> Self {
303        Self {
304            ctrl: FilePluginController::new(MagickWriter::new()),
305            quality_idx: None,
306            bit_depth_idx: None,
307            compress_type_idx: None,
308        }
309    }
310}
311
312impl NDPluginProcess for MagickFileProcessor {
313    fn process_array(&mut self, array: &NDArray, _pool: &NDArrayPool) -> ProcessResult {
314        self.ctrl.process_array(array)
315    }
316
317    fn plugin_type(&self) -> &str {
318        "NDFileMagick"
319    }
320
321    fn register_params(
322        &mut self,
323        base: &mut asyn_rs::port::PortDriverBase,
324    ) -> asyn_rs::error::AsynResult<()> {
325        self.ctrl.register_params(base)?;
326        use asyn_rs::param::ParamType;
327        self.quality_idx = Some(base.create_param("MAGICK_QUALITY", ParamType::Int32)?);
328        self.bit_depth_idx = Some(base.create_param("MAGICK_BIT_DEPTH", ParamType::Int32)?);
329        self.compress_type_idx = Some(base.create_param("MAGICK_COMPRESS_TYPE", ParamType::Int32)?);
330        // Set defaults
331        base.set_int32_param(self.quality_idx.unwrap(), 0, 100)?;
332        base.set_int32_param(self.bit_depth_idx.unwrap(), 0, 8)?;
333        base.set_int32_param(self.compress_type_idx.unwrap(), 0, 0)?;
334        Ok(())
335    }
336
337    fn on_param_change(
338        &mut self,
339        reason: usize,
340        params: &PluginParamSnapshot,
341    ) -> ParamChangeResult {
342        if Some(reason) == self.quality_idx {
343            let q = params.value.as_i32().clamp(1, 100) as u8;
344            self.ctrl.writer.set_quality(q);
345            return ParamChangeResult::empty();
346        }
347        if Some(reason) == self.bit_depth_idx {
348            let d = params.value.as_i32() as u32;
349            self.ctrl.writer.set_bit_depth(d);
350            return ParamChangeResult::empty();
351        }
352        if Some(reason) == self.compress_type_idx {
353            // CompressType stored but not actively used by `image` crate;
354            // format-specific compression is handled by each codec internally.
355            return ParamChangeResult::empty();
356        }
357        self.ctrl.on_param_change(reason, params)
358    }
359}
360
361#[cfg(test)]
362mod tests {
363    use super::*;
364    use std::sync::atomic::{AtomicU32, Ordering};
365
366    static TEST_COUNTER: AtomicU32 = AtomicU32::new(0);
367
368    fn temp_path(ext: &str) -> PathBuf {
369        let n = TEST_COUNTER.fetch_add(1, Ordering::Relaxed);
370        std::env::temp_dir().join(format!("adcore_test_magick_{n}.{ext}"))
371    }
372
373    #[test]
374    fn test_write_read_png_u8() {
375        let path = temp_path("png");
376        let mut writer = MagickWriter::new();
377
378        let mut arr = NDArray::new(
379            vec![NDDimension::new(8), NDDimension::new(8)],
380            NDDataType::UInt8,
381        );
382        if let NDDataBuffer::U8(ref mut v) = arr.data {
383            for i in 0..64 {
384                v[i] = (i * 4) as u8;
385            }
386        }
387
388        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
389        writer.write_file(&arr).unwrap();
390
391        let read_back = writer.read_file().unwrap();
392        assert_eq!(read_back.data.data_type(), NDDataType::UInt8);
393        if let (NDDataBuffer::U8(orig), NDDataBuffer::U8(read)) = (&arr.data, &read_back.data) {
394            assert_eq!(orig, read);
395        }
396
397        writer.close_file().unwrap();
398        std::fs::remove_file(&path).ok();
399    }
400
401    #[test]
402    fn test_write_read_png_u16() {
403        let path = temp_path("png");
404        let mut writer = MagickWriter::new();
405
406        let mut arr = NDArray::new(
407            vec![NDDimension::new(8), NDDimension::new(8)],
408            NDDataType::UInt16,
409        );
410        if let NDDataBuffer::U16(ref mut v) = arr.data {
411            for i in 0..64 {
412                v[i] = (i * 1000) as u16;
413            }
414        }
415
416        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
417        writer.write_file(&arr).unwrap();
418
419        let read_back = writer.read_file().unwrap();
420        assert_eq!(read_back.data.data_type(), NDDataType::UInt16);
421        if let (NDDataBuffer::U16(orig), NDDataBuffer::U16(read)) = (&arr.data, &read_back.data) {
422            assert_eq!(orig, read);
423        }
424
425        writer.close_file().unwrap();
426        std::fs::remove_file(&path).ok();
427    }
428
429    #[test]
430    fn test_write_read_bmp_rgb() {
431        use ad_core_rs::attributes::{NDAttrSource, NDAttrValue, NDAttribute};
432
433        let path = temp_path("bmp");
434        let mut writer = MagickWriter::new();
435
436        let mut arr = NDArray::new(
437            vec![
438                NDDimension::new(3),
439                NDDimension::new(4),
440                NDDimension::new(4),
441            ],
442            NDDataType::UInt8,
443        );
444        arr.attributes.add(NDAttribute {
445            name: "ColorMode".into(),
446            description: "Color Mode".into(),
447            source: NDAttrSource::Driver,
448            value: NDAttrValue::Int32(2), // RGB1
449        });
450        if let NDDataBuffer::U8(ref mut v) = arr.data {
451            for i in 0..48 {
452                v[i] = (i * 5) as u8;
453            }
454        }
455
456        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
457        writer.write_file(&arr).unwrap();
458
459        let read_back = writer.read_file().unwrap();
460        assert_eq!(read_back.dims.len(), 3);
461        assert_eq!(read_back.dims[0].size, 3);
462
463        writer.close_file().unwrap();
464        std::fs::remove_file(&path).ok();
465    }
466
467    #[test]
468    fn test_rejects_unsupported_type() {
469        // F32 is now supported (normalized to U16). Use Float64 as unsupported.
470        let arr = NDArray::new(
471            vec![NDDimension::new(4), NDDimension::new(4)],
472            NDDataType::Float64,
473        );
474        assert!(MagickWriter::array_to_image(&arr).is_err());
475    }
476}
ad_plugins_rs/file_magick.rs

ad_plugins_rs/
file_magick.rs
