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