webarkitlib_rs/
labeling.rs1use crate::types::{ARLabelInfo, ARdouble};
41use log::debug;
42
43pub const AR_LABELING_WORK_SIZE: usize = 1024 * 32;
44
45pub enum LabelingMode {
46 BlackRegion,
47 WhiteRegion,
48}
49
50pub enum ImageProcMode {
51 FrameImage,
52 FieldImage, }
54
55pub fn ar_labeling(
60 image: &[u8],
61 xsize: i32,
62 ysize: i32,
63 mode: LabelingMode,
64 thresh: u8,
65 proc_mode: ImageProcMode,
66 label_info: &mut ARLabelInfo,
67 _debug_mode: bool,
68) -> Result<(), &'static str> {
69
70 let is_region = |pixel: u8| -> bool {
71 match mode {
72 LabelingMode::BlackRegion => pixel <= thresh,
73 LabelingMode::WhiteRegion => pixel > thresh,
74 }
75 };
76
77 let lxsize: usize;
78 let lysize: usize;
79 let row_stride: usize;
80 let _col_stride: usize;
81
82 match proc_mode {
83 ImageProcMode::FrameImage => {
84 lxsize = xsize as usize;
85 lysize = ysize as usize;
86 row_stride = xsize as usize;
87 _col_stride = 1;
88 }
89 ImageProcMode::FieldImage => {
90 lxsize = (xsize / 2) as usize;
91 lysize = (ysize / 2) as usize;
92 row_stride = (xsize * 2) as usize; if ysize < 480 { _col_stride = 1;
95 } else {
96 _col_stride = 2; }
98 }
99 }
100
101 if image.len() < (ysize as usize * xsize as usize) {
102 return Err("Image buffer is too small for given dimensions.");
103 }
104
105 if label_info.label_image.len() < lxsize * lysize {
107 label_info.label_image.resize(lxsize * lysize, 0);
108 }
109 label_info.label_image.fill(0);
110
111 if label_info.work.len() < AR_LABELING_WORK_SIZE {
113 label_info.work.resize(AR_LABELING_WORK_SIZE, 0);
114 }
115 if label_info.work2.len() < AR_LABELING_WORK_SIZE * 7 {
116 label_info.work2.resize(AR_LABELING_WORK_SIZE * 7, 0);
117 }
118
119 let mut wk_max: usize = 0;
120 let work = &mut label_info.work;
121 let work2 = &mut label_info.work2;
122 let label_img = &mut label_info.label_image;
123
124 fn find(work: &mut [i32], mut i: i32) -> i32 {
126 let mut root = i;
127 while work[root as usize - 1] != root {
128 root = work[root as usize - 1];
129 }
130 let mut curr = i;
132 while work[curr as usize - 1] != root {
133 let next = work[curr as usize - 1];
134 work[curr as usize - 1] = root;
135 curr = next;
136 }
137 root
138 }
139
140 fn do_union(work: &mut [i32], m: i32, n: i32) -> i32 {
142 let root_m = find(work, m);
143 let root_n = find(work, n);
144 if root_m < root_n {
145 work[root_n as usize - 1] = root_m;
146 root_m
147 } else {
148 work[root_m as usize - 1] = root_n;
149 root_n
150 }
151 }
152
153 for j in 1..lysize - 1 {
155 for i in 1..lxsize - 1 {
156 let source_idx = match proc_mode {
157 ImageProcMode::FrameImage => j * row_stride + i,
158 ImageProcMode::FieldImage => (j * 2 + 1) * xsize as usize + (i * 2),
159 };
160 let pixel = image[source_idx];
161
162 let p_idx = j * lxsize + i;
163
164 if is_region(pixel) {
165 let left_val = label_img[p_idx - 1];
166 let up_val = label_img[p_idx - lxsize];
167 let up_left = label_img[p_idx - lxsize - 1];
168 let up_right = label_img[p_idx - lxsize + 1];
169
170 if up_val > 0 {
171 label_img[p_idx] = up_val;
172 let l = (find(work, up_val as i32) as usize - 1) * 7;
173 work2[l + 0] += 1;
174 work2[l + 1] += i as i32;
175 work2[l + 2] += j as i32;
176 work2[l + 6] = j as i32;
177 } else if up_right > 0 {
178 if up_left > 0 {
179 let final_label = do_union(work, up_right as i32, up_left as i32);
180 label_img[p_idx] = final_label as crate::types::ARLabelingLabelType;
181
182 let l = (final_label as usize - 1) * 7;
183 work2[l + 0] += 1;
184 work2[l + 1] += i as i32;
185 work2[l + 2] += j as i32;
186 work2[l + 6] = j as i32;
187 } else if left_val > 0 {
188 let final_label = do_union(work, up_right as i32, left_val as i32);
189 label_img[p_idx] = final_label as crate::types::ARLabelingLabelType;
190
191 let l = (final_label as usize - 1) * 7;
192 work2[l + 0] += 1;
193 work2[l + 1] += i as i32;
194 work2[l + 2] += j as i32;
195 work2[l + 6] = j as i32;
196 } else {
197 label_img[p_idx] = up_right;
198 let l = (find(work, up_right as i32) as usize - 1) * 7;
199 work2[l + 0] += 1;
200 work2[l + 1] += i as i32;
201 work2[l + 2] += j as i32;
202 if work2[l + 3] > i as i32 { work2[l + 3] = i as i32; }
203 work2[l + 6] = j as i32;
204 }
205 } else if up_left > 0 {
206 label_img[p_idx] = up_left;
207 let l = (find(work, up_left as i32) as usize - 1) * 7;
208 work2[l + 0] += 1;
209 work2[l + 1] += i as i32;
210 work2[l + 2] += j as i32;
211 if work2[l + 4] < i as i32 { work2[l + 4] = i as i32; }
212 work2[l + 6] = j as i32;
213 } else if left_val > 0 {
214 label_img[p_idx] = left_val;
215 let l = (find(work, left_val as i32) as usize - 1) * 7;
216 work2[l + 0] += 1;
217 work2[l + 1] += i as i32;
218 work2[l + 2] += j as i32;
219 if work2[l + 4] < i as i32 { work2[l + 4] = i as i32; }
220 } else {
221 wk_max += 1;
222 if wk_max > AR_LABELING_WORK_SIZE {
223 return Err("Labeling work array overflow");
224 }
225 work[wk_max - 1] = wk_max as i32;
226 label_img[p_idx] = wk_max as crate::types::ARLabelingLabelType;
227
228 let l = (wk_max - 1) * 7;
229 work2[l + 0] = 1; work2[l + 1] = i as i32; work2[l + 2] = j as i32; work2[l + 3] = i as i32; work2[l + 4] = i as i32; work2[l + 5] = j as i32; work2[l + 6] = j as i32; }
237 } else {
238 label_img[p_idx] = 0;
239 }
240 }
241 }
242
243 let mut num_labels = 0;
245 for i in 1..=wk_max {
246 if work[i - 1] == i as i32 {
247 num_labels += 1;
248 work[i - 1] = num_labels;
249 } else {
250 work[i - 1] = work[work[i - 1] as usize - 1]; }
252 }
253
254 for i in 1..=wk_max {
256 let mut root = i as i32;
257 while work[root as usize - 1] > num_labels || work[work[root as usize - 1] as usize - 1] != work[root as usize - 1] {
258 root = work[root as usize - 1];
260 }
261 work[i - 1] = work[root as usize - 1];
262 }
263
264 label_info.label_num = num_labels;
265 if label_info.label_num == 0 {
266 return Ok(());
267 }
268
269 if label_info.area.len() < num_labels as usize { label_info.area.resize(num_labels as usize, 0); }
271 if label_info.pos.len() < num_labels as usize { label_info.pos.resize(num_labels as usize, [0.0; 2]); }
272 if label_info.clip.len() < num_labels as usize { label_info.clip.resize(num_labels as usize, [0; 4]); }
273
274 label_info.area.fill(0);
275 label_info.pos.fill([0.0; 2]);
276 for clip in label_info.clip.iter_mut() {
277 clip[0] = lxsize as i32;
278 clip[1] = 0;
279 clip[2] = lysize as i32;
280 clip[3] = 0;
281 }
282
283 for i in 0..wk_max {
284 let dest_label = work[i] as usize - 1;
285
286 let area = work2[i * 7 + 0];
287 let pos_x = work2[i * 7 + 1];
288 let pos_y = work2[i * 7 + 2];
289 let clip_xmin = work2[i * 7 + 3];
290 let clip_xmax = work2[i * 7 + 4];
291 let clip_ymin = work2[i * 7 + 5];
292 let clip_ymax = work2[i * 7 + 6];
293
294 label_info.area[dest_label] += area;
295 label_info.pos[dest_label][0] += pos_x as ARdouble;
296 label_info.pos[dest_label][1] += pos_y as ARdouble;
297
298 if label_info.clip[dest_label][0] > clip_xmin { label_info.clip[dest_label][0] = clip_xmin; }
299 if label_info.clip[dest_label][1] < clip_xmax { label_info.clip[dest_label][1] = clip_xmax; }
300 if label_info.clip[dest_label][2] > clip_ymin { label_info.clip[dest_label][2] = clip_ymin; }
301 if label_info.clip[dest_label][3] < clip_ymax { label_info.clip[dest_label][3] = clip_ymax; }
302 }
303
304 for i in 0..num_labels as usize {
305 if label_info.area[i] > 0 {
306 label_info.pos[i][0] /= label_info.area[i] as ARdouble;
307 label_info.pos[i][1] /= label_info.area[i] as ARdouble;
308 }
309 }
310
311 Ok(())
312}
313
314#[cfg(test)]
315mod tests {
316 use super::*;
317
318 #[test]
319 fn test_ar_labeling_simple_square() {
320 let mut image = vec![255u8; 64];
322 for j in 2..6 {
323 for i in 2..6 {
324 image[j * 8 + i] = 0;
325 }
326 }
327
328 let mut info = ARLabelInfo::default();
329
330 ar_labeling(
331 &image,
332 8,
333 8,
334 LabelingMode::BlackRegion,
335 100, ImageProcMode::FrameImage,
337 &mut info,
338 false
339 ).unwrap();
340
341 assert_eq!(info.label_num, 1);
342
343 assert_eq!(info.area[0], 16);
345
346 assert!((info.pos[0][0] - 3.5).abs() < f64::EPSILON);
348 assert!((info.pos[0][1] - 3.5).abs() < f64::EPSILON);
349
350 assert_eq!(info.clip[0][0], 2); assert_eq!(info.clip[0][1], 5); assert_eq!(info.clip[0][2], 2); assert_eq!(info.clip[0][3], 5); }
356}