1 files changed, 20 insertions, 4 deletions

diff --git a/src/main.c b/src/main.c
index 3fb681a..95b3e6d 100644
--- a/src/main.c
+++ b/src/main.c
@@ -120,6 +120,11 @@ int main()
   //-----------------------------------------------
   void* buffer = malloc(STRIP_LENGTH*sizeof(uint8_t));
   uint8_t* raster_8 = calloc(width*height*channels,sizeof(uint8_t));
+  if (raster_8 == NULL) {
+    fprintf(stderr, "Memory allocation error\n");
+    TIFFClose(tif);
+    return 1;
+  }
   for (size_t y = 0; y < STRIP_COUNT; y++) {
     tmsize_t strip_size = TIFFReadRawStrip(tif, y, buffer, STRIP_LENGTH);
     assert(strip_size == STRIP_LENGTH);
@@ -130,17 +135,22 @@ int main()
   free(buffer);
   //-----------------------------------------------
 
+  struct timespec ts_start, ts_end;
+  timespec_get(&ts_start, TIME_UTC);
   //-----------------------------------------------
   //-FLOOD-FILL-SEGMENTATION-----------------------
   //-CONTIGUOUS-REGION-FINDING---------------------
   //-----------------------------------------------
-  struct timespec ts_start, ts_end;
-  char buff[100];
-  timespec_get(&ts_start, TIME_UTC);
   // Flood fill
   uint16_t starting_label = 1;
   uint16_t *labels = NULL;
   labels = (uint16_t*)calloc(width*height,sizeof(uint16_t));
+  if (labels == NULL) {
+    fprintf(stderr, "Memory allocation error\n");
+    free(raster_8);
+    TIFFClose(tif);
+    return 1;
+  }
   for (size_t y = 0; y < height; y++) {
     for (size_t x = 0; x < width; x++) {
       size_t coord = x + y*width;
@@ -149,25 +159,31 @@ int main()
       }
     }
   }
+  // (When treating this as RGBA, last bits should ensure opaque)
+  //  This assumes 4096 (2^12) > labels
   for (size_t y = 0; y < height; y++) {
     for (size_t x = 0; x < width; x++) {
       labels[x + y*width] |= 0xFF000000;
     }
   }
+  //-----------------------------------------------
   timespec_get(&ts_end, TIME_UTC);
   printf("Time difference: %.3fms\n", 1000*diff_timespec(&ts_end, &ts_start));
   printf("N_labels: %u\n", starting_label-1);
-  //-----------------------------------------------
 
   uint32_t *raster = (uint32_t*)_TIFFmalloc(width*height*sizeof(uint32_t));
   if (raster == NULL) {
     fprintf(stderr, "Memory allocation error\n");
+    free(labels);
+    free(raster_8);
     TIFFClose(tif);
     return 1;
   }
 
   if (!TIFFReadRGBAImage(tif, width, height, raster, 0)) {
     fprintf(stderr, "Failed to read TIFF image\n");
+    free(labels);
+    free(raster_8);
     _TIFFfree(raster);
     TIFFClose(tif);
     return 1;