ad_plugins/
color_convert.rs

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

ad_plugins/
color_convert.rs
