1 files changed, 37 insertions, 0 deletions
diff --git a/src/lib/seg/mask_data.c b/src/lib/seg/mask_data.c
index 2f3a51a..8c4b037 100644
--- a/src/lib/seg/mask_data.c
+++ b/src/lib/seg/mask_data.c
@@ -1,4 +1,5 @@
 #include <lib/seg/mask_data.h>
+#include <lib/seg/util.h>
 
 #include <stdio.h>
 
@@ -224,3 +225,39 @@ struct AVLNode* get_small_labels(struct AVLNode* removal_tree, struct AVLNode* l
   }
   return return_tree;
 }
+
+// Get mask label data
+struct AVLNode* get_mask_data(uint16_t* masks, uint32_t width, uint32_t height)
+{
+  struct AVLNode* root = NULL;
+  for (size_t y = 0; y < height; y++) {
+    for (size_t x = 0; x < width; x++) {
+      size_t coord = x + y*width;
+      if (masks[coord] != 0) {
+	root = increase_label_area_alloc(root, masks[coord]);
+	if (is_on_mask_boundary(masks, width, height, x, y)) {
+	  increase_label_perimeter(root, masks[coord]);
+	}
+      }
+    }
+  }
+  return root;
+}
+
+// Filter out small masks in mask
+void filter_small_masks(uint16_t* masks, uint32_t width, uint32_t height, size_t min_area, size_t min_perimeter)
+{
+  struct AVLNode* root = get_mask_data(masks, width, height);
+  struct AVLNode* small_label_tree = NULL;
+  small_label_tree = get_small_labels(NULL, root, min_area, min_perimeter);
+  for (size_t y = 0; y < height; y++) {
+    for (size_t x = 0; x < width; x++) {
+      size_t coord = x + y*width;
+      if (in_uint16_t_tree(small_label_tree, masks[coord])) {
+	masks[coord] = 0;
+      }
+    }
+  }
+  free_avl_tree(small_label_tree);
+  free_avl_tree_nodes(root);
+}