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
15pub 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 let color_mode = array.info().color_mode;
57
58 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
167pub 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 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 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 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 #[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 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 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 assert_eq!(&data[0..2], &[0xFF, 0xD8]);
333 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 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 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 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 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 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 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 let dims = jpeg_roundtrip_dims("jpeg_rgb3", NDColorMode::RGB3, vec![5, 4, 3]);
495 assert_eq!(dims, vec![3, 5, 4]);
496 }
497}