ad_plugins_rs/
color_convert.rs

1use std::sync::Arc;
2
3#[cfg(feature = "parallel")]
4use crate::par_util;
5#[cfg(feature = "parallel")]
6use rayon::prelude::*;
7
8use ad_core_rs::color::{self, NDBayerPattern, NDColorMode};
9use ad_core_rs::ndarray::{NDArray, NDDataBuffer, NDDataType, NDDimension};
10use ad_core_rs::ndarray_pool::NDArrayPool;
11use ad_core_rs::plugin::runtime::{NDPluginProcess, ProcessResult};
12
13/// Simple Bayer demosaic using bilinear interpolation.
14pub fn bayer_to_rgb1(src: &NDArray, pattern: NDBayerPattern) -> Option<NDArray> {
15    if src.dims.len() != 2 {
16        return None;
17    }
18    let w = src.dims[0].size;
19    let h = src.dims[1].size;
20
21    // Read dimension offsets to adjust the bayer phase when offset is odd
22    let offset_x = src.dims[0].offset;
23    let offset_y = src.dims[1].offset;
24
25    // Pre-compute source values into a flat f64 vec for efficient random access
26    let n = w * h;
27    let src_vals: Vec<f64> = (0..n)
28        .map(|i| src.data.get_as_f64(i).unwrap_or(0.0))
29        .collect();
30    let get_val = |x: usize, y: usize| -> f64 { src_vals[y * w + x] };
31
32    let mut r = vec![0.0f64; n];
33    let mut g = vec![0.0f64; n];
34    let mut b = vec![0.0f64; n];
35
36    // Determine which color each pixel position has, flipping phase for odd offsets
37    let (mut r_row_even, mut r_col_even) = match pattern {
38        NDBayerPattern::RGGB => (true, true),
39        NDBayerPattern::GBRG => (true, false),
40        NDBayerPattern::GRBG => (false, true),
41        NDBayerPattern::BGGR => (false, false),
42    };
43    if offset_x % 2 != 0 {
44        r_col_even = !r_col_even;
45    }
46    if offset_y % 2 != 0 {
47        r_row_even = !r_row_even;
48    }
49
50    // Helper to demosaic a single row into (r, g, b) slices
51    let demosaic_row = |y: usize, r_row: &mut [f64], g_row: &mut [f64], b_row: &mut [f64]| {
52        let even_row = (y % 2 == 0) == r_row_even;
53        for x in 0..w {
54            let val = get_val(x, y);
55            let even_col = (x % 2 == 0) == r_col_even;
56
57            match (even_row, even_col) {
58                (true, true) => {
59                    r_row[x] = val;
60                    let mut gsum = 0.0;
61                    let mut gc = 0;
62                    if x > 0 {
63                        gsum += get_val(x - 1, y);
64                        gc += 1;
65                    }
66                    if x < w - 1 {
67                        gsum += get_val(x + 1, y);
68                        gc += 1;
69                    }
70                    if y > 0 {
71                        gsum += get_val(x, y - 1);
72                        gc += 1;
73                    }
74                    if y < h - 1 {
75                        gsum += get_val(x, y + 1);
76                        gc += 1;
77                    }
78                    g_row[x] = if gc > 0 { gsum / gc as f64 } else { 0.0 };
79                    let mut bsum = 0.0;
80                    let mut bc = 0;
81                    if x > 0 && y > 0 {
82                        bsum += get_val(x - 1, y - 1);
83                        bc += 1;
84                    }
85                    if x < w - 1 && y > 0 {
86                        bsum += get_val(x + 1, y - 1);
87                        bc += 1;
88                    }
89                    if x > 0 && y < h - 1 {
90                        bsum += get_val(x - 1, y + 1);
91                        bc += 1;
92                    }
93                    if x < w - 1 && y < h - 1 {
94                        bsum += get_val(x + 1, y + 1);
95                        bc += 1;
96                    }
97                    b_row[x] = if bc > 0 { bsum / bc as f64 } else { 0.0 };
98                }
99                (true, false) | (false, true) => {
100                    g_row[x] = val;
101                    if even_row {
102                        let mut rsum = 0.0;
103                        let mut rc = 0;
104                        if x > 0 {
105                            rsum += get_val(x - 1, y);
106                            rc += 1;
107                        }
108                        if x < w - 1 {
109                            rsum += get_val(x + 1, y);
110                            rc += 1;
111                        }
112                        r_row[x] = if rc > 0 { rsum / rc as f64 } else { 0.0 };
113                        let mut bsum = 0.0;
114                        let mut bc = 0;
115                        if y > 0 {
116                            bsum += get_val(x, y - 1);
117                            bc += 1;
118                        }
119                        if y < h - 1 {
120                            bsum += get_val(x, y + 1);
121                            bc += 1;
122                        }
123                        b_row[x] = if bc > 0 { bsum / bc as f64 } else { 0.0 };
124                    } else {
125                        let mut bsum = 0.0;
126                        let mut bc = 0;
127                        if x > 0 {
128                            bsum += get_val(x - 1, y);
129                            bc += 1;
130                        }
131                        if x < w - 1 {
132                            bsum += get_val(x + 1, y);
133                            bc += 1;
134                        }
135                        b_row[x] = if bc > 0 { bsum / bc as f64 } else { 0.0 };
136                        let mut rsum = 0.0;
137                        let mut rc = 0;
138                        if y > 0 {
139                            rsum += get_val(x, y - 1);
140                            rc += 1;
141                        }
142                        if y < h - 1 {
143                            rsum += get_val(x, y + 1);
144                            rc += 1;
145                        }
146                        r_row[x] = if rc > 0 { rsum / rc as f64 } else { 0.0 };
147                    }
148                }
149                (false, false) => {
150                    b_row[x] = val;
151                    let mut gsum = 0.0;
152                    let mut gc = 0;
153                    if x > 0 {
154                        gsum += get_val(x - 1, y);
155                        gc += 1;
156                    }
157                    if x < w - 1 {
158                        gsum += get_val(x + 1, y);
159                        gc += 1;
160                    }
161                    if y > 0 {
162                        gsum += get_val(x, y - 1);
163                        gc += 1;
164                    }
165                    if y < h - 1 {
166                        gsum += get_val(x, y + 1);
167                        gc += 1;
168                    }
169                    g_row[x] = if gc > 0 { gsum / gc as f64 } else { 0.0 };
170                    let mut rsum = 0.0;
171                    let mut rc = 0;
172                    if x > 0 && y > 0 {
173                        rsum += get_val(x - 1, y - 1);
174                        rc += 1;
175                    }
176                    if x < w - 1 && y > 0 {
177                        rsum += get_val(x + 1, y - 1);
178                        rc += 1;
179                    }
180                    if x > 0 && y < h - 1 {
181                        rsum += get_val(x - 1, y + 1);
182                        rc += 1;
183                    }
184                    if x < w - 1 && y < h - 1 {
185                        rsum += get_val(x + 1, y + 1);
186                        rc += 1;
187                    }
188                    r_row[x] = if rc > 0 { rsum / rc as f64 } else { 0.0 };
189                }
190            }
191        }
192    };
193
194    #[cfg(feature = "parallel")]
195    let use_parallel = par_util::should_parallelize(n);
196    #[cfg(not(feature = "parallel"))]
197    let use_parallel = false;
198
199    if use_parallel {
200        #[cfg(feature = "parallel")]
201        {
202            // Split r, g, b into per-row mutable slices and process in parallel
203            let r_rows: Vec<&mut [f64]> = r.chunks_mut(w).collect();
204            let g_rows: Vec<&mut [f64]> = g.chunks_mut(w).collect();
205            let b_rows: Vec<&mut [f64]> = b.chunks_mut(w).collect();
206
207            par_util::thread_pool().install(|| {
208                r_rows
209                    .into_par_iter()
210                    .zip(g_rows.into_par_iter())
211                    .zip(b_rows.into_par_iter())
212                    .enumerate()
213                    .for_each(|(y, ((r_row, g_row), b_row))| {
214                        demosaic_row(y, r_row, g_row, b_row);
215                    });
216            });
217        }
218    } else {
219        for y in 0..h {
220            let row_start = y * w;
221            let row_end = row_start + w;
222            demosaic_row(
223                y,
224                &mut r[row_start..row_end],
225                &mut g[row_start..row_end],
226                &mut b[row_start..row_end],
227            );
228        }
229    }
230
231    // Build RGB1 interleaved output
232    let out_data = match src.data.data_type() {
233        NDDataType::UInt8 => {
234            let mut out = vec![0u8; n * 3];
235            for i in 0..n {
236                out[i * 3] = r[i].clamp(0.0, 255.0) as u8;
237                out[i * 3 + 1] = g[i].clamp(0.0, 255.0) as u8;
238                out[i * 3 + 2] = b[i].clamp(0.0, 255.0) as u8;
239            }
240            NDDataBuffer::U8(out)
241        }
242        NDDataType::UInt16 => {
243            let mut out = vec![0u16; n * 3];
244            for i in 0..n {
245                out[i * 3] = r[i].clamp(0.0, 65535.0) as u16;
246                out[i * 3 + 1] = g[i].clamp(0.0, 65535.0) as u16;
247                out[i * 3 + 2] = b[i].clamp(0.0, 65535.0) as u16;
248            }
249            NDDataBuffer::U16(out)
250        }
251        _ => return None,
252    };
253
254    let dims = vec![
255        NDDimension::new(3),
256        NDDimension::new(w),
257        NDDimension::new(h),
258    ];
259    let mut arr = NDArray::new(dims, src.data.data_type());
260    arr.data = out_data;
261    arr.unique_id = src.unique_id;
262    arr.timestamp = src.timestamp;
263    arr.attributes = src.attributes.clone();
264    Some(arr)
265}
266
267/// Generate a jet colormap lookup table (256 entries, RGB).
268///
269/// Maps scalar values 0..255 through blue -> cyan -> green -> yellow -> red.
270fn jet_colormap() -> [[u8; 3]; 256] {
271    let mut lut = [[0u8; 3]; 256];
272    for i in 0..256 {
273        let v = i as f64 / 255.0;
274        // Jet colormap: blue -> cyan -> green -> yellow -> red
275        let r = (1.5 - (4.0 * v - 3.0).abs()).clamp(0.0, 1.0);
276        let g = (1.5 - (4.0 * v - 2.0).abs()).clamp(0.0, 1.0);
277        let b = (1.5 - (4.0 * v - 1.0).abs()).clamp(0.0, 1.0);
278        lut[i] = [(r * 255.0) as u8, (g * 255.0) as u8, (b * 255.0) as u8];
279    }
280    lut
281}
282
283/// Convert a mono UInt8 image to RGB1 using false color (jet colormap).
284///
285/// Only supports 2D UInt8 arrays. Each pixel value is mapped through the jet
286/// colormap LUT to produce a pseudo-color RGB1 output.
287fn false_color_mono_to_rgb1(src: &NDArray) -> Option<NDArray> {
288    if src.dims.len() != 2 || src.data.data_type() != NDDataType::UInt8 {
289        return None;
290    }
291
292    let w = src.dims[0].size;
293    let h = src.dims[1].size;
294    let n = w * h;
295    let lut = jet_colormap();
296
297    let src_slice = src.data.as_u8_slice();
298    let mut out = vec![0u8; n * 3];
299    for i in 0..n {
300        let val = src_slice[i] as usize;
301        let [r, g, b] = lut[val];
302        out[i * 3] = r;
303        out[i * 3 + 1] = g;
304        out[i * 3 + 2] = b;
305    }
306
307    let dims = vec![
308        NDDimension::new(3),
309        NDDimension::new(w),
310        NDDimension::new(h),
311    ];
312    let mut arr = NDArray::new(dims, NDDataType::UInt8);
313    arr.data = NDDataBuffer::U8(out);
314    arr.unique_id = src.unique_id;
315    arr.timestamp = src.timestamp;
316    arr.attributes = src.attributes.clone();
317    Some(arr)
318}
319
320/// Detect the color mode of an NDArray.
321///
322/// Checks the `ColorMode` NDAttribute first (required for YUV422/YUV411 which are
323/// 2D arrays indistinguishable from Mono, and YUV444/Bayer which share dimensions
324/// with RGB1/Mono). Falls back to dimension-based detection.
325fn detect_color_mode(array: &NDArray) -> NDColorMode {
326    if let Some(attr) = array.attributes.get("ColorMode") {
327        if let Some(v) = attr.value.as_i64() {
328            return NDColorMode::from_i32(v as i32);
329        }
330    }
331    match array.dims.len() {
332        0 | 1 => NDColorMode::Mono,
333        2 => NDColorMode::Mono,
334        3 => {
335            if array.dims[0].size == 3 {
336                NDColorMode::RGB1
337            } else if array.dims[1].size == 3 {
338                NDColorMode::RGB2
339            } else if array.dims[2].size == 3 {
340                NDColorMode::RGB3
341            } else {
342                NDColorMode::Mono
343            }
344        }
345        _ => NDColorMode::Mono,
346    }
347}
348
349/// Color convert plugin configuration.
350#[derive(Debug, Clone)]
351pub struct ColorConvertConfig {
352    pub target_mode: NDColorMode,
353    pub bayer_pattern: NDBayerPattern,
354    /// False color mode: 0=off, 1=Rainbow, 2=Iron. Nonzero is treated as enabled.
355    pub false_color: i32,
356}
357
358/// Pure color conversion processing logic.
359pub struct ColorConvertProcessor {
360    config: ColorConvertConfig,
361    color_mode_out_idx: Option<usize>,
362    false_color_idx: Option<usize>,
363}
364
365impl ColorConvertProcessor {
366    pub fn new(config: ColorConvertConfig) -> Self {
367        Self {
368            config,
369            color_mode_out_idx: None,
370            false_color_idx: None,
371        }
372    }
373}
374
375impl NDPluginProcess for ColorConvertProcessor {
376    fn register_params(
377        &mut self,
378        base: &mut asyn_rs::port::PortDriverBase,
379    ) -> asyn_rs::error::AsynResult<()> {
380        use asyn_rs::param::ParamType;
381        base.create_param("COLOR_MODE_OUT", ParamType::Int32)?;
382        base.create_param("FALSE_COLOR", ParamType::Int32)?;
383        self.color_mode_out_idx = base.find_param("COLOR_MODE_OUT");
384        self.false_color_idx = base.find_param("FALSE_COLOR");
385        Ok(())
386    }
387
388    fn on_param_change(
389        &mut self,
390        reason: usize,
391        params: &ad_core_rs::plugin::runtime::PluginParamSnapshot,
392    ) -> ad_core_rs::plugin::runtime::ParamChangeResult {
393        if Some(reason) == self.color_mode_out_idx {
394            self.config.target_mode = NDColorMode::from_i32(params.value.as_i32());
395        } else if Some(reason) == self.false_color_idx {
396            self.config.false_color = params.value.as_i32();
397        }
398        ad_core_rs::plugin::runtime::ParamChangeResult::updates(vec![])
399    }
400
401    fn process_array(&mut self, array: &NDArray, _pool: &NDArrayPool) -> ProcessResult {
402        let src_mode = detect_color_mode(array);
403        let target = self.config.target_mode;
404
405        // Same mode - passthrough
406        if src_mode == target {
407            return ProcessResult::arrays(vec![Arc::new(array.clone())]);
408        }
409
410        // Step 1: Convert source to RGB1 intermediate
411        let rgb1 = match src_mode {
412            NDColorMode::RGB1 => Some(array.clone()),
413            NDColorMode::Mono => {
414                if self.config.false_color != 0 {
415                    false_color_mono_to_rgb1(array).or_else(|| color::mono_to_rgb1(array).ok())
416                } else {
417                    color::mono_to_rgb1(array).ok()
418                }
419            }
420            NDColorMode::Bayer => bayer_to_rgb1(array, self.config.bayer_pattern),
421            NDColorMode::RGB2 | NDColorMode::RGB3 => {
422                color::convert_rgb_layout(array, src_mode, NDColorMode::RGB1).ok()
423            }
424            NDColorMode::YUV444 => color::yuv444_to_rgb1(array).ok(),
425            NDColorMode::YUV422 => color::yuv422_to_rgb1(array).ok(),
426            NDColorMode::YUV411 => color::yuv411_to_rgb1(array).ok(),
427        };
428
429        let rgb1 = match rgb1 {
430            Some(r) => r,
431            None => return ProcessResult::empty(),
432        };
433
434        // Step 2: Convert RGB1 intermediate to target
435        let result = match target {
436            NDColorMode::RGB1 => Some(rgb1),
437            NDColorMode::Mono => color::rgb1_to_mono(&rgb1).ok(),
438            NDColorMode::Bayer => None,
439            NDColorMode::RGB2 | NDColorMode::RGB3 => {
440                color::convert_rgb_layout(&rgb1, NDColorMode::RGB1, target).ok()
441            }
442            NDColorMode::YUV444 => color::rgb1_to_yuv444(&rgb1).ok(),
443            NDColorMode::YUV422 => color::rgb1_to_yuv422(&rgb1).ok(),
444            NDColorMode::YUV411 => color::rgb1_to_yuv411(&rgb1).ok(),
445        };
446
447        match result {
448            Some(mut out) => {
449                // C++: set ColorMode attribute on output array
450                let color_mode_val = match target {
451                    NDColorMode::Mono => 0i32,
452                    NDColorMode::Bayer => 1,
453                    NDColorMode::RGB1 => 2,
454                    NDColorMode::RGB2 => 3,
455                    NDColorMode::RGB3 => 4,
456                    NDColorMode::YUV444 => 5,
457                    NDColorMode::YUV422 => 6,
458                    NDColorMode::YUV411 => 7,
459                };
460                use ad_core_rs::attributes::{NDAttrSource, NDAttrValue, NDAttribute};
461                out.attributes.add(NDAttribute {
462                    name: "ColorMode".into(),
463                    description: "Color Mode".into(),
464                    source: NDAttrSource::Driver,
465                    value: NDAttrValue::Int32(color_mode_val),
466                });
467                ProcessResult::arrays(vec![Arc::new(out)])
468            }
469            None => ProcessResult::empty(),
470        }
471    }
472
473    fn plugin_type(&self) -> &str {
474        "NDPluginColorConvert"
475    }
476}
477
478#[cfg(test)]
479mod tests {
480    use super::*;
481
482    #[test]
483    fn test_bayer_to_rgb1_basic() {
484        // 4x4 RGGB bayer pattern
485        let mut arr = NDArray::new(
486            vec![NDDimension::new(4), NDDimension::new(4)],
487            NDDataType::UInt8,
488        );
489        if let NDDataBuffer::U8(ref mut v) = arr.data {
490            // Simple pattern: all pixels = 128
491            for i in 0..16 {
492                v[i] = 128;
493            }
494        }
495
496        let rgb = bayer_to_rgb1(&arr, NDBayerPattern::RGGB).unwrap();
497        assert_eq!(rgb.dims.len(), 3);
498        assert_eq!(rgb.dims[0].size, 3); // color
499        assert_eq!(rgb.dims[1].size, 4); // x
500        assert_eq!(rgb.dims[2].size, 4); // y
501    }
502
503    #[test]
504    fn test_color_convert_processor_bayer() {
505        let config = ColorConvertConfig {
506            target_mode: NDColorMode::RGB1,
507            bayer_pattern: NDBayerPattern::RGGB,
508            false_color: 0,
509        };
510        let mut proc = ColorConvertProcessor::new(config);
511        let pool = NDArrayPool::new(1_000_000);
512
513        let mut arr = NDArray::new(
514            vec![NDDimension::new(4), NDDimension::new(4)],
515            NDDataType::UInt8,
516        );
517        if let NDDataBuffer::U8(ref mut v) = arr.data {
518            for i in 0..16 {
519                v[i] = 128;
520            }
521        }
522
523        let result = proc.process_array(&arr, &pool);
524        assert_eq!(result.output_arrays.len(), 1);
525        assert_eq!(result.output_arrays[0].dims[0].size, 3); // RGB color dim
526    }
527
528    #[test]
529    fn test_false_color_conversion() {
530        let config = ColorConvertConfig {
531            target_mode: NDColorMode::RGB1,
532            bayer_pattern: NDBayerPattern::RGGB,
533            false_color: 1,
534        };
535        let mut proc = ColorConvertProcessor::new(config);
536        let pool = NDArrayPool::new(1_000_000);
537
538        // Create a 4x4 mono UInt8 image with a gradient
539        let mut arr = NDArray::new(
540            vec![NDDimension::new(4), NDDimension::new(4)],
541            NDDataType::UInt8,
542        );
543        if let NDDataBuffer::U8(ref mut v) = arr.data {
544            for i in 0..16 {
545                v[i] = (i * 17) as u8; // 0, 17, 34, ... 255
546            }
547        }
548
549        let result = proc.process_array(&arr, &pool);
550        assert_eq!(result.output_arrays.len(), 1);
551        let out = &result.output_arrays[0];
552        assert_eq!(out.dims.len(), 3);
553        assert_eq!(out.dims[0].size, 3); // color
554        assert_eq!(out.dims[1].size, 4); // x
555        assert_eq!(out.dims[2].size, 4); // y
556
557        // Verify false color: pixel 0 (value=0) should be blue, pixel 15 (value=255) should be red
558        let lut = jet_colormap();
559        if let NDDataBuffer::U8(ref v) = out.data {
560            // First pixel (value=0)
561            assert_eq!(v[0], lut[0][0]); // R
562            assert_eq!(v[1], lut[0][1]); // G
563            assert_eq!(v[2], lut[0][2]); // B
564            // Last pixel (value=255)
565            let last = 15 * 3;
566            assert_eq!(v[last], lut[255][0]); // R
567            assert_eq!(v[last + 1], lut[255][1]); // G
568            assert_eq!(v[last + 2], lut[255][2]); // B
569        } else {
570            panic!("expected UInt8 output");
571        }
572    }
573
574    #[test]
575    fn test_rgb1_to_rgb2_conversion() {
576        let config = ColorConvertConfig {
577            target_mode: NDColorMode::RGB2,
578            bayer_pattern: NDBayerPattern::RGGB,
579            false_color: 0,
580        };
581        let mut proc = ColorConvertProcessor::new(config);
582        let pool = NDArrayPool::new(1_000_000);
583
584        // Create RGB1 image: dims [3, 4, 4]
585        let mut arr = NDArray::new(
586            vec![
587                NDDimension::new(3),
588                NDDimension::new(4),
589                NDDimension::new(4),
590            ],
591            NDDataType::UInt8,
592        );
593        if let NDDataBuffer::U8(ref mut v) = arr.data {
594            for i in 0..v.len() {
595                v[i] = (i % 256) as u8;
596            }
597        }
598
599        let result = proc.process_array(&arr, &pool);
600        assert_eq!(result.output_arrays.len(), 1);
601        let out = &result.output_arrays[0];
602        assert_eq!(out.dims.len(), 3);
603        // RGB2 has color dim in position 1
604        assert_eq!(out.dims[1].size, 3);
605    }
606
607    #[test]
608    fn test_rgb2_to_mono_conversion() {
609        let config = ColorConvertConfig {
610            target_mode: NDColorMode::Mono,
611            bayer_pattern: NDBayerPattern::RGGB,
612            false_color: 0,
613        };
614        let mut proc = ColorConvertProcessor::new(config);
615        let pool = NDArrayPool::new(1_000_000);
616
617        // Create RGB2 image: dims [4, 3, 4]
618        let mut arr = NDArray::new(
619            vec![
620                NDDimension::new(4),
621                NDDimension::new(3),
622                NDDimension::new(4),
623            ],
624            NDDataType::UInt8,
625        );
626        if let NDDataBuffer::U8(ref mut v) = arr.data {
627            for i in 0..v.len() {
628                v[i] = 128;
629            }
630        }
631
632        let result = proc.process_array(&arr, &pool);
633        assert_eq!(result.output_arrays.len(), 1);
634        let out = &result.output_arrays[0];
635        // Mono output should be 2D
636        assert_eq!(out.dims.len(), 2);
637    }
638
639    #[test]
640    fn test_detect_color_mode() {
641        // 2D -> Mono
642        let arr2d = NDArray::new(
643            vec![NDDimension::new(4), NDDimension::new(4)],
644            NDDataType::UInt8,
645        );
646        assert_eq!(detect_color_mode(&arr2d), NDColorMode::Mono);
647
648        // 3D with color dim first -> RGB1
649        let arr_rgb1 = NDArray::new(
650            vec![
651                NDDimension::new(3),
652                NDDimension::new(4),
653                NDDimension::new(4),
654            ],
655            NDDataType::UInt8,
656        );
657        assert_eq!(detect_color_mode(&arr_rgb1), NDColorMode::RGB1);
658
659        // 3D with color dim second -> RGB2
660        let arr_rgb2 = NDArray::new(
661            vec![
662                NDDimension::new(4),
663                NDDimension::new(3),
664                NDDimension::new(4),
665            ],
666            NDDataType::UInt8,
667        );
668        assert_eq!(detect_color_mode(&arr_rgb2), NDColorMode::RGB2);
669
670        // 3D with color dim last -> RGB3
671        let arr_rgb3 = NDArray::new(
672            vec![
673                NDDimension::new(4),
674                NDDimension::new(4),
675                NDDimension::new(3),
676            ],
677            NDDataType::UInt8,
678        );
679        assert_eq!(detect_color_mode(&arr_rgb3), NDColorMode::RGB3);
680    }
681
682    #[test]
683    fn test_same_mode_passthrough() {
684        let config = ColorConvertConfig {
685            target_mode: NDColorMode::Mono,
686            bayer_pattern: NDBayerPattern::RGGB,
687            false_color: 0,
688        };
689        let mut proc = ColorConvertProcessor::new(config);
690        let pool = NDArrayPool::new(1_000_000);
691
692        // 2D mono input with Mono target -> passthrough
693        let mut arr = NDArray::new(
694            vec![NDDimension::new(4), NDDimension::new(4)],
695            NDDataType::UInt8,
696        );
697        arr.unique_id = 42;
698        if let NDDataBuffer::U8(ref mut v) = arr.data {
699            for i in 0..16 {
700                v[i] = i as u8;
701            }
702        }
703
704        let result = proc.process_array(&arr, &pool);
705        assert_eq!(result.output_arrays.len(), 1);
706        assert_eq!(result.output_arrays[0].unique_id, 42);
707        assert_eq!(result.output_arrays[0].dims.len(), 2);
708    }
709
710    fn set_color_mode_attr(arr: &mut NDArray, mode: NDColorMode) {
711        use ad_core_rs::attributes::{NDAttrSource, NDAttrValue, NDAttribute};
712        arr.attributes.add(NDAttribute {
713            name: "ColorMode".to_string(),
714            description: String::new(),
715            source: NDAttrSource::Driver,
716            value: NDAttrValue::Int32(mode as i32),
717        });
718    }
719
720    #[test]
721    fn test_bayer_to_mono_via_rgb1() {
722        let config = ColorConvertConfig {
723            target_mode: NDColorMode::Mono,
724            bayer_pattern: NDBayerPattern::RGGB,
725            false_color: 0,
726        };
727        let mut proc = ColorConvertProcessor::new(config);
728        let pool = NDArrayPool::new(1_000_000);
729
730        let mut arr = NDArray::new(
731            vec![NDDimension::new(4), NDDimension::new(4)],
732            NDDataType::UInt8,
733        );
734        set_color_mode_attr(&mut arr, NDColorMode::Bayer);
735        if let NDDataBuffer::U8(ref mut v) = arr.data {
736            for i in 0..16 {
737                v[i] = 128;
738            }
739        }
740
741        let result = proc.process_array(&arr, &pool);
742        assert_eq!(result.output_arrays.len(), 1);
743        assert_eq!(result.output_arrays[0].dims.len(), 2);
744    }
745
746    #[test]
747    fn test_rgb1_to_yuv444_conversion() {
748        let config = ColorConvertConfig {
749            target_mode: NDColorMode::YUV444,
750            bayer_pattern: NDBayerPattern::RGGB,
751            false_color: 0,
752        };
753        let mut proc = ColorConvertProcessor::new(config);
754        let pool = NDArrayPool::new(1_000_000);
755
756        let mut arr = NDArray::new(
757            vec![
758                NDDimension::new(3),
759                NDDimension::new(4),
760                NDDimension::new(4),
761            ],
762            NDDataType::UInt8,
763        );
764        if let NDDataBuffer::U8(ref mut v) = arr.data {
765            for i in 0..v.len() {
766                v[i] = (i % 256) as u8;
767            }
768        }
769
770        let result = proc.process_array(&arr, &pool);
771        assert_eq!(result.output_arrays.len(), 1);
772        let out = &result.output_arrays[0];
773        assert_eq!(out.dims.len(), 3);
774        assert_eq!(out.dims[0].size, 3);
775    }
776
777    #[test]
778    fn test_yuv422_to_rgb1_conversion() {
779        let config = ColorConvertConfig {
780            target_mode: NDColorMode::RGB1,
781            bayer_pattern: NDBayerPattern::RGGB,
782            false_color: 0,
783        };
784        let mut proc = ColorConvertProcessor::new(config);
785        let pool = NDArrayPool::new(1_000_000);
786
787        // packed_x=8 means 4 pixels wide, 2 rows
788        let mut arr = NDArray::new(
789            vec![NDDimension::new(8), NDDimension::new(2)],
790            NDDataType::UInt8,
791        );
792        set_color_mode_attr(&mut arr, NDColorMode::YUV422);
793        if let NDDataBuffer::U8(ref mut v) = arr.data {
794            // UYVY pattern: U Y0 V Y1
795            let uyvy: [u8; 16] = [
796                128, 100, 128, 150, 128, 200, 128, 50, 128, 128, 128, 128, 128, 64, 128, 192,
797            ];
798            v[..16].copy_from_slice(&uyvy);
799        }
800
801        let result = proc.process_array(&arr, &pool);
802        assert_eq!(result.output_arrays.len(), 1);
803        let out = &result.output_arrays[0];
804        assert_eq!(out.dims[0].size, 3);
805        assert_eq!(out.dims[1].size, 4);
806        assert_eq!(out.dims[2].size, 2);
807    }
808
809    #[test]
810    fn test_mono_to_yuv422_conversion() {
811        let config = ColorConvertConfig {
812            target_mode: NDColorMode::YUV422,
813            bayer_pattern: NDBayerPattern::RGGB,
814            false_color: 0,
815        };
816        let mut proc = ColorConvertProcessor::new(config);
817        let pool = NDArrayPool::new(1_000_000);
818
819        let mut arr = NDArray::new(
820            vec![NDDimension::new(4), NDDimension::new(2)],
821            NDDataType::UInt8,
822        );
823        if let NDDataBuffer::U8(ref mut v) = arr.data {
824            for i in 0..8 {
825                v[i] = (i * 30) as u8;
826            }
827        }
828
829        let result = proc.process_array(&arr, &pool);
830        assert_eq!(result.output_arrays.len(), 1);
831        let out = &result.output_arrays[0];
832        assert_eq!(out.dims.len(), 2);
833        assert_eq!(out.dims[0].size, 8); // packed_x = 4*2
834    }
835
836    #[test]
837    fn test_yuv444_to_mono_conversion() {
838        let config = ColorConvertConfig {
839            target_mode: NDColorMode::Mono,
840            bayer_pattern: NDBayerPattern::RGGB,
841            false_color: 0,
842        };
843        let mut proc = ColorConvertProcessor::new(config);
844        let pool = NDArrayPool::new(1_000_000);
845
846        let mut arr = NDArray::new(
847            vec![
848                NDDimension::new(3),
849                NDDimension::new(4),
850                NDDimension::new(4),
851            ],
852            NDDataType::UInt8,
853        );
854        set_color_mode_attr(&mut arr, NDColorMode::YUV444);
855        if let NDDataBuffer::U8(ref mut v) = arr.data {
856            for i in 0..v.len() {
857                v[i] = 128;
858            }
859        }
860
861        let result = proc.process_array(&arr, &pool);
862        assert_eq!(result.output_arrays.len(), 1);
863        let out = &result.output_arrays[0];
864        assert_eq!(out.dims.len(), 2);
865        assert_eq!(out.dims[0].size, 4);
866        assert_eq!(out.dims[1].size, 4);
867    }
868
869    #[test]
870    fn test_detect_color_mode_with_attribute() {
871        let mut arr = NDArray::new(
872            vec![NDDimension::new(8), NDDimension::new(2)],
873            NDDataType::UInt8,
874        );
875        assert_eq!(detect_color_mode(&arr), NDColorMode::Mono);
876
877        set_color_mode_attr(&mut arr, NDColorMode::YUV422);
878        assert_eq!(detect_color_mode(&arr), NDColorMode::YUV422);
879    }
880
881    #[test]
882    fn test_jet_colormap_endpoints() {
883        let lut = jet_colormap();
884        // At v=0: r=clamp(1.5-3.0,0,1)=0, g=clamp(1.5-2.0,0,1)=0, b=clamp(1.5-1.0,0,1)=0.5
885        assert_eq!(lut[0][0], 0); // R
886        assert_eq!(lut[0][1], 0); // G
887        assert_eq!(lut[0][2], 127); // B (0.5 * 255 = 127)
888
889        // At v=1: r=clamp(1.5-1.0,0,1)=0.5, g=clamp(1.5-2.0,0,1)=0, b=clamp(1.5-3.0,0,1)=0
890        assert_eq!(lut[255][0], 127); // R
891        assert_eq!(lut[255][1], 0); // G
892        assert_eq!(lut[255][2], 0); // B
893    }
894}
ad_plugins_rs/color_convert.rs

ad_plugins_rs/
color_convert.rs
