001/*
002 * Copyright (C) Photon Vision.
003 *
004 * This program is free software: you can redistribute it and/or modify
005 * it under the terms of the GNU General Public License as published by
006 * the Free Software Foundation, either version 3 of the License, or
007 * (at your option) any later version.
008 *
009 * This program is distributed in the hope that it will be useful,
010 * but WITHOUT ANY WARRANTY; without even the implied warranty of
011 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
012 * GNU General Public License for more details.
013 *
014 * You should have received a copy of the GNU General Public License
015 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
016 */
017
018package org.photonvision.vision.pipe.impl;
019
020import java.util.List;
021import org.opencv.core.Mat;
022import org.opencv.core.MatOfPoint2f;
023import org.opencv.core.Point;
024import org.opencv.core.Scalar;
025import org.opencv.core.Size;
026import org.opencv.core.TermCriteria;
027import org.opencv.imgproc.Imgproc;
028import org.opencv.objdetect.Objdetect;
029import org.photonvision.vision.aruco.ArucoDetectionResult;
030import org.photonvision.vision.aruco.PhotonArucoDetector;
031import org.photonvision.vision.opencv.CVMat;
032import org.photonvision.vision.opencv.Releasable;
033import org.photonvision.vision.pipe.CVPipe;
034
035public class ArucoDetectionPipe
036        extends CVPipe<CVMat, List<ArucoDetectionResult>, ArucoDetectionPipeParams>
037        implements Releasable {
038    // ArucoDetector wrapper class
039    private final PhotonArucoDetector photonDetector = new PhotonArucoDetector();
040
041    // Ratio multiplied with image size and added to refinement window size
042    private static final double kRefineWindowImageRatio = 0.004;
043    // Ratio multiplied with max marker diagonal length and added to refinement window size
044    private static final double kRefineWindowMarkerRatio = 0.03;
045
046    @Override
047    protected List<ArucoDetectionResult> process(CVMat in) {
048        var imgMat = in.getMat();
049
050        // Sanity check -- image should not be empty
051        if (imgMat.empty()) {
052            // give up is best we can do here
053            return List.of();
054        }
055
056        var detections = photonDetector.detect(imgMat);
057        // manually do corner refinement ourselves
058        if (params.useCornerRefinement) {
059            for (var detection : detections) {
060                double[] xCorners = detection.getXCorners();
061                double[] yCorners = detection.getYCorners();
062                Point[] cornerPoints =
063                        new Point[] {
064                            new Point(xCorners[0], yCorners[0]),
065                            new Point(xCorners[1], yCorners[1]),
066                            new Point(xCorners[2], yCorners[2]),
067                            new Point(xCorners[3], yCorners[3])
068                        };
069                double bltr =
070                        Math.hypot(
071                                cornerPoints[2].x - cornerPoints[0].x, cornerPoints[2].y - cornerPoints[0].y);
072                double brtl =
073                        Math.hypot(
074                                cornerPoints[3].x - cornerPoints[1].x, cornerPoints[3].y - cornerPoints[1].y);
075                double minDiag = Math.min(bltr, brtl);
076                int halfWindowLength =
077                        (int) Math.ceil(kRefineWindowImageRatio * Math.min(imgMat.rows(), imgMat.cols()));
078                halfWindowLength += (int) (minDiag * kRefineWindowMarkerRatio);
079                // dont do refinement on small markers
080                if (halfWindowLength < 4) continue;
081                var halfWindowSize = new Size(halfWindowLength, halfWindowLength);
082                var ptsMat = new MatOfPoint2f(cornerPoints);
083                var criteria =
084                        new TermCriteria(3, params.refinementMaxIterations, params.refinementMinErrorPx);
085                Imgproc.cornerSubPix(imgMat, ptsMat, halfWindowSize, new Size(-1, -1), criteria);
086                cornerPoints = ptsMat.toArray();
087                for (int i = 0; i < cornerPoints.length; i++) {
088                    var pt = cornerPoints[i];
089                    xCorners[i] = pt.x;
090                    yCorners[i] = pt.y;
091                    // If we want to debug the refinement window, draw a rectangle on the image
092                    if (params.debugRefineWindow) {
093                        drawCornerRefineWindow(imgMat, pt, halfWindowLength);
094                    }
095                }
096            }
097        }
098        return List.of(detections);
099    }
100
101    @Override
102    public void setParams(ArucoDetectionPipeParams newParams) {
103        if (this.params == null || !this.params.equals(newParams)) {
104            System.out.println("Changing tag family to " + newParams.tagFamily);
105            photonDetector
106                    .getDetector()
107                    .setDictionary(Objdetect.getPredefinedDictionary(newParams.tagFamily));
108            var detectParams = photonDetector.getParams();
109
110            detectParams.set_adaptiveThreshWinSizeMin(newParams.threshMinSize);
111            detectParams.set_adaptiveThreshWinSizeStep(newParams.threshStepSize);
112            detectParams.set_adaptiveThreshWinSizeMax(newParams.threshMaxSize);
113            detectParams.set_adaptiveThreshConstant(newParams.threshConstant);
114
115            detectParams.set_errorCorrectionRate(newParams.errorCorrectionRate);
116
117            detectParams.set_useAruco3Detection(newParams.useAruco3);
118            detectParams.set_minSideLengthCanonicalImg(newParams.aruco3MinCanonicalImgSide);
119            detectParams.set_minMarkerLengthRatioOriginalImg((float) newParams.aruco3MinMarkerSideRatio);
120
121            photonDetector.setParams(detectParams);
122        }
123
124        super.setParams(newParams);
125    }
126
127    public PhotonArucoDetector getPhotonDetector() {
128        return photonDetector;
129    }
130
131    private void drawCornerRefineWindow(Mat outputMat, Point corner, int windowSize) {
132        int thickness = (int) (Math.ceil(Math.max(outputMat.cols(), outputMat.rows()) * 0.003));
133        var pt1 = new Point(corner.x - windowSize, corner.y - windowSize);
134        var pt2 = new Point(corner.x + windowSize, corner.y + windowSize);
135        Imgproc.rectangle(outputMat, pt1, pt2, new Scalar(0, 0, 255), thickness);
136    }
137
138    @Override
139    public void release() {
140        photonDetector.release();
141    }
142}