Skip to main content

ad_plugins_rs/
file_jpeg.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 jpeg_encoder::{ColorType as JpegColorType, Encoder as JpegEncoder};
14
15/// JPEG file writer using `jpeg-encoder` for encoding and `jpeg-decoder` for decoding.
16pub struct JpegWriter {
17    current_path: Option<PathBuf>,
18    pub(crate) quality: u8,
19}
20
21impl JpegWriter {
22    pub fn new(quality: u8) -> Self {
23        Self {
24            current_path: None,
25            quality,
26        }
27    }
28
29    pub fn set_quality(&mut self, quality: u8) {
30        self.quality = quality;
31    }
32}
33
34impl NDFileWriter for JpegWriter {
35    fn open_file(&mut self, path: &Path, _mode: NDFileMode, array: &NDArray) -> ADResult<()> {
36        let dt = array.data.data_type();
37        if dt != NDDataType::UInt8 && dt != NDDataType::Int8 {
38            return Err(ADError::UnsupportedConversion(
39                "JPEG only supports UInt8/Int8 data".into(),
40            ));
41        }
42        self.current_path = Some(path.to_path_buf());
43        Ok(())
44    }
45
46    fn write_file(&mut self, array: &NDArray) -> ADResult<()> {
47        let path = self
48            .current_path
49            .as_ref()
50            .ok_or_else(|| ADError::UnsupportedConversion("no file open".into()))?;
51
52        // The ColorMode *attribute* is the only source of truth, defaulting to
53        // Mono when it is absent — C `int colorMode = NDColorModeMono;`
54        // (NDFileJPEG.cpp:29) overwritten only by the attribute (:52-53).
55        // `info().color_mode` is that rule's single owner.
56        let color_mode = array.info().color_mode;
57
58        // C's `openFile` structure chain (NDFileJPEG.cpp:55-84). Each 3-D branch
59        // requires the *attribute* to name the layout, so a 3-D array with no
60        // ColorMode attribute (colorMode stays Mono) matches nothing and C returns
61        // asynError (:79-84). Inferring the layout from the dimensions instead
62        // made such an array look like RGB1 and write a file.
63        //
64        // C also takes `this->colorMode` from the branch it took, not from the
65        // attribute: the 2-D branch forces Mono (:60). And unlike NDFileTIFF
66        // (:180) there is no ndims == 1 branch here — a 1-D array is an error.
67        let color_mode = match array.dims.as_slice() {
68            [_, _] => NDColorMode::Mono,
69            [c, _, _] if c.size == 3 && color_mode == NDColorMode::RGB1 => NDColorMode::RGB1,
70            [_, c, _] if c.size == 3 && color_mode == NDColorMode::RGB2 => NDColorMode::RGB2,
71            [_, _, c] if c.size == 3 && color_mode == NDColorMode::RGB3 => NDColorMode::RGB3,
72            _ => {
73                return Err(ADError::InvalidDimensions(
74                    "unsupported array structure".into(),
75                ));
76            }
77        };
78
79        let is_rgb = matches!(
80            color_mode,
81            NDColorMode::RGB1 | NDColorMode::RGB2 | NDColorMode::RGB3
82        );
83        let src = if is_rgb && color_mode != NDColorMode::RGB1 {
84            &convert_rgb_layout(array, color_mode, NDColorMode::RGB1)?
85        } else {
86            array
87        };
88
89        let info = src.info();
90        let width = info.x_size;
91        let height = info.y_size;
92
93        let data: Vec<u8> = match &src.data {
94            NDDataBuffer::U8(v) => v.clone(),
95            NDDataBuffer::I8(v) => v.iter().map(|&b| b as u8).collect(),
96            _ => {
97                return Err(ADError::UnsupportedConversion(
98                    "JPEG only supports UInt8/Int8".into(),
99                ));
100            }
101        };
102
103        let color_type = if info.color_size == 3 {
104            JpegColorType::Rgb
105        } else {
106            JpegColorType::Luma
107        };
108
109        let mut buf = Vec::new();
110        let encoder = JpegEncoder::new(&mut buf, self.quality);
111        encoder
112            .encode(&data, width as u16, height as u16, color_type)
113            .map_err(|e| ADError::UnsupportedConversion(format!("JPEG encode error: {}", e)))?;
114
115        std::fs::write(path, &buf)?;
116        Ok(())
117    }
118
119    fn read_file(&mut self) -> ADResult<NDArray> {
120        let path = self
121            .current_path
122            .as_ref()
123            .ok_or_else(|| ADError::UnsupportedConversion("no file open".into()))?;
124
125        let file_data = std::fs::read(path)?;
126        let mut decoder = jpeg_decoder::Decoder::new(&file_data[..]);
127        let pixels = decoder
128            .decode()
129            .map_err(|e| ADError::UnsupportedConversion(format!("JPEG decode error: {}", e)))?;
130        let info = decoder.info().unwrap();
131
132        let (width, height) = (info.width as usize, info.height as usize);
133
134        let dims = match info.pixel_format {
135            jpeg_decoder::PixelFormat::L8 => {
136                vec![NDDimension::new(width), NDDimension::new(height)]
137            }
138            jpeg_decoder::PixelFormat::RGB24 => {
139                vec![
140                    NDDimension::new(3),
141                    NDDimension::new(width),
142                    NDDimension::new(height),
143                ]
144            }
145            _ => {
146                return Err(ADError::UnsupportedConversion(
147                    "unsupported JPEG pixel format".into(),
148                ));
149            }
150        };
151
152        let mut arr = NDArray::new(dims, NDDataType::UInt8);
153        arr.data = NDDataBuffer::U8(pixels);
154        Ok(arr)
155    }
156
157    fn close_file(&mut self) -> ADResult<()> {
158        self.current_path = None;
159        Ok(())
160    }
161
162    fn supports_multiple_arrays(&self) -> bool {
163        false
164    }
165}
166
167/// JPEG file processor wrapping FilePluginController<JpegWriter>.
168pub struct JpegFileProcessor {
169    ctrl: FilePluginController<JpegWriter>,
170    jpeg_quality_idx: Option<usize>,
171}
172
173impl JpegFileProcessor {
174    pub fn new(quality: u8) -> Self {
175        Self {
176            ctrl: FilePluginController::new(JpegWriter::new(quality)),
177            jpeg_quality_idx: None,
178        }
179    }
180}
181
182impl Default for JpegFileProcessor {
183    fn default() -> Self {
184        Self::new(50)
185    }
186}
187
188impl NDPluginProcess for JpegFileProcessor {
189    fn process_array(&mut self, array: &NDArray, _pool: &NDArrayPool) -> ProcessResult {
190        self.ctrl.process_array(array)
191    }
192
193    fn plugin_type(&self) -> &str {
194        "NDFileJPEG"
195    }
196
197    /// C `NDPluginFile.cpp:948` (base of every file writer) sets
198    /// `NDArrayCallbacks = 0`: file plugins write to disk, not downstream.
199    fn does_array_callbacks(&self) -> bool {
200        false
201    }
202
203    fn register_params(
204        &mut self,
205        base: &mut asyn_rs::port::PortDriverBase,
206    ) -> asyn_rs::error::AsynResult<()> {
207        self.ctrl.register_params(base)?;
208        use asyn_rs::param::ParamType;
209        let idx = base.create_param("JPEG_QUALITY", ParamType::Int32)?;
210        // Seed the readback PV with the actual encoder default (C++ NDFileJPEG.cpp:327
211        // sets NDFileJPEGQuality default to 50). Without this the PV reads 0 while the
212        // encoder uses its constructed default, so PV and effective quality disagree.
213        base.set_int32_param(idx, 0, i32::from(self.ctrl.writer.quality))?;
214        self.jpeg_quality_idx = Some(idx);
215        Ok(())
216    }
217
218    fn on_param_change(
219        &mut self,
220        reason: usize,
221        params: &PluginParamSnapshot,
222    ) -> ParamChangeResult {
223        // JPEG-specific: quality change
224        if Some(reason) == self.jpeg_quality_idx {
225            let q = params.value.as_i32().clamp(1, 100) as u8;
226            self.ctrl.writer.set_quality(q);
227            return ParamChangeResult::empty();
228        }
229        self.ctrl.on_param_change(reason, params)
230    }
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236    use ad_core_rs::ndarray::{NDDataBuffer, NDDimension};
237    use std::sync::atomic::{AtomicU32, Ordering};
238
239    static TEST_COUNTER: AtomicU32 = AtomicU32::new(0);
240
241    fn temp_path(prefix: &str) -> PathBuf {
242        let n = TEST_COUNTER.fetch_add(1, Ordering::Relaxed);
243        std::env::temp_dir().join(format!("adcore_test_{}_{}.jpg", prefix, n))
244    }
245
246    /// R8-75, JPEG writer: same defect as the cited TIFF site. C defaults
247    /// `colorMode` to Mono (NDFileJPEG.cpp:29), overwrites it only from the
248    /// attribute (:52-53), and each 3-D branch requires the attribute (:61, :67,
249    /// :73) — so a 3-D array without ColorMode returns asynError (:79-84). The
250    /// port inferred RGB1 from the dims. A 1-D array is an error here too: C has
251    /// no ndims == 1 branch in this writer.
252    #[test]
253    fn test_r8_75_3d_without_colormode_attribute_is_an_error() {
254        use ad_core_rs::attributes::{NDAttrSource, NDAttrValue, NDAttribute};
255        use ad_core_rs::color::NDColorMode;
256
257        let rgb1_dims = || {
258            vec![
259                NDDimension::new(3),
260                NDDimension::new(4),
261                NDDimension::new(4),
262            ]
263        };
264
265        let arr = NDArray::new(rgb1_dims(), NDDataType::UInt8);
266        let path = temp_path("jpeg_3d_no_colormode");
267        let mut writer = JpegWriter::new(90);
268        writer
269            .open_file(&path, NDFileMode::Single, &arr)
270            .expect("open");
271        let err = writer.write_file(&arr).unwrap_err();
272        assert!(
273            matches!(err, ADError::InvalidDimensions(_)),
274            "3-D without ColorMode must be rejected, got {err:?}"
275        );
276        std::fs::remove_file(&path).ok();
277
278        // C has no ndims == 1 branch (unlike NDFileTIFF.cpp:180).
279        let arr = NDArray::new(vec![NDDimension::new(16)], NDDataType::UInt8);
280        let path = temp_path("jpeg_1d");
281        let mut writer = JpegWriter::new(90);
282        writer
283            .open_file(&path, NDFileMode::Single, &arr)
284            .expect("open");
285        assert!(matches!(
286            writer.write_file(&arr).unwrap_err(),
287            ADError::InvalidDimensions(_)
288        ));
289        std::fs::remove_file(&path).ok();
290
291        // Positive control: WITH ColorMode=RGB1 the same array writes.
292        let mut arr = NDArray::new(rgb1_dims(), NDDataType::UInt8);
293        arr.attributes.add(NDAttribute::new_static(
294            "ColorMode",
295            "",
296            NDAttrSource::Driver,
297            NDAttrValue::Int32(NDColorMode::RGB1 as i32),
298        ));
299        let path = temp_path("jpeg_3d_rgb1");
300        let mut writer = JpegWriter::new(90);
301        writer
302            .open_file(&path, NDFileMode::Single, &arr)
303            .expect("open");
304        writer
305            .write_file(&arr)
306            .expect("3-D WITH ColorMode=RGB1 must still write");
307        assert!(path.exists());
308        std::fs::remove_file(&path).ok();
309    }
310
311    #[test]
312    fn test_write_u8() {
313        let path = temp_path("jpeg");
314        let mut writer = JpegWriter::new(90);
315
316        let mut arr = NDArray::new(
317            vec![NDDimension::new(8), NDDimension::new(8)],
318            NDDataType::UInt8,
319        );
320        if let NDDataBuffer::U8(ref mut v) = arr.data {
321            for i in 0..64 {
322                v[i] = (i * 4) as u8;
323            }
324        }
325
326        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
327        writer.write_file(&arr).unwrap();
328        writer.close_file().unwrap();
329
330        let data = std::fs::read(&path).unwrap();
331        // Check JPEG SOI marker
332        assert_eq!(&data[0..2], &[0xFF, 0xD8]);
333        // Check JPEG EOI marker at end
334        assert_eq!(&data[data.len() - 2..], &[0xFF, 0xD9]);
335
336        std::fs::remove_file(&path).ok();
337    }
338
339    #[test]
340    fn test_rejects_non_u8() {
341        let path = temp_path("jpeg_u16");
342        let mut writer = JpegWriter::new(90);
343
344        let arr = NDArray::new(
345            vec![NDDimension::new(4), NDDimension::new(4)],
346            NDDataType::UInt16,
347        );
348
349        let result = writer.open_file(&path, NDFileMode::Single, &arr);
350        assert!(result.is_err());
351    }
352
353    #[test]
354    fn test_quality_affects_size() {
355        let path_high = temp_path("jpeg_hi");
356        let path_low = temp_path("jpeg_lo");
357
358        let mut arr = NDArray::new(
359            vec![NDDimension::new(32), NDDimension::new(32)],
360            NDDataType::UInt8,
361        );
362        if let NDDataBuffer::U8(ref mut v) = arr.data {
363            for i in 0..v.len() {
364                v[i] = (i % 256) as u8;
365            }
366        }
367
368        let mut writer_high = JpegWriter::new(95);
369        writer_high
370            .open_file(&path_high, NDFileMode::Single, &arr)
371            .unwrap();
372        writer_high.write_file(&arr).unwrap();
373        writer_high.close_file().unwrap();
374
375        let mut writer_low = JpegWriter::new(10);
376        writer_low
377            .open_file(&path_low, NDFileMode::Single, &arr)
378            .unwrap();
379        writer_low.write_file(&arr).unwrap();
380        writer_low.close_file().unwrap();
381
382        let size_high = std::fs::metadata(&path_high).unwrap().len();
383        let size_low = std::fs::metadata(&path_low).unwrap().len();
384        assert!(
385            size_high > size_low,
386            "high quality ({}) should be larger than low quality ({})",
387            size_high,
388            size_low
389        );
390
391        std::fs::remove_file(&path_high).ok();
392        std::fs::remove_file(&path_low).ok();
393    }
394
395    #[test]
396    fn test_default_quality_is_50() {
397        // C++ NDFileJPEG.cpp:327 default quality is 50.
398        assert_eq!(JpegFileProcessor::default().ctrl.writer.quality, 50);
399        assert_eq!(JpegWriter::new(50).quality, 50);
400    }
401
402    #[test]
403    fn test_register_params_seeds_quality_pv() {
404        use asyn_rs::port::{PortDriverBase, PortFlags};
405        let mut base = PortDriverBase::new("jpeg_param_test", 1, PortFlags::default());
406        let mut proc = JpegFileProcessor::new(50);
407        proc.register_params(&mut base).unwrap();
408        let idx = proc.jpeg_quality_idx.unwrap();
409        // Readback PV must equal the encoder's effective quality, not 0.
410        assert_eq!(base.get_int32_param(idx, 0).unwrap(), 50);
411    }
412
413    #[test]
414    fn test_roundtrip_luma() {
415        let path = temp_path("jpeg_rt");
416        let mut writer = JpegWriter::new(100);
417
418        let mut arr = NDArray::new(
419            vec![NDDimension::new(8), NDDimension::new(8)],
420            NDDataType::UInt8,
421        );
422        if let NDDataBuffer::U8(ref mut v) = arr.data {
423            // Use uniform value so JPEG compression is lossless at quality 100
424            for i in 0..64 {
425                v[i] = 128;
426            }
427        }
428
429        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
430        writer.write_file(&arr).unwrap();
431
432        let read_back = writer.read_file().unwrap();
433        assert_eq!(read_back.data.data_type(), NDDataType::UInt8);
434        if let NDDataBuffer::U8(ref v) = read_back.data {
435            // With uniform input at max quality, decoded values should be close
436            for &px in v.iter() {
437                assert!(
438                    (px as i16 - 128).unsigned_abs() < 5,
439                    "pixel {} too far from 128",
440                    px
441                );
442            }
443        }
444
445        writer.close_file().unwrap();
446        std::fs::remove_file(&path).ok();
447    }
448
449    /// Write an RGB-mode array to JPEG and return the decoded array's dims.
450    fn jpeg_roundtrip_dims(prefix: &str, mode: NDColorMode, dims: Vec<usize>) -> Vec<usize> {
451        use ad_core_rs::attributes::{NDAttrSource, NDAttrValue, NDAttribute};
452        let path = temp_path(prefix);
453        let mut writer = JpegWriter::new(95);
454
455        let mut arr = NDArray::new(
456            dims.iter().map(|&d| NDDimension::new(d)).collect(),
457            NDDataType::UInt8,
458        );
459        arr.attributes.add(NDAttribute::new_static(
460            "ColorMode",
461            "Color mode",
462            NDAttrSource::Driver,
463            NDAttrValue::Int32(mode as i32),
464        ));
465        if let NDDataBuffer::U8(ref mut v) = arr.data {
466            for (i, x) in v.iter_mut().enumerate() {
467                *x = (i % 256) as u8;
468            }
469        }
470
471        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
472        writer.write_file(&arr).unwrap();
473        let read_back = writer.read_file().unwrap();
474        writer.close_file().unwrap();
475        std::fs::remove_file(&path).ok();
476        read_back.dims.iter().map(|d| d.size).collect()
477    }
478
479    #[test]
480    fn test_adp12_rgb2_jpeg_written_as_rgb_not_grayscale() {
481        // RGB2 [x=5, c=3, y=4]: C writes width=5, height=4, 3 JCS_RGB
482        // components (NDFileJPEG.cpp:67-78). The Rust converts to RGB1 first;
483        // before the fix the stale ColorMode=RGB2 attribute made info() read
484        // the RGB1 dims as width=3, color=5 -> a 3x4 grayscale JPEG. Decoded
485        // dims must be RGB24 [3, 5, 4], not grayscale [3, 4].
486        let dims = jpeg_roundtrip_dims("jpeg_rgb2", NDColorMode::RGB2, vec![5, 3, 4]);
487        assert_eq!(dims, vec![3, 5, 4]);
488    }
489
490    #[test]
491    fn test_adp12_rgb3_jpeg_written_as_rgb_not_grayscale() {
492        // RGB3 [x=5, y=4, c=3]: C writes width=5, height=4, 3 components
493        // (NDFileJPEG.cpp:72-78). Decoded dims must be RGB24 [3, 5, 4].
494        let dims = jpeg_roundtrip_dims("jpeg_rgb3", NDColorMode::RGB3, vec![5, 4, 3]);
495        assert_eq!(dims, vec![3, 5, 4]);
496    }
497}