Update WatershedOperation (#607)

This is a breaking change. Any saved pipelines with a watershed operation will fail to load.

Author: SamCarlberg

core/src/main/java/edu/wpi/grip/core/operations/composite/WatershedOperation.java

@@ -10,6 +10,9 @@
 
 import com.google.common.collect.ImmutableList;
 
+import org.bytedeco.javacpp.opencv_core;
+
+import java.util.ArrayList;
 import java.util.List;
 
 import static org.bytedeco.javacpp.opencv_core.CV_32SC1;
@@ -19,15 +22,20 @@
 import static org.bytedeco.javacpp.opencv_core.Mat;
 import static org.bytedeco.javacpp.opencv_core.MatVector;
 import static org.bytedeco.javacpp.opencv_core.Point;
+import static org.bytedeco.javacpp.opencv_core.Point2f;
 import static org.bytedeco.javacpp.opencv_core.Scalar;
-import static org.bytedeco.javacpp.opencv_core.bitwise_not;
+import static org.bytedeco.javacpp.opencv_imgproc.CV_CHAIN_APPROX_TC89_KCOS;
 import static org.bytedeco.javacpp.opencv_imgproc.CV_FILLED;
+import static org.bytedeco.javacpp.opencv_imgproc.CV_RETR_EXTERNAL;
 import static org.bytedeco.javacpp.opencv_imgproc.circle;
 import static org.bytedeco.javacpp.opencv_imgproc.drawContours;
+import static org.bytedeco.javacpp.opencv_imgproc.findContours;
+import static org.bytedeco.javacpp.opencv_imgproc.pointPolygonTest;
 import static org.bytedeco.javacpp.opencv_imgproc.watershed;
 
 /**
- * GRIP {@link Operation} for {@link org.bytedeco.javacpp.opencv_imgproc#watershed}.
+ * GRIP {@link Operation} for
+ * {@link org.bytedeco.javacpp.opencv_imgproc#watershed}.
  */
 public class WatershedOperation implements Operation {
 
@@ -40,21 +48,25 @@ public class WatershedOperation implements Operation {
           .build();
 
   private final SocketHint<Mat> srcHint = SocketHints.Inputs.createMatSocketHint("Input", false);
-  private final SocketHint<ContoursReport> contoursHint = new SocketHint.Builder<>(ContoursReport
-      .class)
-      .identifier("Contours")
-      .initialValueSupplier(ContoursReport::new)
-      .build();
+  private final SocketHint<ContoursReport> contoursHint =
+      new SocketHint.Builder<>(ContoursReport.class)
+          .identifier("Contours")
+          .initialValueSupplier(ContoursReport::new)
+          .build();
 
-  private final SocketHint<Mat> outputHint = SocketHints.Inputs.createMatSocketHint("Output", true);
+  private final SocketHint<ContoursReport> outputHint =
+      new SocketHint.Builder<>(ContoursReport.class)
+          .identifier("Features")
+          .initialValueSupplier(ContoursReport::new)
+          .build();
 
   private final InputSocket<Mat> srcSocket;
   private final InputSocket<ContoursReport> contoursSocket;
-  private final OutputSocket<Mat> outputSocket;
+  private final OutputSocket<ContoursReport> outputSocket;
 
   @SuppressWarnings("JavadocMethod")
-  public WatershedOperation(InputSocket.Factory inputSocketFactory, OutputSocket.Factory
-      outputSocketFactory) {
+  public WatershedOperation(InputSocket.Factory inputSocketFactory,
+                            OutputSocket.Factory outputSocketFactory) {
     srcSocket = inputSocketFactory.create(srcHint);
     contoursSocket = inputSocketFactory.create(contoursHint);
     outputSocket = outputSocketFactory.create(outputHint);
@@ -85,29 +97,85 @@ public void perform() {
     final ContoursReport contourReport = contoursSocket.getValue().get();
     final MatVector contours = contourReport.getContours();
 
+    final int maxMarkers = 253;
+    if (contours.size() > maxMarkers) {
+      throw new IllegalArgumentException(
+          "A maximum of " + maxMarkers + " contours can be used as markers."
+              + " Filter contours before connecting them to this operation if this keeps happening."
+              + " The contours must also all be external; nested contours will not work");
+    }
+
     final Mat markers = new Mat(input.size(), CV_32SC1, new Scalar(0.0));
     final Mat output = new Mat(markers.size(), CV_8UC1, new Scalar(0.0));
 
     try {
       // draw foreground markers (these have to be different colors)
       for (int i = 0; i < contours.size(); i++) {
-        drawContours(markers, contours, i, Scalar.all((i + 1) * (255 / contours.size())),
-            CV_FILLED, LINE_8, null, 2, null);
+        drawContours(markers, contours, i, Scalar.all(i + 1), CV_FILLED, LINE_8, null, 2, null);
       }
 
       // draw background marker a different color from the foreground markers
-      // TODO maybe make this configurable? There may be something in the corner
-      circle(markers, new Point(5, 5), 3, Scalar.WHITE, -1, LINE_8, 0);
+      Point backgroundLabel = fromPoint2f(findBackgroundMarker(markers, contours));
+      circle(markers, backgroundLabel, 1, Scalar.WHITE, -1, LINE_8, 0);
 
+      // Perform watershed
       watershed(input, markers);
       markers.convertTo(output, CV_8UC1);
-      bitwise_not(output, output); // watershed inverts colors; invert them back
 
-      outputSocket.setValue(output);
+      List<Mat> contourList = new ArrayList<>();
+      for (int i = 1; i < contours.size(); i++) {
+        Mat dst = new Mat();
+        output.copyTo(dst, opencv_core.equals(markers, i).asMat());
+        MatVector contour = new MatVector(); // vector with a single element
+        findContours(dst, contour, CV_RETR_EXTERNAL, CV_CHAIN_APPROX_TC89_KCOS);
+        assert contour.size() == 1;
+        contourList.add(contour.get(0).clone());
+        contour.get(0).deallocate();
+        contour.deallocate();
+      }
+      MatVector foundContours = new MatVector(contourList.toArray(new Mat[contourList.size()]));
+      outputSocket.setValue(new ContoursReport(foundContours, output.rows(), output.cols()));
     } finally {
       // make sure that the working mat is freed to avoid a memory leak
       markers.release();
     }
   }
 
+  /**
+   * Finds the first available point to place a background marker for the watershed operation.
+   */
+  private static Point2f findBackgroundMarker(Mat markers, MatVector contours) {
+    final int cols = markers.cols();
+    final int rows = markers.rows();
+    final int minDist = 5;
+    Point2f backgroundLabel = new Point2f();
+    boolean found = false;
+    // Don't place use a marker anywhere within 5 pixels of the edge of the image,
+    // or within 5 pixels of a contour
+    for (int x = minDist; x < cols - minDist && !found; x++) {
+      for (int y = minDist; y < rows - minDist && !found; y++) {
+        backgroundLabel.x(x);
+        backgroundLabel.y(y);
+        boolean isOpen = true;
+        for (int c = 0; c < contours.size(); c++) {
+          isOpen = pointPolygonTest(contours.get(c), backgroundLabel, true) <= -minDist;
+          if (!isOpen) {
+            // We know (x,y) is in a contour, don't need to check if it's in any others
+            break;
+          }
+        }
+        found = isOpen;
+      }
+    }
+    if (!found) {
+      // Should only happen if the image is clogged with contours
+      throw new IllegalStateException("Could not find a point for the background label");
+    }
+    return backgroundLabel;
+  }
+
+  private static Point fromPoint2f(Point2f p) {
+    return new Point((int) p.x(), (int) p.y());
+  }
+
 }