Closing up label gaps

2 files changed, 86 insertions, 9 deletions

diff --git a/src/lib/seg/util.c b/src/lib/seg/util.c
index 22a75cc..530c326 100644
--- a/src/lib/seg/util.c
+++ b/src/lib/seg/util.c
@@ -4,6 +4,7 @@
 #include <assert.h>
 #include <stdio.h>
 #include <stdlib.h>
+#include <string.h>
 
 // Convert x,y coords to linear coordinate
 size_t xy_to_coord(size_t x, size_t y, uint32_t width, uint32_t height)
@@ -93,6 +94,81 @@ uint16_t* dilate(uint16_t* mask, uint32_t width, uint32_t height)
   return new_mask;
 }
 
+// Erode masks by one 4-connected pixel
+uint16_t* erode(uint16_t* mask, uint32_t width, uint32_t height)
+{
+  uint16_t *new_mask = (uint16_t*)calloc(width*height,sizeof(uint16_t));
+  for (size_t y = 0; y < height; y++) {
+    for (size_t x = 0; x < width; x++) {
+      size_t current_position = xy_to_coord(x, y, width, height);
+      if (mask[current_position] != 0) {
+	new_mask[current_position] = mask[current_position];
+	continue;
+      }
+      bool_t erode = FALSE;
+      size_t proposed_position;
+      if (x != 0) {
+	proposed_position = xy_to_coord(x-1, y, width, height);
+	if (mask[proposed_position] == 0) {
+	  new_mask[current_position] = mask[proposed_position];
+	  continue;
+	}
+      }
+      if ((x+1) != width) {
+	proposed_position = xy_to_coord(x+1, y, width, height);
+	if (mask[proposed_position] == 0) {
+	  new_mask[current_position] = mask[proposed_position];
+	  continue;
+	}
+      }
+      if (y != 0) {
+	proposed_position = xy_to_coord(x, y-1, width, height);
+	if (mask[proposed_position] == 0) {
+	  new_mask[current_position] = mask[proposed_position];
+	  continue;
+	}
+      }
+      if ((y+1) != height) {
+	proposed_position = xy_to_coord(x, y+1, width, height);
+	if (mask[proposed_position] == 0) {
+	  new_mask[current_position] = mask[proposed_position];
+	  continue;
+	}
+      }
+    }
+  }
+  return new_mask;
+}
+
+// Close up masks by N-pixels
+uint16_t* closeup(uint16_t* mask, uint32_t width, uint32_t height, size_t num_pixels)
+{
+  uint16_t *new_mask = (uint16_t*)calloc(width*height,sizeof(uint16_t));
+  memcpy(new_mask, mask, width*height*sizeof(uint16_t));
+  for (size_t count = 0; count < num_pixels; count++) {
+    uint16_t *new_labels = dilate(new_mask, width, height);
+    free(new_mask);
+    new_mask = new_labels;
+  }
+  for (size_t count = 0; count < num_pixels; count++) {
+    uint16_t *new_labels = erode(new_mask, width, height);
+    free(new_mask);
+    new_mask = new_labels;
+  }
+  // Retain original mask at the very least
+  for (size_t y = 0; y < height; y++) {
+    for (size_t x = 0; x < width; x++) {
+      size_t coord = x + y*width;
+      if (mask[coord] != 0) {
+	if (new_mask[coord] != mask[coord]) {
+	  new_mask[coord] = mask[coord];
+	}
+      }
+    }
+  }
+  return new_mask;
+}
+
 // Combine Label Masks
 // For all empty spaces in the destination, put the extra label if it exists
 // Allocates an array if destination is unallocated
diff --git a/src/main.c b/src/main.c
index 5746cd8..7c914e5 100644
--- a/src/main.c
+++ b/src/main.c
@@ -17,7 +17,7 @@
 
 #define OFFSET 16
 
-#define N_DILATIONS 0
+#define N_DILATIONS 10
 
 #define MIN_AREA 200
 #define MIN_PERIMETER 50
@@ -104,14 +104,6 @@ int main(int argc, char** argv)
   free(labels);
   printf("%u labels found\n", starting_label-1);
   printf("Mask dimensions: %u %u\n", width, height);
-  TIME(ts_start);
-  for (uint16_t count = 0; count < N_DILATIONS; count++) {
-    uint16_t *new_labels = dilate(masks, width, height);
-    free(masks);
-    masks = new_labels;
-  }
-  TIME(ts_end);
-  printf("Dilation took %f ms\n", 1000*diff_time(&ts_end, &ts_start));
   // Get the area/ perimeter of each label
   root = NULL;
   TIME(ts_info_start);
@@ -191,6 +183,15 @@ int main(int argc, char** argv)
 #endif
   free_avl_tree_nodes(root);
 
+  TIME(ts_start);
+  uint16_t *new_labels = closeup(masks, width, height, N_DILATIONS);
+  if (new_labels != NULL) {
+    free(masks);
+    masks = new_labels;
+  }
+  TIME(ts_end);
+  printf("Dilation took %f ms\n", 1000*diff_time(&ts_end, &ts_start));
+
   //-----------------------------------------------
 
 #ifdef VISUAL