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 from its dimensionality.
220///
221/// 2D arrays are always treated as Mono (Bayer cannot be distinguished from Mono
222/// by dimensions alone). For 3D arrays, the dimension with size 3 determines the
223/// RGB layout variant.
224fn detect_color_mode(array: &NDArray) -> NDColorMode {
225    match array.dims.len() {
226        0 | 1 => NDColorMode::Mono,
227        2 => NDColorMode::Mono, // 2D is always mono (or bayer, but we can't tell)
228        3 => {
229            // Check color dimension
230            if array.dims[0].size == 3 {
231                NDColorMode::RGB1
232            } else if array.dims[1].size == 3 {
233                NDColorMode::RGB2
234            } else if array.dims[2].size == 3 {
235                NDColorMode::RGB3
236            } else {
237                NDColorMode::Mono
238            }
239        }
240        _ => NDColorMode::Mono,
241    }
242}
243
244/// Color convert plugin configuration.
245#[derive(Debug, Clone)]
246pub struct ColorConvertConfig {
247    pub target_mode: NDColorMode,
248    pub bayer_pattern: NDBayerPattern,
249    pub false_color: bool,
250}
251
252/// Pure color conversion processing logic.
253pub struct ColorConvertProcessor {
254    config: ColorConvertConfig,
255}
256
257impl ColorConvertProcessor {
258    pub fn new(config: ColorConvertConfig) -> Self {
259        Self { config }
260    }
261}
262
263impl NDPluginProcess for ColorConvertProcessor {
264    fn process_array(&mut self, array: &NDArray, _pool: &NDArrayPool) -> ProcessResult {
265        let src_mode = detect_color_mode(array);
266        let result = match (src_mode, self.config.target_mode) {
267            // Same mode - passthrough
268            (s, t) if s == t => Some(array.clone()),
269
270            // To Mono
271            (NDColorMode::RGB1 | NDColorMode::RGB2 | NDColorMode::RGB3, NDColorMode::Mono) => {
272                // If not RGB1, convert to RGB1 first, then to mono
273                let rgb1 = if src_mode != NDColorMode::RGB1 {
274                    color::convert_rgb_layout(array, src_mode, NDColorMode::RGB1).ok()
275                } else {
276                    Some(array.clone())
277                };
278                rgb1.and_then(|a| color::rgb1_to_mono(&a).ok())
279            }
280
281            // Mono to RGB with false color
282            (NDColorMode::Mono, NDColorMode::RGB1) if self.config.false_color => {
283                false_color_mono_to_rgb1(array)
284            }
285
286            // Mono to any RGB
287            (NDColorMode::Mono, NDColorMode::RGB1) => color::mono_to_rgb1(array).ok(),
288            (NDColorMode::Mono, NDColorMode::RGB2 | NDColorMode::RGB3) => {
289                color::mono_to_rgb1(array).ok().and_then(|a| {
290                    color::convert_rgb_layout(&a, NDColorMode::RGB1, self.config.target_mode).ok()
291                })
292            }
293
294            // Bayer to any RGB
295            (NDColorMode::Bayer, NDColorMode::RGB1) => {
296                bayer_to_rgb1(array, self.config.bayer_pattern)
297            }
298            (NDColorMode::Bayer, NDColorMode::RGB2 | NDColorMode::RGB3) => {
299                bayer_to_rgb1(array, self.config.bayer_pattern).and_then(|a| {
300                    color::convert_rgb_layout(&a, NDColorMode::RGB1, self.config.target_mode).ok()
301                })
302            }
303
304            // RGB to RGB (layout conversion)
305            (
306                NDColorMode::RGB1 | NDColorMode::RGB2 | NDColorMode::RGB3,
307                NDColorMode::RGB1 | NDColorMode::RGB2 | NDColorMode::RGB3,
308            ) => color::convert_rgb_layout(array, src_mode, self.config.target_mode).ok(),
309
310            _ => None,
311        };
312        match result {
313            Some(out) => ProcessResult::arrays(vec![Arc::new(out)]),
314            None => ProcessResult::empty(),
315        }
316    }
317
318    fn plugin_type(&self) -> &str {
319        "NDPluginColorConvert"
320    }
321}
322
323#[cfg(test)]
324mod tests {
325    use super::*;
326
327    #[test]
328    fn test_bayer_to_rgb1_basic() {
329        // 4x4 RGGB bayer pattern
330        let mut arr = NDArray::new(
331            vec![NDDimension::new(4), NDDimension::new(4)],
332            NDDataType::UInt8,
333        );
334        if let NDDataBuffer::U8(ref mut v) = arr.data {
335            // Simple pattern: all pixels = 128
336            for i in 0..16 {
337                v[i] = 128;
338            }
339        }
340
341        let rgb = bayer_to_rgb1(&arr, NDBayerPattern::RGGB).unwrap();
342        assert_eq!(rgb.dims.len(), 3);
343        assert_eq!(rgb.dims[0].size, 3); // color
344        assert_eq!(rgb.dims[1].size, 4); // x
345        assert_eq!(rgb.dims[2].size, 4); // y
346    }
347
348    #[test]
349    fn test_color_convert_processor_bayer() {
350        let config = ColorConvertConfig {
351            target_mode: NDColorMode::RGB1,
352            bayer_pattern: NDBayerPattern::RGGB,
353            false_color: false,
354        };
355        let mut proc = ColorConvertProcessor::new(config);
356        let pool = NDArrayPool::new(1_000_000);
357
358        let mut arr = NDArray::new(
359            vec![NDDimension::new(4), NDDimension::new(4)],
360            NDDataType::UInt8,
361        );
362        if let NDDataBuffer::U8(ref mut v) = arr.data {
363            for i in 0..16 {
364                v[i] = 128;
365            }
366        }
367
368        let result = proc.process_array(&arr, &pool);
369        assert_eq!(result.output_arrays.len(), 1);
370        assert_eq!(result.output_arrays[0].dims[0].size, 3); // RGB color dim
371    }
372
373    #[test]
374    fn test_false_color_conversion() {
375        let config = ColorConvertConfig {
376            target_mode: NDColorMode::RGB1,
377            bayer_pattern: NDBayerPattern::RGGB,
378            false_color: true,
379        };
380        let mut proc = ColorConvertProcessor::new(config);
381        let pool = NDArrayPool::new(1_000_000);
382
383        // Create a 4x4 mono UInt8 image with a gradient
384        let mut arr = NDArray::new(
385            vec![NDDimension::new(4), NDDimension::new(4)],
386            NDDataType::UInt8,
387        );
388        if let NDDataBuffer::U8(ref mut v) = arr.data {
389            for i in 0..16 {
390                v[i] = (i * 17) as u8; // 0, 17, 34, ... 255
391            }
392        }
393
394        let result = proc.process_array(&arr, &pool);
395        assert_eq!(result.output_arrays.len(), 1);
396        let out = &result.output_arrays[0];
397        assert_eq!(out.dims.len(), 3);
398        assert_eq!(out.dims[0].size, 3); // color
399        assert_eq!(out.dims[1].size, 4); // x
400        assert_eq!(out.dims[2].size, 4); // y
401
402        // Verify false color: pixel 0 (value=0) should be blue, pixel 15 (value=255) should be red
403        let lut = jet_colormap();
404        if let NDDataBuffer::U8(ref v) = out.data {
405            // First pixel (value=0)
406            assert_eq!(v[0], lut[0][0]); // R
407            assert_eq!(v[1], lut[0][1]); // G
408            assert_eq!(v[2], lut[0][2]); // B
409            // Last pixel (value=255)
410            let last = 15 * 3;
411            assert_eq!(v[last], lut[255][0]); // R
412            assert_eq!(v[last + 1], lut[255][1]); // G
413            assert_eq!(v[last + 2], lut[255][2]); // B
414        } else {
415            panic!("expected UInt8 output");
416        }
417    }
418
419    #[test]
420    fn test_rgb1_to_rgb2_conversion() {
421        let config = ColorConvertConfig {
422            target_mode: NDColorMode::RGB2,
423            bayer_pattern: NDBayerPattern::RGGB,
424            false_color: false,
425        };
426        let mut proc = ColorConvertProcessor::new(config);
427        let pool = NDArrayPool::new(1_000_000);
428
429        // Create RGB1 image: dims [3, 4, 4]
430        let mut arr = NDArray::new(
431            vec![NDDimension::new(3), NDDimension::new(4), NDDimension::new(4)],
432            NDDataType::UInt8,
433        );
434        if let NDDataBuffer::U8(ref mut v) = arr.data {
435            for i in 0..v.len() {
436                v[i] = (i % 256) as u8;
437            }
438        }
439
440        let result = proc.process_array(&arr, &pool);
441        assert_eq!(result.output_arrays.len(), 1);
442        let out = &result.output_arrays[0];
443        assert_eq!(out.dims.len(), 3);
444        // RGB2 has color dim in position 1
445        assert_eq!(out.dims[1].size, 3);
446    }
447
448    #[test]
449    fn test_rgb2_to_mono_conversion() {
450        let config = ColorConvertConfig {
451            target_mode: NDColorMode::Mono,
452            bayer_pattern: NDBayerPattern::RGGB,
453            false_color: false,
454        };
455        let mut proc = ColorConvertProcessor::new(config);
456        let pool = NDArrayPool::new(1_000_000);
457
458        // Create RGB2 image: dims [4, 3, 4]
459        let mut arr = NDArray::new(
460            vec![NDDimension::new(4), NDDimension::new(3), NDDimension::new(4)],
461            NDDataType::UInt8,
462        );
463        if let NDDataBuffer::U8(ref mut v) = arr.data {
464            for i in 0..v.len() {
465                v[i] = 128;
466            }
467        }
468
469        let result = proc.process_array(&arr, &pool);
470        assert_eq!(result.output_arrays.len(), 1);
471        let out = &result.output_arrays[0];
472        // Mono output should be 2D
473        assert_eq!(out.dims.len(), 2);
474    }
475
476    #[test]
477    fn test_detect_color_mode() {
478        // 2D -> Mono
479        let arr2d = NDArray::new(
480            vec![NDDimension::new(4), NDDimension::new(4)],
481            NDDataType::UInt8,
482        );
483        assert_eq!(detect_color_mode(&arr2d), NDColorMode::Mono);
484
485        // 3D with color dim first -> RGB1
486        let arr_rgb1 = NDArray::new(
487            vec![NDDimension::new(3), NDDimension::new(4), NDDimension::new(4)],
488            NDDataType::UInt8,
489        );
490        assert_eq!(detect_color_mode(&arr_rgb1), NDColorMode::RGB1);
491
492        // 3D with color dim second -> RGB2
493        let arr_rgb2 = NDArray::new(
494            vec![NDDimension::new(4), NDDimension::new(3), NDDimension::new(4)],
495            NDDataType::UInt8,
496        );
497        assert_eq!(detect_color_mode(&arr_rgb2), NDColorMode::RGB2);
498
499        // 3D with color dim last -> RGB3
500        let arr_rgb3 = NDArray::new(
501            vec![NDDimension::new(4), NDDimension::new(4), NDDimension::new(3)],
502            NDDataType::UInt8,
503        );
504        assert_eq!(detect_color_mode(&arr_rgb3), NDColorMode::RGB3);
505    }
506
507    #[test]
508    fn test_same_mode_passthrough() {
509        let config = ColorConvertConfig {
510            target_mode: NDColorMode::Mono,
511            bayer_pattern: NDBayerPattern::RGGB,
512            false_color: false,
513        };
514        let mut proc = ColorConvertProcessor::new(config);
515        let pool = NDArrayPool::new(1_000_000);
516
517        // 2D mono input with Mono target -> passthrough
518        let mut arr = NDArray::new(
519            vec![NDDimension::new(4), NDDimension::new(4)],
520            NDDataType::UInt8,
521        );
522        arr.unique_id = 42;
523        if let NDDataBuffer::U8(ref mut v) = arr.data {
524            for i in 0..16 {
525                v[i] = i as u8;
526            }
527        }
528
529        let result = proc.process_array(&arr, &pool);
530        assert_eq!(result.output_arrays.len(), 1);
531        assert_eq!(result.output_arrays[0].unique_id, 42);
532        assert_eq!(result.output_arrays[0].dims.len(), 2);
533    }
534
535    #[test]
536    fn test_jet_colormap_endpoints() {
537        let lut = jet_colormap();
538        // 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
539        assert_eq!(lut[0][0], 0); // R
540        assert_eq!(lut[0][1], 0); // G
541        assert_eq!(lut[0][2], 127); // B (0.5 * 255 = 127)
542
543        // 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
544        assert_eq!(lut[255][0], 127); // R
545        assert_eq!(lut[255][1], 0); // G
546        assert_eq!(lut[255][2], 0); // B
547    }
548}
ad_plugins/color_convert.rs

ad_plugins/
color_convert.rs
