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

ad_plugins_rs/
color_convert.rs
