Cache unionfind structure across detect calls and lazy-initialize size array

Author: bouk

apriltag.c

@@ -411,6 +411,8 @@ void apriltag_detector_destroy(apriltag_detector_t *td)
     apriltag_detector_clear_families(td);
 
     zarray_destroy(td->tag_families);
+    if (td->cached_uf)
+        unionfind_destroy(td->cached_uf);
     free(td);
 }
 

apriltag.h

@@ -39,6 +39,7 @@ extern "C" {
 #include "common/workerpool.h"
 #include "common/timeprofile.h"
 #include "common/pthreads_cross.h"
+#include "common/unionfind.h"
 
 #define APRILTAG_TASKS_PER_THREAD_TARGET 10
 
@@ -188,6 +189,9 @@ struct apriltag_detector
 
     // Used for thread safety.
     pthread_mutex_t mutex;
+
+    // Cached unionfind structure (reused across detect calls)
+    unionfind_t *cached_uf;
 };
 
 // Represents the detection of a tag. These are returned to the user

apriltag_quad_thresh.c

@@ -1511,7 +1511,17 @@ image_u8_t *threshold_bayer(apriltag_detector_t *td, image_u8_t *im)
 }
 
 unionfind_t* connected_components(apriltag_detector_t *td, image_u8_t* threshim, int w, int h, int ts) {
-    unionfind_t *uf = unionfind_create(w * h);
+    uint32_t maxid = w * h;
+    if (td->cached_uf) {
+        if (td->cached_uf->maxid >= maxid) {
+            unionfind_reset(td->cached_uf);
+        } else {
+            unionfind_resize(td->cached_uf, maxid);
+        }
+    } else {
+        td->cached_uf = unionfind_create(maxid);
+    }
+    unionfind_t *uf = td->cached_uf;
 
     if (td->nthreads <= 1) {
         do_unionfind_first_line(uf, threshim, w, ts);
@@ -2004,8 +2014,6 @@ zarray_t *apriltag_quad_thresh(apriltag_detector_t *td, image_u8_t *im)
 
     timeprofile_stamp(td->tp, "fit quads to clusters");
 
-    unionfind_destroy(uf);
-
     for (int i = 0; i < zarray_size(clusters); i++) {
         zarray_t *cluster;
         zarray_get(clusters, i, &cluster);

common/unionfind.h

@@ -44,14 +44,28 @@ struct unionfind
     uint32_t *size;
 };
 
+static inline void unionfind_reset(unionfind_t *uf)
+{
+    memset(uf->parent, 0xff, (uf->maxid+1) * sizeof(uint32_t));
+    // uf->size is lazily initialized in unionfind_get_representative
+}
+
 static inline unionfind_t *unionfind_create(uint32_t maxid)
 {
     unionfind_t *uf = (unionfind_t*) calloc(1, sizeof(unionfind_t));
     uf->maxid = maxid;
     uf->parent = (uint32_t *) malloc((maxid+1) * sizeof(uint32_t) * 2);
-    memset(uf->parent, 0xff, (maxid+1) * sizeof(uint32_t));
     uf->size = uf->parent + (maxid+1);
-    memset(uf->size, 0, (maxid+1) * sizeof(uint32_t));
+    unionfind_reset(uf);
+    return uf;
+}
+
+static inline unionfind_t *unionfind_resize(unionfind_t *uf, uint32_t maxid)
+{
+    uf->maxid = maxid;
+    uf->parent = (uint32_t *) realloc(uf->parent, (maxid+1) * sizeof(uint32_t) * 2);
+    uf->size = uf->parent + (maxid+1);
+    unionfind_reset(uf);
     return uf;
 }
 
@@ -85,6 +99,7 @@ static inline uint32_t unionfind_get_representative(unionfind_t *uf, uint32_t id
     // unititialized node, so set to self
     if (uf->parent[id] == 0xffffffff) {
         uf->parent[id] = id;
+        uf->size[id] = 0;
         return id;
     }
 

test/CMakeLists.txt

@@ -48,3 +48,11 @@ endif()
 
 add_test(NAME test_quick_decode COMMAND test_quick_decode)
 
+add_executable(test_multiple_sizes test_multiple_sizes.c)
+target_link_libraries(test_multiple_sizes ${PROJECT_NAME})
+
+add_test(NAME test_multiple_sizes
+         COMMAND $<TARGET_FILE:test_multiple_sizes> data/33369213973_9d9bb4cc96_c.jpg
+         WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+)
+

test/test_multiple_sizes.c

@@ -0,0 +1,56 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <apriltag.h>
+#include <tag36h11.h>
+#include <common/pjpeg.h>
+
+// This test verifies that the cached data inside the detector handle size changes
+
+int main(int argc, char *argv[])
+{
+    if (argc != 2) {
+        fprintf(stderr, "Usage: %s <image.jpg>\n", argv[0]);
+        return EXIT_FAILURE;
+    }
+
+    pjpeg_t *pjpeg = pjpeg_create_from_file(argv[1], 0, NULL);
+    if (pjpeg == NULL) {
+        fprintf(stderr, "Failed to load %s\n", argv[1]);
+        return EXIT_FAILURE;
+    }
+    image_u8_t *im = pjpeg_to_u8_baseline(pjpeg);
+
+    apriltag_family_t *tf = tag36h11_create();
+    apriltag_detector_t *td = apriltag_detector_create();
+    apriltag_detector_add_family(td, tf);
+    td->refine_edges = false;
+
+    // First detect with quad_decimate=2
+    td->quad_decimate = 2;
+    zarray_t *dets1 = apriltag_detector_detect(td, im);
+    int n1 = zarray_size(dets1);
+    printf("decimate=2: %d detections\n", n1);
+    apriltag_detections_destroy(dets1);
+
+    // Now detect with quad_decimate=1
+    td->quad_decimate = 1;
+    zarray_t *dets2 = apriltag_detector_detect(td, im);
+    int n2 = zarray_size(dets2);
+    printf("decimate=1: %d detections\n", n2);
+    apriltag_detections_destroy(dets2);
+
+    // Then detect with quad_decimate=2 again
+    td->quad_decimate = 2;
+    zarray_t *dets3 = apriltag_detector_detect(td, im);
+    int n3 = zarray_size(dets3);
+    printf("decimate=2: %d detections\n", n3);
+    apriltag_detections_destroy(dets3);
+
+    image_u8_destroy(im);
+    pjpeg_destroy(pjpeg);
+    apriltag_detector_destroy(td);
+    tag36h11_destroy(tf);
+
+    printf("PASS\n");
+    return EXIT_SUCCESS;
+}