feat: Inverse watershed to labelling pixels in findPeaks2DRegion (#12)

* feat: add inverse watershed to labelling

* feat: add labelling options to findPeaks2DRegion

src/__tests__/ccLabelling.test.js

@@ -1,4 +1,4 @@
-import { ccLabeling as labeling } from '../ccLabeling';
+import { floodFillLabelling as labeling } from '../floodFillLabelling';
 
 let rows = 8;
 let cols = 8;

src/drainLabelling.js

@@ -0,0 +1,78 @@
+import DisjointSet from 'ml-disjoint-set';
+
+const direction8X = [-1, -1, 0, 1, -1, 0, 1, 1];
+const direction8Y = [0, -1, -1, -1, 1, 1, 1, 0];
+const neighbours8 = [null, null, null, null, null, null, null, null];
+
+const direction4X = [-1, 0, 1, 0];
+const direction4Y = [0, -1, 0, 1];
+const neighbours4 = [null, null, null, null];
+
+export function drainLabelling(data, mask, options = {}) {
+  const { neighbours = 8, width, height } = options;
+
+  let directionX;
+  let directionY;
+  let neighboursList;
+  if (neighbours === 8) {
+    directionX = direction8X;
+    directionY = direction8Y;
+    neighboursList = neighbours8;
+  } else if (neighbours === 4) {
+    directionX = direction4X;
+    directionY = direction4Y;
+    neighboursList = neighbours4;
+  } else {
+    throw new RangeError(`unsupported neighbours count: ${neighbours}`);
+  }
+
+  let sorted = new Array(height * width);
+  for (let i = 0, index = 0; i < height; i++) {
+    for (let j = 0; j < width; j++, index++) {
+      sorted[index] = { value: data[index], row: i, col: j, mask: mask[index] };
+    }
+  }
+
+  sorted.sort((a, b) => a.value - b.value);
+
+  const size = mask.length;
+  const labels = new Array(size);
+  const pixels = new Int16Array(size);
+  const linked = new DisjointSet();
+
+  for (let i = 0, currentLabel = 1; i < mask.length; i++) {
+    let element = sorted[i];
+    if (!element.mask) continue;
+
+    let { row, col, value } = element;
+    let index = col + row * width;
+    let label = labels[index];
+    if (!label) {
+      labels[index] = linked.add(currentLabel++);
+    }
+
+    for (let k = 0; k < neighboursList.length; k++) {
+      let ii = col + directionX[k];
+      let jj = row + directionY[k];
+      if (ii >= 0 && jj >= 0 && ii < width && jj < height) {
+        let neighbor = labels[ii + jj * width];
+        if (!neighbor) {
+          let neighborValue = data[ii + jj * width];
+          if (value < neighborValue) {
+            labels[ii + jj * width] = labels[index];
+          }
+        }
+      }
+    }
+  }
+
+  for (let j = 0; j < height; j++) {
+    for (let i = 0; i < width; i++) {
+      let index = i + j * width;
+      if (mask[index]) {
+        pixels[index] = linked.find(labels[index]).value;
+      }
+    }
+  }
+  return pixels;
+}

src/ccLabeling.js → src/floodFillLabelling.js

@@ -8,7 +8,7 @@ const direction8X = [-1, -1, 0, 1];
 const direction8Y = [0, -1, -1, -1];
 const neighbours8 = [null, null, null, null];
 
-export function ccLabeling(mask, width, height, options) {
+export function floodFillLabelling(mask, width, height, options) {
   options = options || {};
   const neighbours = options.neighbours || 8;
 
@@ -38,7 +38,6 @@ export function ccLabeling(mask, width, height, options) {
       // true means out of background
       let smallestNeighbor = null;
       index = i + j * width;
-
       if (mask[index]) {
         for (let k = 0; k < neighboursList.length; k++) {
           let ii = i + directionX[k];
@@ -71,7 +70,6 @@ export function ccLabeling(mask, width, height, options) {
       }
     }
   }
-
   for (let j = 0; j < height; j++) {
     for (let i = 0; i < width; i++) {
       index = i + j * width;

src/index.js

@@ -1,6 +1,7 @@
 import * as convolution from 'ml-matrix-convolution';
 
-import { ccLabeling as labeling } from './ccLabeling';
+import { drainLabelling } from './drainLabelling';
+import { floodFillLabelling } from './floodFillLabelling';
 
 const smallFilter = [
   [0, 0, 1, 2, 2, 2, 1, 0, 0],
@@ -19,7 +20,8 @@ const smallFilter = [
  * @param {Array<Array>} input - matrix to get the local maxima
  * @param {Object} [options = {}] - options of the method.
  * @param {Array<Array>} [options.nStdDev = 3] - number of times of the standard deviations for the noise level.Float64Array
- * @param {Array<Array>} [options.kernel] - kernel to the convolution step
+ * @param {Array<Array>} [options.kernel] - kernel to the convolution step.
+ * @param {string} [options.labelling = 'drain'] - select the labelling algorithm to assign pixels.
  * @param {Array<Array>} [options.originalData] - original data useful when the original matrix has values and the input matrix has absolute ones
  * @param {Array<Array>} [options.filteredData] - convoluted data, if it is defined the convolution step is skipped
  */
@@ -31,10 +33,11 @@ export function findPeaks2DRegion(input, options = {}) {
     filteredData,
     rows: nRows,
     cols: nCols,
+    labelling = 'drain',
   } = options;
 
   let flatten = convolution.matrix2Array(input);
-  let inputData = flatten.data;
+  let data = flatten.data;
 
   if (!nRows || !nCols) {
     nRows = flatten.rows;
@@ -46,7 +49,7 @@ export function findPeaks2DRegion(input, options = {}) {
   }
 
   let cs = filteredData;
-  if (!cs) cs = convolution.fft(inputData, kernel, options);
+  if (!cs) cs = convolution.fft(data, kernel, options);
 
   let threshold = 0;
   for (let i = nCols * nRows - 2; i >= 0; i--) {
@@ -61,13 +64,27 @@ export function findPeaks2DRegion(input, options = {}) {
     }
   }
 
-  let pixels = labeling(bitmask, nCols, nRows, { neighbours: 8 });
-  let peakList = extractPeaks(pixels, {
-    data: inputData,
+  let pixels;
+  switch (labelling.toLowerCase()) {
+    case 'drain':
+      pixels = drainLabelling(cs, bitmask, {
+        neighbours: 8,
+        width: nCols,
+        height: nRows,
+      });
+      break;
+    case 'floodfill':
+      pixels = floodFillLabelling(bitmask, nCols, nRows, { neighbours: 8 });
+      break;
+    default:
+      throw new Error(`labelling ${labelling} does not support`);
+  }
+
+  return extractPeaks(pixels, {
+    data,
     nCols,
     originalData,
   });
-  return peakList;
 }
 /**
  Detects all the 2D-peaks in the given spectrum based on the Max logic.
@@ -84,7 +101,7 @@ export function findPeaks2DMax(input, options) {
   } = options;
 
   let flatten = convolution.matrix2Array(input);
-  let inputData = flatten.data;
+  let data = flatten.data;
 
   if (!nRows || !nCols) {
     nRows = flatten.rows;
@@ -96,7 +113,7 @@ export function findPeaks2DMax(input, options) {
   }
 
   let cs = filteredData;
-  if (!cs) cs = convolution.fft(inputData, kernel, options);
+  if (!cs) cs = convolution.fft(data, kernel, options);
 
   let threshold = 0;
   for (let i = nCols * nRows - 2; i >= 0; i--) {
@@ -154,6 +171,7 @@ function extractPeaks(pixels, options) {
         tmp.x += col * data[i];
         tmp.y += row * data[i];
         tmp.z += originalData[i];
+        tmp.count++;
         if (col < tmp.minX) tmp.minX = col;
         if (col > tmp.maxX) tmp.maxX = col;
         if (row < tmp.minY) tmp.minY = row;
@@ -163,6 +181,7 @@ function extractPeaks(pixels, options) {
           x: col * data[i],
           y: row * data[i],
           z: originalData[i],
+          count: 1,
           minX: col,
           maxX: col,
           minY: row,
@@ -179,5 +198,6 @@ function extractPeaks(pixels, options) {
     peakList[i].x /= zValue;
     peakList[i].y /= zValue;
   }
+
   return peakList;
 }