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::codecs::png::{CompressionType as PngCompression, FilterType as PngFilter};
14use image::{DynamicImage, ImageEncoder, ImageFormat};
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum MagickCompression {
20 None = 0,
21 BZip = 1,
22 Fax = 2,
23 Group4 = 3,
24 Jpeg = 4,
25 Lzw = 5,
26 Rle = 6,
27 Zip = 7,
28}
29
30impl MagickCompression {
31 fn from_index(idx: i32) -> Self {
32 match idx {
33 1 => Self::BZip,
34 2 => Self::Fax,
35 3 => Self::Group4,
36 4 => Self::Jpeg,
37 5 => Self::Lzw,
38 6 => Self::Rle,
39 7 => Self::Zip,
40 _ => Self::None,
41 }
42 }
43}
44
45pub struct MagickWriter {
50 current_path: Option<PathBuf>,
51 quality: u8,
52 bit_depth: u32,
53 compress_type: MagickCompression,
54}
55
56impl MagickWriter {
57 pub fn new() -> Self {
58 Self {
59 current_path: None,
60 quality: 100,
61 bit_depth: 0,
65 compress_type: MagickCompression::None,
66 }
67 }
68
69 pub fn set_quality(&mut self, q: u8) {
70 self.quality = q;
71 }
72
73 pub fn set_bit_depth(&mut self, depth: u32) {
74 self.bit_depth = depth;
75 }
76
77 pub fn set_compress_type(&mut self, idx: i32) {
78 self.compress_type = MagickCompression::from_index(idx);
79 }
80
81 fn color_mode(array: &NDArray) -> NDColorMode {
82 array
83 .attributes
84 .get("ColorMode")
85 .and_then(|attr| attr.value.as_i64())
86 .map(|v| NDColorMode::from_i32(v as i32))
87 .unwrap_or_else(|| match array.dims.as_slice() {
88 [a, _, _] if a.size == 3 => NDColorMode::RGB1,
89 [_, b, _] if b.size == 3 => NDColorMode::RGB2,
90 [_, _, c] if c.size == 3 => NDColorMode::RGB3,
91 _ => NDColorMode::Mono,
92 })
93 }
94
95 fn array_to_image(array: &NDArray, bit_depth: u32) -> ADResult<DynamicImage> {
101 let img = Self::array_to_image_native(array)?;
102 Ok(Self::apply_bit_depth(img, bit_depth))
103 }
104
105 fn apply_bit_depth(img: DynamicImage, bit_depth: u32) -> DynamicImage {
107 if bit_depth == 0 {
108 return img;
110 }
111 let is_rgb = matches!(
112 img,
113 DynamicImage::ImageRgb8(_) | DynamicImage::ImageRgb16(_)
114 );
115 if bit_depth <= 8 {
116 if is_rgb {
117 DynamicImage::ImageRgb8(img.to_rgb8())
118 } else {
119 DynamicImage::ImageLuma8(img.to_luma8())
120 }
121 } else {
122 if is_rgb {
123 DynamicImage::ImageRgb16(img.to_rgb16())
124 } else {
125 DynamicImage::ImageLuma16(img.to_luma16())
126 }
127 }
128 }
129
130 fn array_to_image_native(array: &NDArray) -> ADResult<DynamicImage> {
132 let info = array.info();
133 let width = info.x_size as u32;
134 let height = info.y_size as u32;
135 let color = Self::color_mode(array);
136 let is_rgb = matches!(
137 color,
138 NDColorMode::RGB1 | NDColorMode::RGB2 | NDColorMode::RGB3
139 );
140
141 let src = if is_rgb && color != NDColorMode::RGB1 {
143 &convert_rgb_layout(array, color, NDColorMode::RGB1)?
144 } else {
145 array
146 };
147
148 match &src.data {
149 NDDataBuffer::U8(v) => {
150 if is_rgb {
151 image::RgbImage::from_raw(width, height, v.clone())
152 .map(DynamicImage::ImageRgb8)
153 .ok_or_else(|| {
154 ADError::UnsupportedConversion("RGB8 buffer size mismatch".into())
155 })
156 } else {
157 image::GrayImage::from_raw(width, height, v.clone())
158 .map(DynamicImage::ImageLuma8)
159 .ok_or_else(|| {
160 ADError::UnsupportedConversion("Gray8 buffer size mismatch".into())
161 })
162 }
163 }
164 NDDataBuffer::I8(v) => {
165 let u8_data: Vec<u8> = v.iter().map(|&b| b as u8).collect();
166 if is_rgb {
167 image::RgbImage::from_raw(width, height, u8_data)
168 .map(DynamicImage::ImageRgb8)
169 .ok_or_else(|| {
170 ADError::UnsupportedConversion("RGB8 buffer size mismatch".into())
171 })
172 } else {
173 image::GrayImage::from_raw(width, height, u8_data)
174 .map(DynamicImage::ImageLuma8)
175 .ok_or_else(|| {
176 ADError::UnsupportedConversion("Gray8 buffer size mismatch".into())
177 })
178 }
179 }
180 NDDataBuffer::U16(v) => {
181 if is_rgb {
182 image::ImageBuffer::<image::Rgb<u16>, Vec<u16>>::from_raw(
183 width,
184 height,
185 v.clone(),
186 )
187 .map(DynamicImage::ImageRgb16)
188 .ok_or_else(|| {
189 ADError::UnsupportedConversion("RGB16 buffer size mismatch".into())
190 })
191 } else {
192 image::ImageBuffer::<image::Luma<u16>, Vec<u16>>::from_raw(
193 width,
194 height,
195 v.clone(),
196 )
197 .map(DynamicImage::ImageLuma16)
198 .ok_or_else(|| {
199 ADError::UnsupportedConversion("Gray16 buffer size mismatch".into())
200 })
201 }
202 }
203 NDDataBuffer::I16(v) => {
204 let u16_data: Vec<u16> = v.iter().map(|&b| b as u16).collect();
205 if is_rgb {
206 image::ImageBuffer::<image::Rgb<u16>, Vec<u16>>::from_raw(
207 width, height, u16_data,
208 )
209 .map(DynamicImage::ImageRgb16)
210 .ok_or_else(|| {
211 ADError::UnsupportedConversion("RGB16 buffer size mismatch".into())
212 })
213 } else {
214 image::ImageBuffer::<image::Luma<u16>, Vec<u16>>::from_raw(
215 width, height, u16_data,
216 )
217 .map(DynamicImage::ImageLuma16)
218 .ok_or_else(|| {
219 ADError::UnsupportedConversion("Gray16 buffer size mismatch".into())
220 })
221 }
222 }
223 NDDataBuffer::F32(v) => {
224 let mut min = f32::INFINITY;
227 let mut max = f32::NEG_INFINITY;
228 for &f in v {
229 if f.is_finite() {
230 min = min.min(f);
231 max = max.max(f);
232 }
233 }
234 let range = if min.is_finite() && max > min {
235 max - min
236 } else {
237 1.0
238 };
239 let offset = if min.is_finite() { min } else { 0.0 };
240 let u16_data: Vec<u16> = v
241 .iter()
242 .map(|&f| {
243 let norm = ((f - offset) / range).clamp(0.0, 1.0);
244 (norm * 65535.0).round() as u16
245 })
246 .collect();
247 if is_rgb {
248 image::ImageBuffer::<image::Rgb<u16>, Vec<u16>>::from_raw(
249 width, height, u16_data,
250 )
251 .map(DynamicImage::ImageRgb16)
252 .ok_or_else(|| {
253 ADError::UnsupportedConversion("RGB16 buffer size mismatch".into())
254 })
255 } else {
256 image::ImageBuffer::<image::Luma<u16>, Vec<u16>>::from_raw(
257 width, height, u16_data,
258 )
259 .map(DynamicImage::ImageLuma16)
260 .ok_or_else(|| {
261 ADError::UnsupportedConversion("Gray16 buffer size mismatch".into())
262 })
263 }
264 }
265 _ => Err(ADError::UnsupportedConversion(format!(
266 "NDFileMagick: unsupported data type {:?}, use UInt8, Int8, UInt16, Int16, or Float32",
267 src.data.data_type()
268 ))),
269 }
270 }
271}
272
273impl NDFileWriter for MagickWriter {
274 fn open_file(&mut self, path: &Path, _mode: NDFileMode, _array: &NDArray) -> ADResult<()> {
275 self.current_path = Some(path.to_path_buf());
276 Ok(())
277 }
278
279 fn write_file(&mut self, array: &NDArray) -> ADResult<()> {
280 let path = self
281 .current_path
282 .as_ref()
283 .ok_or_else(|| ADError::UnsupportedConversion("no file open".into()))?;
284
285 let img = Self::array_to_image(array, self.bit_depth)?;
286
287 let format = ImageFormat::from_path(path).unwrap_or(ImageFormat::Png);
289
290 match format {
291 ImageFormat::Jpeg => {
292 let mut buf = Vec::new();
294 let encoder =
295 image::codecs::jpeg::JpegEncoder::new_with_quality(&mut buf, self.quality);
296 img.write_with_encoder(encoder).map_err(|e| {
297 ADError::UnsupportedConversion(format!("Magick encode error: {e}"))
298 })?;
299 std::fs::write(path, &buf)?;
300 }
301 ImageFormat::Png => {
302 let compression = match self.compress_type {
306 MagickCompression::None => PngCompression::Uncompressed,
307 MagickCompression::Zip | MagickCompression::BZip => PngCompression::Best,
308 _ => PngCompression::default(),
309 };
310 let mut buf = Vec::new();
311 let encoder = image::codecs::png::PngEncoder::new_with_quality(
312 &mut buf,
313 compression,
314 PngFilter::Adaptive,
315 );
316 let rgb = img.color();
317 encoder
318 .write_image(img.as_bytes(), img.width(), img.height(), rgb.into())
319 .map_err(|e| {
320 ADError::UnsupportedConversion(format!("Magick PNG encode error: {e}"))
321 })?;
322 std::fs::write(path, &buf)?;
323 }
324 _ => {
325 img.save(path).map_err(|e| {
330 ADError::UnsupportedConversion(format!("Magick save error: {e}"))
331 })?;
332 }
333 }
334
335 Ok(())
336 }
337
338 fn read_file(&mut self) -> ADResult<NDArray> {
339 let path = self
340 .current_path
341 .as_ref()
342 .ok_or_else(|| ADError::UnsupportedConversion("no file open".into()))?;
343
344 let img = image::open(path)
345 .map_err(|e| ADError::UnsupportedConversion(format!("Magick read error: {e}")))?;
346
347 let width = img.width() as usize;
348 let height = img.height() as usize;
349
350 match img {
351 DynamicImage::ImageLuma8(buf) => {
352 let mut arr = NDArray::new(
353 vec![NDDimension::new(width), NDDimension::new(height)],
354 NDDataType::UInt8,
355 );
356 arr.data = NDDataBuffer::U8(buf.into_raw());
357 Ok(arr)
358 }
359 DynamicImage::ImageRgb8(buf) => {
360 let mut arr = NDArray::new(
361 vec![
362 NDDimension::new(3),
363 NDDimension::new(width),
364 NDDimension::new(height),
365 ],
366 NDDataType::UInt8,
367 );
368 arr.data = NDDataBuffer::U8(buf.into_raw());
369 Ok(arr)
370 }
371 DynamicImage::ImageLuma16(buf) => {
372 let mut arr = NDArray::new(
373 vec![NDDimension::new(width), NDDimension::new(height)],
374 NDDataType::UInt16,
375 );
376 arr.data = NDDataBuffer::U16(buf.into_raw());
377 Ok(arr)
378 }
379 DynamicImage::ImageRgb16(buf) => {
380 let mut arr = NDArray::new(
381 vec![
382 NDDimension::new(3),
383 NDDimension::new(width),
384 NDDimension::new(height),
385 ],
386 NDDataType::UInt16,
387 );
388 arr.data = NDDataBuffer::U16(buf.into_raw());
389 Ok(arr)
390 }
391 other => {
392 let rgb = other.to_rgb8();
394 let mut arr = NDArray::new(
395 vec![
396 NDDimension::new(3),
397 NDDimension::new(width),
398 NDDimension::new(height),
399 ],
400 NDDataType::UInt8,
401 );
402 arr.data = NDDataBuffer::U8(rgb.into_raw());
403 Ok(arr)
404 }
405 }
406 }
407
408 fn close_file(&mut self) -> ADResult<()> {
409 self.current_path = None;
410 Ok(())
411 }
412
413 fn supports_multiple_arrays(&self) -> bool {
414 false
415 }
416}
417
418pub struct MagickFileProcessor {
420 ctrl: FilePluginController<MagickWriter>,
421 quality_idx: Option<usize>,
422 bit_depth_idx: Option<usize>,
423 compress_type_idx: Option<usize>,
424}
425
426impl MagickFileProcessor {
427 pub fn new() -> Self {
428 Self {
429 ctrl: FilePluginController::new(MagickWriter::new()),
430 quality_idx: None,
431 bit_depth_idx: None,
432 compress_type_idx: None,
433 }
434 }
435}
436
437impl NDPluginProcess for MagickFileProcessor {
438 fn process_array(&mut self, array: &NDArray, _pool: &NDArrayPool) -> ProcessResult {
439 self.ctrl.process_array(array)
440 }
441
442 fn plugin_type(&self) -> &str {
443 "NDFileMagick"
444 }
445
446 fn register_params(
447 &mut self,
448 base: &mut asyn_rs::port::PortDriverBase,
449 ) -> asyn_rs::error::AsynResult<()> {
450 self.ctrl.register_params(base)?;
451 use asyn_rs::param::ParamType;
452 self.quality_idx = Some(base.create_param("MAGICK_QUALITY", ParamType::Int32)?);
453 self.bit_depth_idx = Some(base.create_param("MAGICK_BIT_DEPTH", ParamType::Int32)?);
454 self.compress_type_idx = Some(base.create_param("MAGICK_COMPRESS_TYPE", ParamType::Int32)?);
455 base.set_int32_param(self.quality_idx.unwrap(), 0, 100)?;
457 base.set_int32_param(self.bit_depth_idx.unwrap(), 0, 8)?;
458 base.set_int32_param(self.compress_type_idx.unwrap(), 0, 0)?;
459 Ok(())
460 }
461
462 fn on_param_change(
463 &mut self,
464 reason: usize,
465 params: &PluginParamSnapshot,
466 ) -> ParamChangeResult {
467 if Some(reason) == self.quality_idx {
468 let q = params.value.as_i32().clamp(1, 100) as u8;
469 self.ctrl.writer.set_quality(q);
470 return ParamChangeResult::empty();
471 }
472 if Some(reason) == self.bit_depth_idx {
473 let d = params.value.as_i32() as u32;
474 self.ctrl.writer.set_bit_depth(d);
475 return ParamChangeResult::empty();
476 }
477 if Some(reason) == self.compress_type_idx {
478 self.ctrl.writer.set_compress_type(params.value.as_i32());
479 return ParamChangeResult::empty();
480 }
481 self.ctrl.on_param_change(reason, params)
482 }
483}
484
485#[cfg(test)]
486mod tests {
487 use super::*;
488 use std::sync::atomic::{AtomicU32, Ordering};
489
490 static TEST_COUNTER: AtomicU32 = AtomicU32::new(0);
491
492 fn temp_path(ext: &str) -> PathBuf {
493 let n = TEST_COUNTER.fetch_add(1, Ordering::Relaxed);
494 std::env::temp_dir().join(format!("adcore_test_magick_{n}.{ext}"))
495 }
496
497 #[test]
498 fn test_write_read_png_u8() {
499 let path = temp_path("png");
500 let mut writer = MagickWriter::new();
501
502 let mut arr = NDArray::new(
503 vec![NDDimension::new(8), NDDimension::new(8)],
504 NDDataType::UInt8,
505 );
506 if let NDDataBuffer::U8(ref mut v) = arr.data {
507 for i in 0..64 {
508 v[i] = (i * 4) as u8;
509 }
510 }
511
512 writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
513 writer.write_file(&arr).unwrap();
514
515 let read_back = writer.read_file().unwrap();
516 assert_eq!(read_back.data.data_type(), NDDataType::UInt8);
517 if let (NDDataBuffer::U8(orig), NDDataBuffer::U8(read)) = (&arr.data, &read_back.data) {
518 assert_eq!(orig, read);
519 }
520
521 writer.close_file().unwrap();
522 std::fs::remove_file(&path).ok();
523 }
524
525 #[test]
526 fn test_write_read_png_u16() {
527 let path = temp_path("png");
528 let mut writer = MagickWriter::new();
529
530 let mut arr = NDArray::new(
531 vec![NDDimension::new(8), NDDimension::new(8)],
532 NDDataType::UInt16,
533 );
534 if let NDDataBuffer::U16(ref mut v) = arr.data {
535 for i in 0..64 {
536 v[i] = (i * 1000) as u16;
537 }
538 }
539
540 writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
541 writer.write_file(&arr).unwrap();
542
543 let read_back = writer.read_file().unwrap();
544 assert_eq!(read_back.data.data_type(), NDDataType::UInt16);
545 if let (NDDataBuffer::U16(orig), NDDataBuffer::U16(read)) = (&arr.data, &read_back.data) {
546 assert_eq!(orig, read);
547 }
548
549 writer.close_file().unwrap();
550 std::fs::remove_file(&path).ok();
551 }
552
553 #[test]
554 fn test_write_read_bmp_rgb() {
555 use ad_core_rs::attributes::{NDAttrSource, NDAttrValue, NDAttribute};
556
557 let path = temp_path("bmp");
558 let mut writer = MagickWriter::new();
559
560 let mut arr = NDArray::new(
561 vec![
562 NDDimension::new(3),
563 NDDimension::new(4),
564 NDDimension::new(4),
565 ],
566 NDDataType::UInt8,
567 );
568 arr.attributes.add(NDAttribute::new_static(
569 "ColorMode",
570 "Color Mode",
571 NDAttrSource::Driver,
572 NDAttrValue::Int32(2), ));
574 if let NDDataBuffer::U8(ref mut v) = arr.data {
575 for i in 0..48 {
576 v[i] = (i * 5) as u8;
577 }
578 }
579
580 writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
581 writer.write_file(&arr).unwrap();
582
583 let read_back = writer.read_file().unwrap();
584 assert_eq!(read_back.dims.len(), 3);
585 assert_eq!(read_back.dims[0].size, 3);
586
587 writer.close_file().unwrap();
588 std::fs::remove_file(&path).ok();
589 }
590
591 #[test]
592 fn test_rejects_unsupported_type() {
593 let arr = NDArray::new(
595 vec![NDDimension::new(4), NDDimension::new(4)],
596 NDDataType::Float64,
597 );
598 assert!(MagickWriter::array_to_image(&arr, 8).is_err());
599 }
600
601 #[test]
602 fn test_bit_depth_controls_output_depth() {
603 let mut arr = NDArray::new(
605 vec![NDDimension::new(4), NDDimension::new(4)],
606 NDDataType::UInt16,
607 );
608 if let NDDataBuffer::U16(ref mut v) = arr.data {
609 for (i, x) in v.iter_mut().enumerate() {
610 *x = (i * 4000) as u16;
611 }
612 }
613 let img8 = MagickWriter::array_to_image(&arr, 8).unwrap();
614 assert!(matches!(img8, DynamicImage::ImageLuma8(_)));
615 let img16 = MagickWriter::array_to_image(&arr, 16).unwrap();
616 assert!(matches!(img16, DynamicImage::ImageLuma16(_)));
617 }
618
619 #[test]
620 fn test_f32_scales_by_actual_range() {
621 let mut arr = NDArray::new(
623 vec![NDDimension::new(2), NDDimension::new(2)],
624 NDDataType::Float32,
625 );
626 if let NDDataBuffer::F32(ref mut v) = arr.data {
627 v[0] = 100.0;
628 v[1] = 200.0;
629 v[2] = 300.0;
630 v[3] = 400.0;
631 }
632 let img = MagickWriter::array_to_image(&arr, 16).unwrap();
633 if let DynamicImage::ImageLuma16(buf) = img {
634 let raw = buf.into_raw();
635 assert_eq!(raw[0], 0);
637 assert_eq!(raw[3], 65535);
638 assert!(raw[1] > 0 && raw[1] < raw[2]);
639 } else {
640 panic!("expected 16-bit luma image");
641 }
642 }
643
644 #[test]
645 fn test_compress_type_applied_to_png() {
646 let mut arr = NDArray::new(
649 vec![NDDimension::new(64), NDDimension::new(64)],
650 NDDataType::UInt8,
651 );
652 if let NDDataBuffer::U8(ref mut v) = arr.data {
653 for x in v.iter_mut() {
654 *x = 128; }
656 }
657
658 let path_none = temp_path("png");
659 let mut w_none = MagickWriter::new();
660 w_none.set_compress_type(0); w_none
662 .open_file(&path_none, NDFileMode::Single, &arr)
663 .unwrap();
664 w_none.write_file(&arr).unwrap();
665 w_none.close_file().unwrap();
666
667 let path_zip = temp_path("png");
668 let mut w_zip = MagickWriter::new();
669 w_zip.set_compress_type(7); w_zip
671 .open_file(&path_zip, NDFileMode::Single, &arr)
672 .unwrap();
673 w_zip.write_file(&arr).unwrap();
674 w_zip.close_file().unwrap();
675
676 let size_none = std::fs::metadata(&path_none).unwrap().len();
677 let size_zip = std::fs::metadata(&path_zip).unwrap().len();
678 assert!(
679 size_zip < size_none,
680 "Zip ({size_zip}) should be smaller than None ({size_none})"
681 );
682
683 std::fs::remove_file(&path_none).ok();
684 std::fs::remove_file(&path_zip).ok();
685 }
686}