Skip to main content

webarkitlib_rs/
labeling.rs

1/*
2 *  labeling.rs
3 *  WebARKitLib-rs
4 *
5 *  This file is part of WebARKitLib-rs - WebARKit.
6 *
7 *  WebARKitLib-rs is free software: you can redistribute it and/or modify
8 *  it under the terms of the GNU Lesser General Public License as published by
9 *  the Free Software Foundation, either version 3 of the License, or
10 *  (at your option) any later version.
11 *
12 *  WebARKitLib-rs is distributed in the hope that it will be useful,
13 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
14 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15 *  GNU Lesser General Public License for more details.
16 *
17 *  You should have received a copy of the GNU Lesser General Public License
18 *  along with WebARKitLib-rs.  If not, see <http://www.gnu.org/licenses/>.
19 *
20 *  As a special exception, the copyright holders of this library give you
21 *  permission to link this library with independent modules to produce an
22 *  executable, regardless of the license terms of these independent modules, and to
23 *  copy and distribute the resulting executable under terms of your choice,
24 *  provided that you also meet, for each linked independent module, the terms and
25 *  conditions of the license of that module. An independent module is a module
26 *  which is neither derived from nor based on this library. If you modify this
27 *  library, you may extend this exception to your version of the library, but you
28 *  are not obligated to do so. If you do not wish to do so, delete this exception
29 *  statement from your version.
30 *
31 *  Copyright 2026 WebARKit.
32 *
33 *  Author(s): Walter Perdan @kalwalt https://github.com/kalwalt
34 *
35 */
36
37//! Connected Component Labeling Utilities
38//! Translated from ARToolKit C headers (arLabeling.c, arLabelingSub.h)
39
40use 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, // Handles interlaced half-resolution fields
53}
54
55/// Perform connected-component analysis and labeling on an 8-bit luma image.
56///
57/// This collapses the ~14 macro variations of `arLabelingSub*` into one generic safe Rust function,
58/// parameterizing the behavior using enums and generic pixel-checking closures.
59pub 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; // skip every other row
93            if ysize < 480 { // Assuming 'height' in the snippet refers to ysize
94                _col_stride = 1;
95            } else {
96                _col_stride = 2; // skip every other column
97            }
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    // Ensure the label_image buffer is perfectly sized
106    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    // Initialize work equivalence arrays
112    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    // Helper for Union-Find find with path compression
125    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        // Path compression
131        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    // Helper for Union-Find union
141    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    // Scan the inner pixel region (skip boundaries 0, and max-1)
154    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;         // area
230                    work2[l + 1] = i as i32;  // pos[0]
231                    work2[l + 2] = j as i32;  // pos[1]
232                    work2[l + 3] = i as i32;  // clip[0] (xmin)
233                    work2[l + 4] = i as i32;  // clip[1] (xmax)
234                    work2[l + 5] = j as i32;  // clip[2] (ymin)
235                    work2[l + 6] = j as i32;  // clip[3] (ymax)
236                }
237            } else {
238                label_img[p_idx] = 0;
239            }
240        }
241    }
242
243    // Pass 2: Map equivalence table down to dense sequential indexes
244    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]; // This is fine for one level, but let's be thorough
251        }
252    }
253    
254    // Thoroughly flatten the equivalence table
255    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             // If not yet a finalized label index, follow parent
259             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    // Allocate memory and reset results for the second pass
270    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        // Create an 8x8 image where center 4x4 is black (value 0), rest is white (value 255)
321        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, // Threshold 100
336            ImageProcMode::FrameImage,
337            &mut info,
338            false
339        ).unwrap();
340
341        assert_eq!(info.label_num, 1);
342        
343        // The square is 4x4, so area should be 16
344        assert_eq!(info.area[0], 16);
345        
346        // Pos should be the center: average of 2, 3, 4, 5 = 3.5
347        assert!((info.pos[0][0] - 3.5).abs() < f64::EPSILON);
348        assert!((info.pos[0][1] - 3.5).abs() < f64::EPSILON);
349        
350        // Clip coords
351        assert_eq!(info.clip[0][0], 2); // xmin
352        assert_eq!(info.clip[0][1], 5); // xmax
353        assert_eq!(info.clip[0][2], 2); // ymin
354        assert_eq!(info.clip[0][3], 5); // ymax
355    }
356}