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

Author: bouk

apriltag.c

@@ -411,6 +411,9 @@ 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;
     }