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) -> ADResult<NDColorMode> {
95 let attr_mode = array.info().color_mode;
96 Ok(match array.dims.as_slice() {
97 [_, _] => NDColorMode::Mono,
98 [c, _, _] if c.size == 3 && attr_mode == NDColorMode::RGB1 => NDColorMode::RGB1,
99 [_, c, _] if c.size == 3 && attr_mode == NDColorMode::RGB2 => NDColorMode::RGB2,
100 [_, _, c] if c.size == 3 && attr_mode == NDColorMode::RGB3 => NDColorMode::RGB3,
101 _ => {
102 return Err(ADError::InvalidDimensions(
103 "unsupported array structure".into(),
104 ));
105 }
106 })
107 }
108
109 fn array_to_image(array: &NDArray, bit_depth: u32) -> ADResult<DynamicImage> {
115 let img = Self::array_to_image_native(array)?;
116 Ok(Self::apply_bit_depth(img, bit_depth))
117 }
118
119 fn apply_bit_depth(img: DynamicImage, bit_depth: u32) -> DynamicImage {
121 if bit_depth == 0 {
122 return img;
124 }
125 let is_rgb = matches!(
126 img,
127 DynamicImage::ImageRgb8(_) | DynamicImage::ImageRgb16(_)
128 );
129 if bit_depth <= 8 {
130 if is_rgb {
131 DynamicImage::ImageRgb8(img.to_rgb8())
132 } else {
133 DynamicImage::ImageLuma8(img.to_luma8())
134 }
135 } else {
136 if is_rgb {
137 DynamicImage::ImageRgb16(img.to_rgb16())
138 } else {
139 DynamicImage::ImageLuma16(img.to_luma16())
140 }
141 }
142 }
143
144 fn array_to_image_native(array: &NDArray) -> ADResult<DynamicImage> {
146 let info = array.info();
147 let width = info.x_size as u32;
148 let height = info.y_size as u32;
149 let color = Self::color_mode(array)?;
150 let is_rgb = matches!(
151 color,
152 NDColorMode::RGB1 | NDColorMode::RGB2 | NDColorMode::RGB3
153 );
154
155 let src = if is_rgb && color != NDColorMode::RGB1 {
157 &convert_rgb_layout(array, color, NDColorMode::RGB1)?
158 } else {
159 array
160 };
161
162 match &src.data {
163 NDDataBuffer::U8(v) => {
164 if is_rgb {
165 image::RgbImage::from_raw(width, height, v.clone())
166 .map(DynamicImage::ImageRgb8)
167 .ok_or_else(|| {
168 ADError::UnsupportedConversion("RGB8 buffer size mismatch".into())
169 })
170 } else {
171 image::GrayImage::from_raw(width, height, v.clone())
172 .map(DynamicImage::ImageLuma8)
173 .ok_or_else(|| {
174 ADError::UnsupportedConversion("Gray8 buffer size mismatch".into())
175 })
176 }
177 }
178 NDDataBuffer::I8(v) => {
179 let u8_data: Vec<u8> = v.iter().map(|&b| b as u8).collect();
180 if is_rgb {
181 image::RgbImage::from_raw(width, height, u8_data)
182 .map(DynamicImage::ImageRgb8)
183 .ok_or_else(|| {
184 ADError::UnsupportedConversion("RGB8 buffer size mismatch".into())
185 })
186 } else {
187 image::GrayImage::from_raw(width, height, u8_data)
188 .map(DynamicImage::ImageLuma8)
189 .ok_or_else(|| {
190 ADError::UnsupportedConversion("Gray8 buffer size mismatch".into())
191 })
192 }
193 }
194 NDDataBuffer::U16(v) => {
195 if is_rgb {
196 image::ImageBuffer::<image::Rgb<u16>, Vec<u16>>::from_raw(
197 width,
198 height,
199 v.clone(),
200 )
201 .map(DynamicImage::ImageRgb16)
202 .ok_or_else(|| {
203 ADError::UnsupportedConversion("RGB16 buffer size mismatch".into())
204 })
205 } else {
206 image::ImageBuffer::<image::Luma<u16>, Vec<u16>>::from_raw(
207 width,
208 height,
209 v.clone(),
210 )
211 .map(DynamicImage::ImageLuma16)
212 .ok_or_else(|| {
213 ADError::UnsupportedConversion("Gray16 buffer size mismatch".into())
214 })
215 }
216 }
217 NDDataBuffer::I16(v) => {
218 let u16_data: Vec<u16> = v.iter().map(|&b| b as u16).collect();
219 if is_rgb {
220 image::ImageBuffer::<image::Rgb<u16>, Vec<u16>>::from_raw(
221 width, height, u16_data,
222 )
223 .map(DynamicImage::ImageRgb16)
224 .ok_or_else(|| {
225 ADError::UnsupportedConversion("RGB16 buffer size mismatch".into())
226 })
227 } else {
228 image::ImageBuffer::<image::Luma<u16>, Vec<u16>>::from_raw(
229 width, height, u16_data,
230 )
231 .map(DynamicImage::ImageLuma16)
232 .ok_or_else(|| {
233 ADError::UnsupportedConversion("Gray16 buffer size mismatch".into())
234 })
235 }
236 }
237 NDDataBuffer::F32(v) => {
238 let mut min = f32::INFINITY;
241 let mut max = f32::NEG_INFINITY;
242 for &f in v {
243 if f.is_finite() {
244 min = min.min(f);
245 max = max.max(f);
246 }
247 }
248 let range = if min.is_finite() && max > min {
249 max - min
250 } else {
251 1.0
252 };
253 let offset = if min.is_finite() { min } else { 0.0 };
254 let u16_data: Vec<u16> = v
255 .iter()
256 .map(|&f| {
257 let norm = ((f - offset) / range).clamp(0.0, 1.0);
258 (norm * 65535.0).round() as u16
259 })
260 .collect();
261 if is_rgb {
262 image::ImageBuffer::<image::Rgb<u16>, Vec<u16>>::from_raw(
263 width, height, u16_data,
264 )
265 .map(DynamicImage::ImageRgb16)
266 .ok_or_else(|| {
267 ADError::UnsupportedConversion("RGB16 buffer size mismatch".into())
268 })
269 } else {
270 image::ImageBuffer::<image::Luma<u16>, Vec<u16>>::from_raw(
271 width, height, u16_data,
272 )
273 .map(DynamicImage::ImageLuma16)
274 .ok_or_else(|| {
275 ADError::UnsupportedConversion("Gray16 buffer size mismatch".into())
276 })
277 }
278 }
279 _ => Err(ADError::UnsupportedConversion(format!(
280 "NDFileMagick: unsupported data type {:?}, use UInt8, Int8, UInt16, Int16, or Float32",
281 src.data.data_type()
282 ))),
283 }
284 }
285}
286
287impl NDFileWriter for MagickWriter {
288 fn open_file(&mut self, path: &Path, _mode: NDFileMode, _array: &NDArray) -> ADResult<()> {
289 self.current_path = Some(path.to_path_buf());
290 Ok(())
291 }
292
293 fn write_file(&mut self, array: &NDArray) -> ADResult<()> {
294 let path = self
295 .current_path
296 .as_ref()
297 .ok_or_else(|| ADError::UnsupportedConversion("no file open".into()))?;
298
299 let img = Self::array_to_image(array, self.bit_depth)?;
300
301 let format = ImageFormat::from_path(path).unwrap_or(ImageFormat::Png);
303
304 match format {
305 ImageFormat::Jpeg => {
306 let mut buf = Vec::new();
308 let encoder =
309 image::codecs::jpeg::JpegEncoder::new_with_quality(&mut buf, self.quality);
310 img.write_with_encoder(encoder).map_err(|e| {
311 ADError::UnsupportedConversion(format!("Magick encode error: {e}"))
312 })?;
313 std::fs::write(path, &buf)?;
314 }
315 ImageFormat::Png => {
316 let compression = match self.compress_type {
320 MagickCompression::None => PngCompression::Uncompressed,
321 MagickCompression::Zip | MagickCompression::BZip => PngCompression::Best,
322 _ => PngCompression::default(),
323 };
324 let mut buf = Vec::new();
325 let encoder = image::codecs::png::PngEncoder::new_with_quality(
326 &mut buf,
327 compression,
328 PngFilter::Adaptive,
329 );
330 let rgb = img.color();
331 encoder
332 .write_image(img.as_bytes(), img.width(), img.height(), rgb.into())
333 .map_err(|e| {
334 ADError::UnsupportedConversion(format!("Magick PNG encode error: {e}"))
335 })?;
336 std::fs::write(path, &buf)?;
337 }
338 _ => {
339 img.save(path).map_err(|e| {
344 ADError::UnsupportedConversion(format!("Magick save error: {e}"))
345 })?;
346 }
347 }
348
349 Ok(())
350 }
351
352 fn read_file(&mut self) -> ADResult<NDArray> {
353 let path = self
354 .current_path
355 .as_ref()
356 .ok_or_else(|| ADError::UnsupportedConversion("no file open".into()))?;
357
358 let img = image::open(path)
359 .map_err(|e| ADError::UnsupportedConversion(format!("Magick read error: {e}")))?;
360
361 let width = img.width() as usize;
362 let height = img.height() as usize;
363
364 match img {
365 DynamicImage::ImageLuma8(buf) => {
366 let mut arr = NDArray::new(
367 vec![NDDimension::new(width), NDDimension::new(height)],
368 NDDataType::UInt8,
369 );
370 arr.data = NDDataBuffer::U8(buf.into_raw());
371 Ok(arr)
372 }
373 DynamicImage::ImageRgb8(buf) => {
374 let mut arr = NDArray::new(
375 vec![
376 NDDimension::new(3),
377 NDDimension::new(width),
378 NDDimension::new(height),
379 ],
380 NDDataType::UInt8,
381 );
382 arr.data = NDDataBuffer::U8(buf.into_raw());
383 Ok(arr)
384 }
385 DynamicImage::ImageLuma16(buf) => {
386 let mut arr = NDArray::new(
387 vec![NDDimension::new(width), NDDimension::new(height)],
388 NDDataType::UInt16,
389 );
390 arr.data = NDDataBuffer::U16(buf.into_raw());
391 Ok(arr)
392 }
393 DynamicImage::ImageRgb16(buf) => {
394 let mut arr = NDArray::new(
395 vec![
396 NDDimension::new(3),
397 NDDimension::new(width),
398 NDDimension::new(height),
399 ],
400 NDDataType::UInt16,
401 );
402 arr.data = NDDataBuffer::U16(buf.into_raw());
403 Ok(arr)
404 }
405 other => {
406 let rgb = other.to_rgb8();
408 let mut arr = NDArray::new(
409 vec![
410 NDDimension::new(3),
411 NDDimension::new(width),
412 NDDimension::new(height),
413 ],
414 NDDataType::UInt8,
415 );
416 arr.data = NDDataBuffer::U8(rgb.into_raw());
417 Ok(arr)
418 }
419 }
420 }
421
422 fn close_file(&mut self) -> ADResult<()> {
423 self.current_path = None;
424 Ok(())
425 }
426
427 fn supports_multiple_arrays(&self) -> bool {
428 false
429 }
430}
431
432pub struct MagickFileProcessor {
434 ctrl: FilePluginController<MagickWriter>,
435 quality_idx: Option<usize>,
436 bit_depth_idx: Option<usize>,
437 compress_type_idx: Option<usize>,
438}
439
440impl MagickFileProcessor {
441 pub fn new() -> Self {
442 Self {
443 ctrl: FilePluginController::new(MagickWriter::new()),
444 quality_idx: None,
445 bit_depth_idx: None,
446 compress_type_idx: None,
447 }
448 }
449}
450
451impl NDPluginProcess for MagickFileProcessor {
452 fn process_array(&mut self, array: &NDArray, _pool: &NDArrayPool) -> ProcessResult {
453 self.ctrl.process_array(array)
454 }
455
456 fn plugin_type(&self) -> &str {
457 "NDFileMagick"
458 }
459
460 fn does_array_callbacks(&self) -> bool {
463 false
464 }
465
466 fn register_params(
467 &mut self,
468 base: &mut asyn_rs::port::PortDriverBase,
469 ) -> asyn_rs::error::AsynResult<()> {
470 self.ctrl.register_params(base)?;
471 use asyn_rs::param::ParamType;
472 self.quality_idx = Some(base.create_param("MAGICK_QUALITY", ParamType::Int32)?);
473 self.bit_depth_idx = Some(base.create_param("MAGICK_BIT_DEPTH", ParamType::Int32)?);
474 self.compress_type_idx = Some(base.create_param("MAGICK_COMPRESS_TYPE", ParamType::Int32)?);
475 base.set_int32_param(self.quality_idx.unwrap(), 0, 100)?;
477 base.set_int32_param(self.bit_depth_idx.unwrap(), 0, 8)?;
478 base.set_int32_param(self.compress_type_idx.unwrap(), 0, 0)?;
479 Ok(())
480 }
481
482 fn on_param_change(
483 &mut self,
484 reason: usize,
485 params: &PluginParamSnapshot,
486 ) -> ParamChangeResult {
487 if Some(reason) == self.quality_idx {
488 let q = params.value.as_i32().clamp(1, 100) as u8;
489 self.ctrl.writer.set_quality(q);
490 return ParamChangeResult::empty();
491 }
492 if Some(reason) == self.bit_depth_idx {
493 let d = params.value.as_i32() as u32;
494 self.ctrl.writer.set_bit_depth(d);
495 return ParamChangeResult::empty();
496 }
497 if Some(reason) == self.compress_type_idx {
498 self.ctrl.writer.set_compress_type(params.value.as_i32());
499 return ParamChangeResult::empty();
500 }
501 self.ctrl.on_param_change(reason, params)
502 }
503}
504
505#[cfg(test)]
506mod tests {
507 use super::*;
508 use std::sync::atomic::{AtomicU32, Ordering};
509
510 static TEST_COUNTER: AtomicU32 = AtomicU32::new(0);
511
512 fn temp_path(ext: &str) -> PathBuf {
513 let n = TEST_COUNTER.fetch_add(1, Ordering::Relaxed);
514 std::env::temp_dir().join(format!("adcore_test_magick_{n}.{ext}"))
515 }
516
517 #[test]
523 fn test_r8_75_3d_without_colormode_attribute_is_an_error() {
524 use ad_core_rs::attributes::{NDAttrSource, NDAttrValue, NDAttribute};
525
526 let rgb1_dims = || {
527 vec![
528 NDDimension::new(3),
529 NDDimension::new(4),
530 NDDimension::new(4),
531 ]
532 };
533
534 let arr = NDArray::new(rgb1_dims(), NDDataType::UInt8);
535 let path = temp_path("png");
536 let mut writer = MagickWriter::new();
537 writer
538 .open_file(&path, NDFileMode::Single, &arr)
539 .expect("open");
540 let err = writer.write_file(&arr).unwrap_err();
541 assert!(
542 matches!(err, ADError::InvalidDimensions(_)),
543 "3-D without ColorMode must be rejected, got {err:?}"
544 );
545 std::fs::remove_file(&path).ok();
546
547 let mut arr = NDArray::new(rgb1_dims(), NDDataType::UInt8);
549 arr.attributes.add(NDAttribute::new_static(
550 "ColorMode",
551 "",
552 NDAttrSource::Driver,
553 NDAttrValue::Int32(NDColorMode::RGB1 as i32),
554 ));
555 let path = temp_path("png");
556 let mut writer = MagickWriter::new();
557 writer
558 .open_file(&path, NDFileMode::Single, &arr)
559 .expect("open");
560 writer
561 .write_file(&arr)
562 .expect("3-D WITH ColorMode=RGB1 must still write");
563 assert!(path.exists());
564 std::fs::remove_file(&path).ok();
565 }
566
567 #[test]
568 fn test_write_read_png_u8() {
569 let path = temp_path("png");
570 let mut writer = MagickWriter::new();
571
572 let mut arr = NDArray::new(
573 vec![NDDimension::new(8), NDDimension::new(8)],
574 NDDataType::UInt8,
575 );
576 if let NDDataBuffer::U8(ref mut v) = arr.data {
577 for i in 0..64 {
578 v[i] = (i * 4) as u8;
579 }
580 }
581
582 writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
583 writer.write_file(&arr).unwrap();
584
585 let read_back = writer.read_file().unwrap();
586 assert_eq!(read_back.data.data_type(), NDDataType::UInt8);
587 if let (NDDataBuffer::U8(orig), NDDataBuffer::U8(read)) = (&arr.data, &read_back.data) {
588 assert_eq!(orig, read);
589 }
590
591 writer.close_file().unwrap();
592 std::fs::remove_file(&path).ok();
593 }
594
595 #[test]
596 fn test_write_read_png_u16() {
597 let path = temp_path("png");
598 let mut writer = MagickWriter::new();
599
600 let mut arr = NDArray::new(
601 vec![NDDimension::new(8), NDDimension::new(8)],
602 NDDataType::UInt16,
603 );
604 if let NDDataBuffer::U16(ref mut v) = arr.data {
605 for i in 0..64 {
606 v[i] = (i * 1000) as u16;
607 }
608 }
609
610 writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
611 writer.write_file(&arr).unwrap();
612
613 let read_back = writer.read_file().unwrap();
614 assert_eq!(read_back.data.data_type(), NDDataType::UInt16);
615 if let (NDDataBuffer::U16(orig), NDDataBuffer::U16(read)) = (&arr.data, &read_back.data) {
616 assert_eq!(orig, read);
617 }
618
619 writer.close_file().unwrap();
620 std::fs::remove_file(&path).ok();
621 }
622
623 #[test]
624 fn test_write_read_bmp_rgb() {
625 use ad_core_rs::attributes::{NDAttrSource, NDAttrValue, NDAttribute};
626
627 let path = temp_path("bmp");
628 let mut writer = MagickWriter::new();
629
630 let mut arr = NDArray::new(
631 vec![
632 NDDimension::new(3),
633 NDDimension::new(4),
634 NDDimension::new(4),
635 ],
636 NDDataType::UInt8,
637 );
638 arr.attributes.add(NDAttribute::new_static(
639 "ColorMode",
640 "Color Mode",
641 NDAttrSource::Driver,
642 NDAttrValue::Int32(2), ));
644 if let NDDataBuffer::U8(ref mut v) = arr.data {
645 for i in 0..48 {
646 v[i] = (i * 5) as u8;
647 }
648 }
649
650 writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
651 writer.write_file(&arr).unwrap();
652
653 let read_back = writer.read_file().unwrap();
654 assert_eq!(read_back.dims.len(), 3);
655 assert_eq!(read_back.dims[0].size, 3);
656
657 writer.close_file().unwrap();
658 std::fs::remove_file(&path).ok();
659 }
660
661 #[test]
662 fn test_rejects_unsupported_type() {
663 let arr = NDArray::new(
665 vec![NDDimension::new(4), NDDimension::new(4)],
666 NDDataType::Float64,
667 );
668 assert!(MagickWriter::array_to_image(&arr, 8).is_err());
669 }
670
671 #[test]
672 fn test_bit_depth_controls_output_depth() {
673 let mut arr = NDArray::new(
675 vec![NDDimension::new(4), NDDimension::new(4)],
676 NDDataType::UInt16,
677 );
678 if let NDDataBuffer::U16(ref mut v) = arr.data {
679 for (i, x) in v.iter_mut().enumerate() {
680 *x = (i * 4000) as u16;
681 }
682 }
683 let img8 = MagickWriter::array_to_image(&arr, 8).unwrap();
684 assert!(matches!(img8, DynamicImage::ImageLuma8(_)));
685 let img16 = MagickWriter::array_to_image(&arr, 16).unwrap();
686 assert!(matches!(img16, DynamicImage::ImageLuma16(_)));
687 }
688
689 #[test]
690 fn test_f32_scales_by_actual_range() {
691 let mut arr = NDArray::new(
693 vec![NDDimension::new(2), NDDimension::new(2)],
694 NDDataType::Float32,
695 );
696 if let NDDataBuffer::F32(ref mut v) = arr.data {
697 v[0] = 100.0;
698 v[1] = 200.0;
699 v[2] = 300.0;
700 v[3] = 400.0;
701 }
702 let img = MagickWriter::array_to_image(&arr, 16).unwrap();
703 if let DynamicImage::ImageLuma16(buf) = img {
704 let raw = buf.into_raw();
705 assert_eq!(raw[0], 0);
707 assert_eq!(raw[3], 65535);
708 assert!(raw[1] > 0 && raw[1] < raw[2]);
709 } else {
710 panic!("expected 16-bit luma image");
711 }
712 }
713
714 #[test]
715 fn test_compress_type_applied_to_png() {
716 let mut arr = NDArray::new(
719 vec![NDDimension::new(64), NDDimension::new(64)],
720 NDDataType::UInt8,
721 );
722 if let NDDataBuffer::U8(ref mut v) = arr.data {
723 for x in v.iter_mut() {
724 *x = 128; }
726 }
727
728 let path_none = temp_path("png");
729 let mut w_none = MagickWriter::new();
730 w_none.set_compress_type(0); w_none
732 .open_file(&path_none, NDFileMode::Single, &arr)
733 .unwrap();
734 w_none.write_file(&arr).unwrap();
735 w_none.close_file().unwrap();
736
737 let path_zip = temp_path("png");
738 let mut w_zip = MagickWriter::new();
739 w_zip.set_compress_type(7); w_zip
741 .open_file(&path_zip, NDFileMode::Single, &arr)
742 .unwrap();
743 w_zip.write_file(&arr).unwrap();
744 w_zip.close_file().unwrap();
745
746 let size_none = std::fs::metadata(&path_none).unwrap().len();
747 let size_zip = std::fs::metadata(&path_zip).unwrap().len();
748 assert!(
749 size_zip < size_none,
750 "Zip ({size_zip}) should be smaller than None ({size_none})"
751 );
752
753 std::fs::remove_file(&path_none).ok();
754 std::fs::remove_file(&path_zip).ok();
755 }
756}