diff --git a/core/src/main/java/edu/wpi/grip/core/operations/composite/WatershedOperation.java b/core/src/main/java/edu/wpi/grip/core/operations/composite/WatershedOperation.java index cdfd455609..13181916e2 100644 --- a/core/src/main/java/edu/wpi/grip/core/operations/composite/WatershedOperation.java +++ b/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 srcHint = SocketHints.Inputs.createMatSocketHint("Input", false); - private final SocketHint contoursHint = new SocketHint.Builder<>(ContoursReport - .class) - .identifier("Contours") - .initialValueSupplier(ContoursReport::new) - .build(); + private final SocketHint contoursHint = + new SocketHint.Builder<>(ContoursReport.class) + .identifier("Contours") + .initialValueSupplier(ContoursReport::new) + .build(); - private final SocketHint outputHint = SocketHints.Inputs.createMatSocketHint("Output", true); + private final SocketHint outputHint = + new SocketHint.Builder<>(ContoursReport.class) + .identifier("Features") + .initialValueSupplier(ContoursReport::new) + .build(); private final InputSocket srcSocket; private final InputSocket contoursSocket; - private final OutputSocket outputSocket; + private final OutputSocket 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 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()); + } + }