Skip to main content

yscv_imgproc/ops/
bitwise.rs

1//! Bitwise and pixel-level operations (OpenCV-style in_range, bitwise, blending).
2
3use yscv_tensor::Tensor;
4
5use super::super::ImgProcError;
6
7/// Create a binary mask where each pixel is 1.0 if all channels are within
8/// `[lower, upper]` (inclusive), 0.0 otherwise.
9///
10/// Input: `[H, W, C]`, output: `[H, W, 1]`.
11pub fn in_range(image: &Tensor, lower: &[f32], upper: &[f32]) -> Result<Tensor, ImgProcError> {
12    let shape = image.shape();
13    if shape.len() != 3 {
14        return Err(ImgProcError::InvalidImageShape {
15            expected_rank: 3,
16            got: shape.to_vec(),
17        });
18    }
19    let (h, w, c) = (shape[0], shape[1], shape[2]);
20    if lower.len() != c || upper.len() != c {
21        return Err(ImgProcError::InvalidChannelCount {
22            expected: c,
23            got: lower.len(),
24        });
25    }
26
27    let data = image.data();
28    let mut out = vec![0.0f32; h * w];
29    for y in 0..h {
30        for x in 0..w {
31            let base = (y * w + x) * c;
32            let mut inside = true;
33            for ch in 0..c {
34                let v = data[base + ch];
35                if v < lower[ch] || v > upper[ch] {
36                    inside = false;
37                    break;
38                }
39            }
40            out[y * w + x] = if inside { 1.0 } else { 0.0 };
41        }
42    }
43    Ok(Tensor::from_vec(vec![h, w, 1], out)?)
44}
45
46/// Elementwise AND of two single-channel binary masks.
47///
48/// Values > 0.5 are treated as 1, otherwise 0.
49pub fn bitwise_and(a: &Tensor, b: &Tensor) -> Result<Tensor, ImgProcError> {
50    if a.shape() != b.shape() {
51        return Err(ImgProcError::ShapeMismatch {
52            expected: a.shape().to_vec(),
53            got: b.shape().to_vec(),
54        });
55    }
56    let out: Vec<f32> = a
57        .data()
58        .iter()
59        .zip(b.data().iter())
60        .map(|(av, bv)| if *av > 0.5 && *bv > 0.5 { 1.0 } else { 0.0 })
61        .collect();
62    Ok(Tensor::from_vec(a.shape().to_vec(), out)?)
63}
64
65/// Elementwise OR of two single-channel binary masks.
66pub fn bitwise_or(a: &Tensor, b: &Tensor) -> Result<Tensor, ImgProcError> {
67    if a.shape() != b.shape() {
68        return Err(ImgProcError::ShapeMismatch {
69            expected: a.shape().to_vec(),
70            got: b.shape().to_vec(),
71        });
72    }
73    let out: Vec<f32> = a
74        .data()
75        .iter()
76        .zip(b.data().iter())
77        .map(|(av, bv)| if *av > 0.5 || *bv > 0.5 { 1.0 } else { 0.0 })
78        .collect();
79    Ok(Tensor::from_vec(a.shape().to_vec(), out)?)
80}
81
82/// Elementwise NOT of a single-channel binary mask.
83pub fn bitwise_not(a: &Tensor) -> Result<Tensor, ImgProcError> {
84    let out: Vec<f32> = a
85        .data()
86        .iter()
87        .map(|v| if *v > 0.5 { 0.0 } else { 1.0 })
88        .collect();
89    Ok(Tensor::from_vec(a.shape().to_vec(), out)?)
90}
91
92/// Alpha-blend two images: `dst = alpha * a + beta * b + gamma`.
93///
94/// Both images must have the same shape.
95pub fn add_weighted(
96    a: &Tensor,
97    alpha: f32,
98    b: &Tensor,
99    beta: f32,
100    gamma: f32,
101) -> Result<Tensor, ImgProcError> {
102    if a.shape() != b.shape() {
103        return Err(ImgProcError::ShapeMismatch {
104            expected: a.shape().to_vec(),
105            got: b.shape().to_vec(),
106        });
107    }
108    let out: Vec<f32> = a
109        .data()
110        .iter()
111        .zip(b.data().iter())
112        .map(|(av, bv)| alpha * av + beta * bv + gamma)
113        .collect();
114    Ok(Tensor::from_vec(a.shape().to_vec(), out)?)
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn test_in_range() {
123        let img = Tensor::from_vec(
124            vec![2, 2, 3],
125            vec![0.5, 0.5, 0.5, 0.1, 0.1, 0.1, 0.9, 0.9, 0.9, 0.5, 0.5, 0.5],
126        )
127        .unwrap();
128        let mask = in_range(&img, &[0.4, 0.4, 0.4], &[0.6, 0.6, 0.6]).unwrap();
129        assert_eq!(mask.shape(), &[2, 2, 1]);
130        assert_eq!(mask.data(), &[1.0, 0.0, 0.0, 1.0]);
131    }
132
133    #[test]
134    fn test_bitwise_and_or_not() {
135        let a = Tensor::from_vec(vec![4], vec![1.0, 1.0, 0.0, 0.0]).unwrap();
136        let b = Tensor::from_vec(vec![4], vec![1.0, 0.0, 1.0, 0.0]).unwrap();
137        assert_eq!(bitwise_and(&a, &b).unwrap().data(), &[1.0, 0.0, 0.0, 0.0]);
138        assert_eq!(bitwise_or(&a, &b).unwrap().data(), &[1.0, 1.0, 1.0, 0.0]);
139        assert_eq!(bitwise_not(&a).unwrap().data(), &[0.0, 0.0, 1.0, 1.0]);
140    }
141
142    #[test]
143    fn test_add_weighted() {
144        let a = Tensor::from_vec(vec![3], vec![1.0, 2.0, 3.0]).unwrap();
145        let b = Tensor::from_vec(vec![3], vec![4.0, 5.0, 6.0]).unwrap();
146        let out = add_weighted(&a, 0.5, &b, 0.5, 0.0).unwrap();
147        assert_eq!(out.data(), &[2.5, 3.5, 4.5]);
148    }
149}