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.pipeline;
019
020import java.util.Arrays;
021import java.util.List;
022import java.util.stream.Collectors;
023import org.apache.commons.lang3.tuple.Pair;
024import org.opencv.core.Point;
025import org.photonvision.vision.frame.Frame;
026import org.photonvision.vision.frame.FrameThresholdType;
027import org.photonvision.vision.opencv.CVShape;
028import org.photonvision.vision.opencv.Contour;
029import org.photonvision.vision.opencv.ContourShape;
030import org.photonvision.vision.opencv.DualOffsetValues;
031import org.photonvision.vision.pipe.CVPipe.CVPipeResult;
032import org.photonvision.vision.pipe.impl.*;
033import org.photonvision.vision.pipeline.result.CVPipelineResult;
034import org.photonvision.vision.target.PotentialTarget;
035import org.photonvision.vision.target.TrackedTarget;
036
037public class ColoredShapePipeline
038        extends CVPipeline<CVPipelineResult, ColoredShapePipelineSettings> {
039    private final SpeckleRejectPipe speckleRejectPipe = new SpeckleRejectPipe();
040    private final FindContoursPipe findContoursPipe = new FindContoursPipe();
041    private final FindPolygonPipe findPolygonPipe = new FindPolygonPipe();
042    private final FindCirclesPipe findCirclesPipe = new FindCirclesPipe();
043    private final FilterShapesPipe filterShapesPipe = new FilterShapesPipe();
044    private final SortContoursPipe sortContoursPipe = new SortContoursPipe();
045    private final Collect2dTargetsPipe collect2dTargetsPipe = new Collect2dTargetsPipe();
046    private final CornerDetectionPipe cornerDetectionPipe = new CornerDetectionPipe();
047    private final SolvePNPPipe solvePNPPipe = new SolvePNPPipe();
048    private final Draw2dCrosshairPipe draw2dCrosshairPipe = new Draw2dCrosshairPipe();
049    private final Draw2dTargetsPipe draw2DTargetsPipe = new Draw2dTargetsPipe();
050    private final Draw3dTargetsPipe draw3dTargetsPipe = new Draw3dTargetsPipe();
051    private final CalculateFPSPipe calculateFPSPipe = new CalculateFPSPipe();
052
053    private final Point[] rectPoints = new Point[4];
054
055    private static final FrameThresholdType PROCESSING_TYPE = FrameThresholdType.HSV;
056
057    public ColoredShapePipeline() {
058        super(PROCESSING_TYPE);
059        settings = new ColoredShapePipelineSettings();
060    }
061
062    public ColoredShapePipeline(ColoredShapePipelineSettings settings) {
063        super(PROCESSING_TYPE);
064        this.settings = settings;
065    }
066
067    @Override
068    protected void setPipeParamsImpl() {
069        DualOffsetValues dualOffsetValues =
070                new DualOffsetValues(
071                        settings.offsetDualPointA,
072                        settings.offsetDualPointAArea,
073                        settings.offsetDualPointB,
074                        settings.offsetDualPointBArea);
075
076        SpeckleRejectPipe.SpeckleRejectParams speckleRejectParams =
077                new SpeckleRejectPipe.SpeckleRejectParams(settings.contourSpecklePercentage);
078        speckleRejectPipe.setParams(speckleRejectParams);
079
080        FindContoursPipe.FindContoursParams findContoursParams =
081                new FindContoursPipe.FindContoursParams();
082        findContoursPipe.setParams(findContoursParams);
083
084        FindPolygonPipe.FindPolygonPipeParams findPolygonPipeParams =
085                new FindPolygonPipe.FindPolygonPipeParams(settings.accuracyPercentage);
086        findPolygonPipe.setParams(findPolygonPipeParams);
087
088        FindCirclesPipe.FindCirclePipeParams findCirclePipeParams =
089                new FindCirclesPipe.FindCirclePipeParams(
090                        settings.circleDetectThreshold,
091                        settings.contourRadius.getFirst(),
092                        settings.minDist,
093                        settings.contourRadius.getSecond(),
094                        settings.maxCannyThresh,
095                        settings.circleAccuracy,
096                        Math.hypot(frameStaticProperties.imageWidth, frameStaticProperties.imageHeight));
097        findCirclesPipe.setParams(findCirclePipeParams);
098
099        FilterShapesPipe.FilterShapesPipeParams filterShapesPipeParams =
100                new FilterShapesPipe.FilterShapesPipeParams(
101                        settings.contourShape,
102                        settings.contourArea.getFirst(),
103                        settings.contourArea.getSecond(),
104                        settings.contourPerimeter.getFirst(),
105                        settings.contourPerimeter.getSecond(),
106                        frameStaticProperties);
107        filterShapesPipe.setParams(filterShapesPipeParams);
108
109        SortContoursPipe.SortContoursParams sortContoursParams =
110                new SortContoursPipe.SortContoursParams(
111                        settings.contourSortMode,
112                        settings.outputShowMultipleTargets ? MAX_MULTI_TARGET_RESULTS : 1,
113                        frameStaticProperties); // TODO don't hardcode?
114        sortContoursPipe.setParams(sortContoursParams);
115
116        Collect2dTargetsPipe.Collect2dTargetsParams collect2dTargetsParams =
117                new Collect2dTargetsPipe.Collect2dTargetsParams(
118                        settings.offsetRobotOffsetMode,
119                        settings.offsetSinglePoint,
120                        dualOffsetValues,
121                        settings.contourTargetOffsetPointEdge,
122                        settings.contourTargetOrientation,
123                        frameStaticProperties);
124        collect2dTargetsPipe.setParams(collect2dTargetsParams);
125
126        var params =
127                new CornerDetectionPipe.CornerDetectionPipeParameters(
128                        settings.cornerDetectionStrategy,
129                        settings.cornerDetectionUseConvexHulls,
130                        settings.cornerDetectionExactSideCount,
131                        settings.cornerDetectionSideCount,
132                        settings.cornerDetectionAccuracyPercentage);
133        cornerDetectionPipe.setParams(params);
134
135        var solvePNPParams =
136                new SolvePNPPipe.SolvePNPPipeParams(
137                        frameStaticProperties.cameraCalibration, settings.targetModel);
138        solvePNPPipe.setParams(solvePNPParams);
139
140        Draw2dTargetsPipe.Draw2dTargetsParams draw2DTargetsParams =
141                new Draw2dTargetsPipe.Draw2dTargetsParams(
142                        settings.outputShouldDraw,
143                        settings.outputShowMultipleTargets,
144                        settings.streamingFrameDivisor);
145        draw2DTargetsParams.showShape = true;
146        draw2DTargetsParams.showMaximumBox = false;
147        draw2DTargetsParams.showRotatedBox = false;
148        draw2DTargetsPipe.setParams(draw2DTargetsParams);
149
150        Draw2dCrosshairPipe.Draw2dCrosshairParams draw2dCrosshairParams =
151                new Draw2dCrosshairPipe.Draw2dCrosshairParams(
152                        settings.outputShouldDraw,
153                        settings.offsetRobotOffsetMode,
154                        settings.offsetSinglePoint,
155                        dualOffsetValues,
156                        frameStaticProperties,
157                        settings.streamingFrameDivisor,
158                        settings.inputImageRotationMode);
159        draw2dCrosshairPipe.setParams(draw2dCrosshairParams);
160
161        var draw3dTargetsParams =
162                new Draw3dTargetsPipe.Draw3dContoursParams(
163                        settings.outputShouldDraw,
164                        frameStaticProperties.cameraCalibration,
165                        settings.targetModel,
166                        settings.streamingFrameDivisor);
167        draw3dTargetsPipe.setParams(draw3dTargetsParams);
168    }
169
170    @Override
171    protected CVPipelineResult process(Frame frame, ColoredShapePipelineSettings settings) {
172        long sumPipeNanosElapsed = 0L;
173
174        CVPipeResult<List<Contour>> findContoursResult =
175                findContoursPipe.run(frame.processedImage.getMat());
176        sumPipeNanosElapsed += findContoursResult.nanosElapsed;
177
178        CVPipeResult<List<Contour>> speckleRejectResult =
179                speckleRejectPipe.run(findContoursResult.output);
180        sumPipeNanosElapsed += speckleRejectResult.nanosElapsed;
181
182        List<CVShape> shapes = null;
183        if (settings.contourShape == ContourShape.Circle) {
184            CVPipeResult<List<CVShape>> findCirclesResult =
185                    findCirclesPipe.run(Pair.of(frame.processedImage.getMat(), speckleRejectResult.output));
186            sumPipeNanosElapsed += findCirclesResult.nanosElapsed;
187            shapes = findCirclesResult.output;
188        } else {
189            CVPipeResult<List<CVShape>> findPolygonsResult =
190                    findPolygonPipe.run(speckleRejectResult.output);
191            sumPipeNanosElapsed += findPolygonsResult.nanosElapsed;
192            shapes = findPolygonsResult.output;
193        }
194
195        CVPipeResult<List<CVShape>> filterShapeResult = filterShapesPipe.run(shapes);
196        sumPipeNanosElapsed += filterShapeResult.nanosElapsed;
197
198        CVPipeResult<List<PotentialTarget>> sortContoursResult =
199                sortContoursPipe.run(
200                        filterShapeResult.output.stream()
201                                .map(shape -> new PotentialTarget(shape.getContour(), shape))
202                                .collect(Collectors.toList()));
203        sumPipeNanosElapsed += sortContoursResult.nanosElapsed;
204
205        CVPipeResult<List<TrackedTarget>> collect2dTargetsResult =
206                collect2dTargetsPipe.run(sortContoursResult.output);
207        sumPipeNanosElapsed += collect2dTargetsResult.nanosElapsed;
208
209        List<TrackedTarget> targetList;
210
211        if (settings.solvePNPEnabled && settings.contourShape == ContourShape.Circle) {
212            var cornerDetectionResult = cornerDetectionPipe.run(collect2dTargetsResult.output);
213            collect2dTargetsResult.output.forEach(
214                    shape -> {
215                        shape.getMinAreaRect().points(rectPoints);
216                        shape.setTargetCorners(Arrays.asList(rectPoints));
217                    });
218            sumPipeNanosElapsed += cornerDetectionResult.nanosElapsed;
219
220            var solvePNPResult = solvePNPPipe.run(cornerDetectionResult.output);
221            sumPipeNanosElapsed += solvePNPResult.nanosElapsed;
222
223            targetList = solvePNPResult.output;
224        } else {
225            targetList = collect2dTargetsResult.output;
226        }
227
228        var fpsResult = calculateFPSPipe.run(null);
229        var fps = fpsResult.output;
230
231        return new CVPipelineResult(frame.sequenceID, sumPipeNanosElapsed, fps, targetList, frame);
232    }
233}