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        // Detect color mode and convert RGB2/RGB3 to RGB1 if needed
53        let color_mode = array
54            .attributes
55            .get("ColorMode")
56            .and_then(|attr| attr.value.as_i64())
57            .map(|v| NDColorMode::from_i32(v as i32))
58            .unwrap_or_else(|| {
59                if array.dims.len() == 3 {
60                    let d0 = array.dims[0].size;
61                    let d1 = array.dims[1].size;
62                    let d2 = array.dims[2].size;
63                    if d0 == 3 {
64                        NDColorMode::RGB1
65                    } else if d1 == 3 {
66                        NDColorMode::RGB2
67                    } else if d2 == 3 {
68                        NDColorMode::RGB3
69                    } else {
70                        NDColorMode::Mono
71                    }
72                } else {
73                    NDColorMode::Mono
74                }
75            });
76
77        let is_rgb = matches!(
78            color_mode,
79            NDColorMode::RGB1 | NDColorMode::RGB2 | NDColorMode::RGB3
80        );
81        let src = if is_rgb && color_mode != NDColorMode::RGB1 {
82            &convert_rgb_layout(array, color_mode, NDColorMode::RGB1)?
83        } else {
84            array
85        };
86
87        let info = src.info();
88        let width = info.x_size;
89        let height = info.y_size;
90
91        let data: Vec<u8> = match &src.data {
92            NDDataBuffer::U8(v) => v.clone(),
93            NDDataBuffer::I8(v) => v.iter().map(|&b| b as u8).collect(),
94            _ => {
95                return Err(ADError::UnsupportedConversion(
96                    "JPEG only supports UInt8/Int8".into(),
97                ));
98            }
99        };
100
101        let color_type = if info.color_size == 3 {
102            JpegColorType::Rgb
103        } else {
104            JpegColorType::Luma
105        };
106
107        let mut buf = Vec::new();
108        let encoder = JpegEncoder::new(&mut buf, self.quality);
109        encoder
110            .encode(&data, width as u16, height as u16, color_type)
111            .map_err(|e| ADError::UnsupportedConversion(format!("JPEG encode error: {}", e)))?;
112
113        std::fs::write(path, &buf)?;
114        Ok(())
115    }
116
117    fn read_file(&mut self) -> ADResult<NDArray> {
118        let path = self
119            .current_path
120            .as_ref()
121            .ok_or_else(|| ADError::UnsupportedConversion("no file open".into()))?;
122
123        let file_data = std::fs::read(path)?;
124        let mut decoder = jpeg_decoder::Decoder::new(&file_data[..]);
125        let pixels = decoder
126            .decode()
127            .map_err(|e| ADError::UnsupportedConversion(format!("JPEG decode error: {}", e)))?;
128        let info = decoder.info().unwrap();
129
130        let (width, height) = (info.width as usize, info.height as usize);
131
132        let dims = match info.pixel_format {
133            jpeg_decoder::PixelFormat::L8 => {
134                vec![NDDimension::new(width), NDDimension::new(height)]
135            }
136            jpeg_decoder::PixelFormat::RGB24 => {
137                vec![
138                    NDDimension::new(3),
139                    NDDimension::new(width),
140                    NDDimension::new(height),
141                ]
142            }
143            _ => {
144                return Err(ADError::UnsupportedConversion(
145                    "unsupported JPEG pixel format".into(),
146                ));
147            }
148        };
149
150        let mut arr = NDArray::new(dims, NDDataType::UInt8);
151        arr.data = NDDataBuffer::U8(pixels);
152        Ok(arr)
153    }
154
155    fn close_file(&mut self) -> ADResult<()> {
156        self.current_path = None;
157        Ok(())
158    }
159
160    fn supports_multiple_arrays(&self) -> bool {
161        false
162    }
163}
164
165/// JPEG file processor wrapping FilePluginController<JpegWriter>.
166pub struct JpegFileProcessor {
167    ctrl: FilePluginController<JpegWriter>,
168    jpeg_quality_idx: Option<usize>,
169}
170
171impl JpegFileProcessor {
172    pub fn new(quality: u8) -> Self {
173        Self {
174            ctrl: FilePluginController::new(JpegWriter::new(quality)),
175            jpeg_quality_idx: None,
176        }
177    }
178}
179
180impl Default for JpegFileProcessor {
181    fn default() -> Self {
182        Self::new(50)
183    }
184}
185
186impl NDPluginProcess for JpegFileProcessor {
187    fn process_array(&mut self, array: &NDArray, _pool: &NDArrayPool) -> ProcessResult {
188        self.ctrl.process_array(array)
189    }
190
191    fn plugin_type(&self) -> &str {
192        "NDFileJPEG"
193    }
194
195    fn register_params(
196        &mut self,
197        base: &mut asyn_rs::port::PortDriverBase,
198    ) -> asyn_rs::error::AsynResult<()> {
199        self.ctrl.register_params(base)?;
200        use asyn_rs::param::ParamType;
201        let idx = base.create_param("JPEG_QUALITY", ParamType::Int32)?;
202        // Seed the readback PV with the actual encoder default (C++ NDFileJPEG.cpp:327
203        // sets NDFileJPEGQuality default to 50). Without this the PV reads 0 while the
204        // encoder uses its constructed default, so PV and effective quality disagree.
205        base.set_int32_param(idx, 0, i32::from(self.ctrl.writer.quality))?;
206        self.jpeg_quality_idx = Some(idx);
207        Ok(())
208    }
209
210    fn on_param_change(
211        &mut self,
212        reason: usize,
213        params: &PluginParamSnapshot,
214    ) -> ParamChangeResult {
215        // JPEG-specific: quality change
216        if Some(reason) == self.jpeg_quality_idx {
217            let q = params.value.as_i32().clamp(1, 100) as u8;
218            self.ctrl.writer.set_quality(q);
219            return ParamChangeResult::empty();
220        }
221        self.ctrl.on_param_change(reason, params)
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228    use ad_core_rs::ndarray::{NDDataBuffer, NDDimension};
229    use std::sync::atomic::{AtomicU32, Ordering};
230
231    static TEST_COUNTER: AtomicU32 = AtomicU32::new(0);
232
233    fn temp_path(prefix: &str) -> PathBuf {
234        let n = TEST_COUNTER.fetch_add(1, Ordering::Relaxed);
235        std::env::temp_dir().join(format!("adcore_test_{}_{}.jpg", prefix, n))
236    }
237
238    #[test]
239    fn test_write_u8() {
240        let path = temp_path("jpeg");
241        let mut writer = JpegWriter::new(90);
242
243        let mut arr = NDArray::new(
244            vec![NDDimension::new(8), NDDimension::new(8)],
245            NDDataType::UInt8,
246        );
247        if let NDDataBuffer::U8(ref mut v) = arr.data {
248            for i in 0..64 {
249                v[i] = (i * 4) as u8;
250            }
251        }
252
253        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
254        writer.write_file(&arr).unwrap();
255        writer.close_file().unwrap();
256
257        let data = std::fs::read(&path).unwrap();
258        // Check JPEG SOI marker
259        assert_eq!(&data[0..2], &[0xFF, 0xD8]);
260        // Check JPEG EOI marker at end
261        assert_eq!(&data[data.len() - 2..], &[0xFF, 0xD9]);
262
263        std::fs::remove_file(&path).ok();
264    }
265
266    #[test]
267    fn test_rejects_non_u8() {
268        let path = temp_path("jpeg_u16");
269        let mut writer = JpegWriter::new(90);
270
271        let arr = NDArray::new(
272            vec![NDDimension::new(4), NDDimension::new(4)],
273            NDDataType::UInt16,
274        );
275
276        let result = writer.open_file(&path, NDFileMode::Single, &arr);
277        assert!(result.is_err());
278    }
279
280    #[test]
281    fn test_quality_affects_size() {
282        let path_high = temp_path("jpeg_hi");
283        let path_low = temp_path("jpeg_lo");
284
285        let mut arr = NDArray::new(
286            vec![NDDimension::new(32), NDDimension::new(32)],
287            NDDataType::UInt8,
288        );
289        if let NDDataBuffer::U8(ref mut v) = arr.data {
290            for i in 0..v.len() {
291                v[i] = (i % 256) as u8;
292            }
293        }
294
295        let mut writer_high = JpegWriter::new(95);
296        writer_high
297            .open_file(&path_high, NDFileMode::Single, &arr)
298            .unwrap();
299        writer_high.write_file(&arr).unwrap();
300        writer_high.close_file().unwrap();
301
302        let mut writer_low = JpegWriter::new(10);
303        writer_low
304            .open_file(&path_low, NDFileMode::Single, &arr)
305            .unwrap();
306        writer_low.write_file(&arr).unwrap();
307        writer_low.close_file().unwrap();
308
309        let size_high = std::fs::metadata(&path_high).unwrap().len();
310        let size_low = std::fs::metadata(&path_low).unwrap().len();
311        assert!(
312            size_high > size_low,
313            "high quality ({}) should be larger than low quality ({})",
314            size_high,
315            size_low
316        );
317
318        std::fs::remove_file(&path_high).ok();
319        std::fs::remove_file(&path_low).ok();
320    }
321
322    #[test]
323    fn test_default_quality_is_50() {
324        // C++ NDFileJPEG.cpp:327 default quality is 50.
325        assert_eq!(JpegFileProcessor::default().ctrl.writer.quality, 50);
326        assert_eq!(JpegWriter::new(50).quality, 50);
327    }
328
329    #[test]
330    fn test_register_params_seeds_quality_pv() {
331        use asyn_rs::port::{PortDriverBase, PortFlags};
332        let mut base = PortDriverBase::new("jpeg_param_test", 1, PortFlags::default());
333        let mut proc = JpegFileProcessor::new(50);
334        proc.register_params(&mut base).unwrap();
335        let idx = proc.jpeg_quality_idx.unwrap();
336        // Readback PV must equal the encoder's effective quality, not 0.
337        assert_eq!(base.get_int32_param(idx, 0).unwrap(), 50);
338    }
339
340    #[test]
341    fn test_roundtrip_luma() {
342        let path = temp_path("jpeg_rt");
343        let mut writer = JpegWriter::new(100);
344
345        let mut arr = NDArray::new(
346            vec![NDDimension::new(8), NDDimension::new(8)],
347            NDDataType::UInt8,
348        );
349        if let NDDataBuffer::U8(ref mut v) = arr.data {
350            // Use uniform value so JPEG compression is lossless at quality 100
351            for i in 0..64 {
352                v[i] = 128;
353            }
354        }
355
356        writer.open_file(&path, NDFileMode::Single, &arr).unwrap();
357        writer.write_file(&arr).unwrap();
358
359        let read_back = writer.read_file().unwrap();
360        assert_eq!(read_back.data.data_type(), NDDataType::UInt8);
361        if let NDDataBuffer::U8(ref v) = read_back.data {
362            // With uniform input at max quality, decoded values should be close
363            for &px in v.iter() {
364                assert!(
365                    (px as i16 - 128).unsigned_abs() < 5,
366                    "pixel {} too far from 128",
367                    px
368                );
369            }
370        }
371
372        writer.close_file().unwrap();
373        std::fs::remove_file(&path).ok();
374    }
375}
ad_plugins_rs/file_jpeg.rs

ad_plugins_rs/
file_jpeg.rs
