1use yscv_tensor::Tensor;
4
5use super::super::ImgProcError;
6
7pub 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
46pub 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
65pub 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
82pub 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
92pub 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}