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
15pub 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 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 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 let format = ImageFormat::from_path(path).unwrap_or(ImageFormat::Png);
137
138 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 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
234pub 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 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 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}