Closing up label gaps

1 files changed, 76 insertions, 0 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