Skip to main content

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::U16(v) => {
91                if is_rgb {
92                    image::ImageBuffer::<image::Rgb<u16>, Vec<u16>>::from_raw(
93                        width,
94                        height,
95                        v.clone(),
96                    )
97                    .map(DynamicImage::ImageRgb16)
98                    .ok_or_else(|| {
99                        ADError::UnsupportedConversion("RGB16 buffer size mismatch".into())
100                    })
101                } else {
102                    image::ImageBuffer::<image::Luma<u16>, Vec<u16>>::from_raw(
103                        width,
104                        height,
105                        v.clone(),
106                    )
107                    .map(DynamicImage::ImageLuma16)
108                    .ok_or_else(|| {
109                        ADError::UnsupportedConversion("Gray16 buffer size mismatch".into())
110                    })
111                }
112            }
113            _ => Err(ADError::UnsupportedConversion(format!(
114                "NDFileMagick: unsupported data type {:?}, use UInt8 or UInt16",
115                src.data.data_type()
116            ))),
117        }
118    }
119}
120
121impl NDFileWriter for MagickWriter {
122    fn open_file(&mut self, path: &Path, _mode: NDFileMode, _array: &NDArray) -> ADResult<()> {
123        self.current_path = Some(path.to_path_buf());
124        Ok(())
125    }
126
127    fn write_file(&mut self, array: &NDArray) -> ADResult<()> {
128        let path = self
129            .current_path
130            .as_ref()
131            .ok_or_else(|| ADError::UnsupportedConversion("no file open".into()))?;
132
133        let img = Self::array_to_image(array)?;
134
135        // Determine format from extension, default to PNG
136        let format = ImageFormat::from_path(path).unwrap_or(ImageFormat::Png);
137
138        // For JPEG, use quality setting
139        if format == ImageFormat::Jpeg {
140            let mut buf = Vec::new();
141            let encoder =
142                image::codecs::jpeg::JpegEncoder::new_with_quality(&mut buf, self.quality);
143            img.write_with_encoder(encoder)
144                .map_err(|e| ADError::UnsupportedConversion(format!("Magick encode error: {e}")))?;
145            std::fs::write(path, &buf)?;
146        } else {
147            img.save(path)
148                .map_err(|e| ADError::UnsupportedConversion(format!("Magick save error: {e}")))?;
149        }
150
151        Ok(())
152    }
153
154    fn read_file(&mut self) -> ADResult<NDArray> {
155        let path = self
156            .current_path
157            .as_ref()
158            .ok_or_else(|| ADError::UnsupportedConversion("no file open".into()))?;
159
160        let img = image::open(path)
161            .map_err(|e| ADError::UnsupportedConversion(format!("Magick read error: {e}")))?;
162
163        let width = img.width() as usize;
164        let height = img.height() as usize;
165
166        match img {
167            DynamicImage::ImageLuma8(buf) => {
168                let mut arr = NDArray::new(
169                    vec![NDDimension::new(width), NDDimension::new(height)],
170                    NDDataType::UInt8,
171                );
172                arr.data = NDDataBuffer::U8(buf.into_raw());
173                Ok(arr)
174            }
175            DynamicImage::ImageRgb8(buf) => {
176                let mut arr = NDArray::new(
177                    vec![
178                        NDDimension::new(3),
179                        NDDimension::new(width),
180                        NDDimension::new(height),
181                    ],
182                    NDDataType::UInt8,
183                );
184                arr.data = NDDataBuffer::U8(buf.into_raw());
185                Ok(arr)
186            }
187            DynamicImage::ImageLuma16(buf) => {
188                let mut arr = NDArray::new(
189                    vec![NDDimension::new(width), NDDimension::new(height)],
190                    NDDataType::UInt16,
191                );
192                arr.data = NDDataBuffer::U16(buf.into_raw());
193                Ok(arr)
194            }
195            DynamicImage::ImageRgb16(buf) => {
196                let mut arr = NDArray::new(
197                    vec![
198                        NDDimension::new(3),
199                        NDDimension::new(width),
200                        NDDimension::new(height),
201                    ],
202                    NDDataType::UInt16,
203                );
204                arr.data = NDDataBuffer::U16(buf.into_raw());
205                Ok(arr)
206            }
207            other => {
208                // Convert anything else to RGB8
209                let rgb = other.to_rgb8();
210                let mut arr = NDArray::new(
211                    vec![
212                        NDDimension::new(3),
213                        NDDimension::new(width),
214                        NDDimension::new(height),
215                    ],
216                    NDDataType::UInt8,
217                );
218                arr.data = NDDataBuffer::U8(rgb.into_raw());
219                Ok(arr)
220            }
221        }
222    }
223
224    fn close_file(&mut self) -> ADResult<()> {
225        self.current_path = None;
226        Ok(())
227    }
228
229    fn supports_multiple_arrays(&self) -> bool {
230        false
231    }
232}
233
234/// Magick file processor wrapping FilePluginController<MagickWriter>.
235pub struct MagickFileProcessor {
236    ctrl: FilePluginController<MagickWriter>,
237    quality_idx: Option<usize>,
238    bit_depth_idx: Option<usize>,
239    compress_type_idx: Option<usize>,
240}
241
242impl MagickFileProcessor {
243    pub fn new() -> Self {
244        Self {
245            ctrl: FilePluginController::new(MagickWriter::new()),
246            quality_idx: None,
247            bit_depth_idx: None,
248            compress_type_idx: None,
249        }
250    }
251}
252
253impl NDPluginProcess for MagickFileProcessor {
254    fn process_array(&mut self, array: &NDArray, _pool: &NDArrayPool) -> ProcessResult {
255        self.ctrl.process_array(array)
256    }
257
258    fn plugin_type(&self) -> &str {
259        "NDFileMagick"
260    }
261
262    fn register_params(
263        &mut self,
264        base: &mut asyn_rs::port::PortDriverBase,
265    ) -> asyn_rs::error::AsynResult<()> {
266        self.ctrl.register_params(base)?;
267        use asyn_rs::param::ParamType;
268        self.quality_idx = Some(base.create_param("MAGICK_QUALITY", ParamType::Int32)?);
269        self.bit_depth_idx = Some(base.create_param("MAGICK_BIT_DEPTH", ParamType::Int32)?);
270        self.compress_type_idx = Some(base.create_param("MAGICK_COMPRESS_TYPE", ParamType::Int32)?);
271        // Set defaults
272        base.set_int32_param(self.quality_idx.unwrap(), 0, 100)?;
273        base.set_int32_param(self.bit_depth_idx.unwrap(), 0, 8)?;
274        base.set_int32_param(self.compress_type_idx.unwrap(), 0, 0)?;
275        Ok(())
276    }
277
278    fn on_param_change(
279        &mut self,
280        reason: usize,
281        params: &PluginParamSnapshot,
282    ) -> ParamChangeResult {
283        if Some(reason) == self.quality_idx {
284            let q = params.value.as_i32().clamp(1, 100) as u8;
285            self.ctrl.writer.set_quality(q);
286            return ParamChangeResult::empty();
287        }
288        if Some(reason) == self.bit_depth_idx {
289            let d = params.value.as_i32() as u32;
290            self.ctrl.writer.set_bit_depth(d);
291            return ParamChangeResult::empty();
292        }
293        if Some(reason) == self.compress_type_idx {
294            // CompressType stored but not actively used by `image` crate;
295            // format-specific compression is handled by each codec internally.
296            return ParamChangeResult::empty();
297        }
298        self.ctrl.on_param_change(reason, params)
299    }
300}
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305    use std::sync::atomic::{AtomicU32, Ordering};
306
307    static TEST_COUNTER: AtomicU32 = AtomicU32::new(0);
308
309    fn temp_path(ext: &str) -> PathBuf {
310        let n = TEST_COUNTER.fetch_add(1, Ordering::Relaxed);
311        std::env::temp_dir().join(format!("adcore_test_magick_{n}.{ext}"))
312    }
313
314    #[test]
315    fn test_write_read_png_u8() {
316        let path = temp_path("png");
317        let mut writer = MagickWriter::new();
318
319        let mut arr = NDArray::new(
320            vec![NDDimension::new(8), NDDimension::new(8)],
321            NDDataType::UInt8,
322        );
323        if let NDDataBuffer::U8(ref mut v) = arr.data {
324            for i in 0..64 {
325                v[i] = (i * 4) as u8;
326            }
327        }
328
329        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
330        writer.write_file(&arr).unwrap();
331
332        let read_back = writer.read_file().unwrap();
333        assert_eq!(read_back.data.data_type(), NDDataType::UInt8);
334        if let (NDDataBuffer::U8(orig), NDDataBuffer::U8(read)) = (&arr.data, &read_back.data) {
335            assert_eq!(orig, read);
336        }
337
338        writer.close_file().unwrap();
339        std::fs::remove_file(&path).ok();
340    }
341
342    #[test]
343    fn test_write_read_png_u16() {
344        let path = temp_path("png");
345        let mut writer = MagickWriter::new();
346
347        let mut arr = NDArray::new(
348            vec![NDDimension::new(8), NDDimension::new(8)],
349            NDDataType::UInt16,
350        );
351        if let NDDataBuffer::U16(ref mut v) = arr.data {
352            for i in 0..64 {
353                v[i] = (i * 1000) as u16;
354            }
355        }
356
357        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
358        writer.write_file(&arr).unwrap();
359
360        let read_back = writer.read_file().unwrap();
361        assert_eq!(read_back.data.data_type(), NDDataType::UInt16);
362        if let (NDDataBuffer::U16(orig), NDDataBuffer::U16(read)) = (&arr.data, &read_back.data) {
363            assert_eq!(orig, read);
364        }
365
366        writer.close_file().unwrap();
367        std::fs::remove_file(&path).ok();
368    }
369
370    #[test]
371    fn test_write_read_bmp_rgb() {
372        let path = temp_path("bmp");
373        let mut writer = MagickWriter::new();
374
375        let mut arr = NDArray::new(
376            vec![
377                NDDimension::new(3),
378                NDDimension::new(4),
379                NDDimension::new(4),
380            ],
381            NDDataType::UInt8,
382        );
383        if let NDDataBuffer::U8(ref mut v) = arr.data {
384            for i in 0..48 {
385                v[i] = (i * 5) as u8;
386            }
387        }
388
389        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
390        writer.write_file(&arr).unwrap();
391
392        let read_back = writer.read_file().unwrap();
393        assert_eq!(read_back.dims.len(), 3);
394        assert_eq!(read_back.dims[0].size, 3);
395
396        writer.close_file().unwrap();
397        std::fs::remove_file(&path).ok();
398    }
399
400    #[test]
401    fn test_rejects_unsupported_type() {
402        let arr = NDArray::new(
403            vec![NDDimension::new(4), NDDimension::new(4)],
404            NDDataType::Float32,
405        );
406        assert!(MagickWriter::array_to_image(&arr).is_err());
407    }
408}