Filter small labels

3 files changed, 110 insertions, 0 deletions

diff --git a/src/lib/algo/avl_tree.c b/src/lib/algo/avl_tree.c
index 98cd18b..db4fa4b 100644
--- a/src/lib/algo/avl_tree.c
+++ b/src/lib/algo/avl_tree.c
@@ -56,6 +56,75 @@ struct AVLNode* left_rotate(struct AVLNode* parent)
   return child1;
 }
 
+// Create AVL node
+struct AVLNode* create_avl_node(void* data, bool_t (*compare)(void*, void*))
+{
+  struct AVLNode* node = (struct AVLNode*)malloc(sizeof(struct AVLNode));
+  if (node == NULL) {
+    return NULL;
+  }
+  node->data = data;
+  node->compare = compare;
+  node->left = NULL;
+  node->right = NULL;
+  node->height = 1; // Leaf initially
+  return node;
+}
+
+// Insert data into AVL tree
+struct Result avl_insert(struct AVLNode* node, void* data, bool_t (*compare)(void*, void*))
+{
+  struct Result result;
+  // 1. Standard BST insertion
+  if (node == NULL) {
+    return (struct Result) {create_avl_node(data, compare), TRUE};
+  }
+
+  struct MaskData *node_data = (struct MaskData*)node->data;
+  if (node->compare(data, node_data)) {
+    result = avl_insert(node->left, data, compare);
+    if (!result.success) {
+      fprintf(stderr, "Failed to insert!");
+      return result;
+    }
+    node->left = (struct AVLNode*)result.data;
+  } else if (node->compare(node->data, data)) {
+    result = avl_insert(node->right, data, compare);
+    if (!result.success) {
+      fprintf(stderr, "Failed to insert!");
+      return result;
+    }
+    node->right = (struct AVLNode*)result.data;
+  } else {
+    return (struct Result) {node, FALSE};
+  }
+
+  // 2. Update height of the ancestor node
+  node->height = 1 + max_height(get_height(node->left), get_height(node->right));
+
+  ssize_t balance = get_balance_factor(node);
+
+  // 4. If the node becomes unbalanced
+
+  // LeftLeft
+  if ((balance > 1) && node->compare(data, node->left->data)) {
+    return (struct Result) {right_rotate(node), TRUE};
+  }
+  // RightRight
+  if ((balance < -1) && node->compare(node->right->data, data)) {
+    return (struct Result) {left_rotate(node), TRUE};
+  }
+  // LeftRight
+  if ((balance > 1) && node->compare(node->left->data, data)) {
+    return (struct Result) {right_rotate(node), TRUE};
+  }
+  // RightLeft
+  if ((balance < -1) && node->compare(data,node->right->data)) {
+    return (struct Result) {left_rotate(node), TRUE};
+  }
+  return (struct Result) {node, TRUE};
+}
+
 // In-order traversal print pointer
 void print_in_order(struct AVLNode* root)
 {
diff --git a/src/lib/seg/mask_data.c b/src/lib/seg/mask_data.c
index 00e4807..2ff58a8 100644
--- a/src/lib/seg/mask_data.c
+++ b/src/lib/seg/mask_data.c
@@ -173,3 +173,39 @@ struct AVLNode* increase_label_perimeter_alloc(struct AVLNode* root, uint16_t la
   }
   return new_root;
 }
+
+// Comparison of uint16_ts
+bool_t compare_uint16_t(uint16_t* s1, uint16_t* s2)
+{
+  return *s1 < *s2;
+}
+
+// In-order traversal print pointer
+void print_in_order_uint16_t(struct AVLNode* root)
+{
+  if (root != NULL) {
+    print_in_order_uint16_t(root->left);
+    printf("%d ", *((uint16_t*)root->data));
+    print_in_order_uint16_t(root->right);
+  }
+}
+
+// Filter out small masks
+//  Assumption: Contiguous labeling
+struct AVLNode* get_small_labels(struct AVLNode* removal_tree, struct AVLNode* label_tree, size_t min_area, size_t min_perimeter)
+{
+  struct AVLNode* return_tree = removal_tree;
+  if (label_tree != NULL) {
+    return_tree = get_small_labels(return_tree, label_tree->left, min_area, min_perimeter);
+    struct MaskData* node_data = (struct MaskData*)label_tree->data;
+    if ((node_data->area < min_area) || (node_data->perimeter < min_perimeter)) {
+      // Insert
+      struct Result result = avl_insert(return_tree, &node_data->label, (bool_t (*)(void*,void*))compare_uint16_t);
+      if (result.success) {
+	return_tree = result.data;
+      }
+    }
+    return_tree = get_small_labels(return_tree, label_tree->right, min_area, min_perimeter);
+  }
+  return return_tree;
+}
diff --git a/src/main.c b/src/main.c
index e199c43..8ee035f 100644
--- a/src/main.c
+++ b/src/main.c
@@ -180,6 +180,7 @@ int main(int argc, char** argv)
   CloseWindow();
 #else
   if (masks != NULL) {
+    // Get the area/ perimeter of each label
     TIME(ts_start);
     for (size_t y = 0; y < height; y++) {
       for (size_t x = 0; x < width; x++) {
@@ -198,6 +199,10 @@ int main(int argc, char** argv)
     print_label(root);
     printf("\n");
 #endif
+    // Get the smallest labels
+    struct AVLNode* small_label_tree = NULL;
+    small_label_tree = get_small_labels(NULL, root, 50, 50);
+    free_avl_tree(small_label_tree);
     free_avl_tree_nodes(root);
     root = NULL;
     write_array("../out.bin", masks, width*height*sizeof(uint16_t));