Source code

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.awt.*;
021import java.util.List;
022import org.apache.commons.lang3.tuple.Pair;
023import org.opencv.core.*;
024import org.opencv.core.Point;
025import org.opencv.imgproc.Imgproc;
026import org.photonvision.common.logging.LogGroup;
027import org.photonvision.common.logging.Logger;
028import org.photonvision.common.util.ColorHelper;
029import org.photonvision.vision.frame.FrameDivisor;
030import org.photonvision.vision.opencv.CVShape;
031import org.photonvision.vision.opencv.ContourShape;
032import org.photonvision.vision.pipe.MutatingPipe;
033import org.photonvision.vision.target.TrackedTarget;
034
035public class Draw2dTargetsPipe
036        extends MutatingPipe<Pair<Mat, List<TrackedTarget>>, Draw2dTargetsPipe.Draw2dTargetsParams> {
037    MatOfPoint tempMat = new MatOfPoint();
038    private static final Logger logger = new Logger(Draw2dTargetsPipe.class, LogGroup.General);
039
040    @Override
041    protected Void process(Pair<Mat, List<TrackedTarget>> in) {
042        var imRows = in.getLeft().rows();
043        var imCols = in.getLeft().cols();
044        var imageSize = Math.sqrt(imRows * imCols);
045        var textSize = params.kPixelsToText * imageSize;
046        var thickness = params.kPixelsToThickness * imageSize;
047
048        if (!params.shouldDraw) return null;
049
050        if (!in.getRight().isEmpty()
051                && (params.showCentroid
052                        || params.showMaximumBox
053                        || params.showRotatedBox
054                        || params.showShape)) {
055            var centroidColour = ColorHelper.colorToScalar(params.centroidColor);
056            var maximumBoxColour = ColorHelper.colorToScalar(params.maximumBoxColor);
057            var rotatedBoxColour = ColorHelper.colorToScalar(params.rotatedBoxColor);
058            var circleColor = ColorHelper.colorToScalar(params.circleColor);
059            var shapeColour = ColorHelper.colorToScalar(params.shapeOutlineColour);
060
061            for (int i = 0; i < (params.showMultipleTargets ? in.getRight().size() : 1); i++) {
062                Point[] vertices = new Point[4];
063                MatOfPoint contour = new MatOfPoint();
064
065                if (i != 0 && !params.showMultipleTargets) {
066                    break;
067                }
068
069                TrackedTarget target = in.getRight().get(i);
070                RotatedRect r = target.getMinAreaRect();
071
072                if (r == null) continue;
073
074                r.points(vertices);
075                dividePointArray(vertices);
076                contour.fromArray(vertices);
077
078                if (params.shouldShowRotatedBox(target.getShape())) {
079                    Imgproc.drawContours(
080                            in.getLeft(),
081                            List.of(contour),
082                            0,
083                            rotatedBoxColour,
084                            (int) Math.ceil(imageSize * params.kPixelsToBoxThickness));
085                } else if (params.shouldShowCircle(target.getShape())) {
086                    Imgproc.circle(
087                            in.getLeft(),
088                            target.getShape().center,
089                            (int) target.getShape().radius,
090                            circleColor,
091                            (int) Math.ceil(imageSize * params.kPixelsToBoxThickness));
092                } else {
093                    // draw approximate polygon
094                    var poly = target.getApproximateBoundingPolygon();
095
096                    // fall back on the shape's approx poly dp
097                    if (poly == null && target.getShape() != null)
098                        poly = target.getShape().getContour().getApproxPolyDp();
099                    if (poly != null) {
100                        var mat = new MatOfPoint();
101                        mat.fromArray(poly.toArray());
102                        divideMat(mat, mat);
103                        Imgproc.drawContours(
104                                in.getLeft(),
105                                List.of(mat),
106                                -1,
107                                ColorHelper.colorToScalar(params.rotatedBoxColor),
108                                2);
109                        mat.release();
110                    }
111                }
112
113                if (params.showMaximumBox) {
114                    Rect box = Imgproc.boundingRect(contour);
115                    Imgproc.rectangle(
116                            in.getLeft(),
117                            new Point(box.x, box.y),
118                            new Point(box.x + box.width, box.y + box.height),
119                            maximumBoxColour,
120                            (int) Math.ceil(imageSize * params.kPixelsToBoxThickness));
121                }
122
123                if (params.showShape) {
124                    divideMat(target.m_mainContour.mat, tempMat);
125                    Imgproc.drawContours(
126                            in.getLeft(),
127                            List.of(tempMat),
128                            -1,
129                            shapeColour,
130                            (int) Math.ceil(imageSize * params.kPixelsToBoxThickness));
131                }
132
133                if (params.showContourNumber) {
134                    var center = target.m_mainContour.getCenterPoint();
135                    var textPos =
136                            new Point(
137                                    center.x + params.kPixelsToOffset * imageSize,
138                                    center.y - params.kPixelsToOffset * imageSize);
139                    dividePoint(textPos);
140
141                    int id = target.getFiducialId();
142                    var contourNumber = String.valueOf(id == -1 ? i : id);
143
144                    Imgproc.putText(
145                            in.getLeft(),
146                            contourNumber,
147                            textPos,
148                            0,
149                            textSize,
150                            ColorHelper.colorToScalar(params.textColor),
151                            (int) thickness);
152                }
153
154                if (params.showCentroid) {
155                    Point centroid = target.getTargetOffsetPoint().clone();
156                    dividePoint(centroid);
157                    var crosshairRadius = (int) (imageSize * params.kPixelsToCentroidRadius);
158                    var x = centroid.x;
159                    var y = centroid.y;
160                    Point xMax = new Point(x + crosshairRadius, y);
161                    Point xMin = new Point(x - crosshairRadius, y);
162                    Point yMax = new Point(x, y + crosshairRadius);
163                    Point yMin = new Point(x, y - crosshairRadius);
164
165                    Imgproc.line(
166                            in.getLeft(),
167                            xMax,
168                            xMin,
169                            centroidColour,
170                            (int) Math.ceil(imageSize * params.kPixelsToBoxThickness));
171                    Imgproc.line(
172                            in.getLeft(),
173                            yMax,
174                            yMin,
175                            centroidColour,
176                            (int) Math.ceil(imageSize * params.kPixelsToBoxThickness));
177                }
178            }
179        }
180
181        return null;
182    }
183
184    private void divideMat(MatOfPoint src, MatOfPoint dst) {
185        var hull = src.toArray();
186        for (Point point : hull) {
187            dividePoint(point);
188        }
189        dst.fromArray(hull);
190    }
191
192    private void divideMat(MatOfPoint2f src, MatOfPoint dst) {
193        var hull = src.toArray();
194        for (Point point : hull) {
195            dividePoint(point);
196        }
197        dst.fromArray(hull);
198    }
199
200    /** Scale a given point list by the current frame divisor. the point list is mutated! */
201    private void dividePointList(List<Point> points) {
202        for (var p : points) {
203            dividePoint(p);
204        }
205    }
206
207    /** Scale a given point array by the current frame divisor. the point list is mutated! */
208    private void dividePointArray(Point[] points) {
209        for (var p : points) {
210            dividePoint(p);
211        }
212    }
213
214    private void dividePoint(Point p) {
215        p.x = p.x / (double) params.divisor.value;
216        p.y = p.y / (double) params.divisor.value;
217    }
218
219    public static class Draw2dTargetsParams {
220        public double kPixelsToText = 0.0025;
221        public double kPixelsToThickness = 0.008;
222        public double kPixelsToOffset = 0.04;
223        public double kPixelsToBoxThickness = 0.007;
224        public double kPixelsToCentroidRadius = 0.03;
225        public boolean showCentroid = true;
226        public boolean showRotatedBox = true;
227        public boolean showShape = false;
228        public boolean showMaximumBox = true;
229        public boolean showContourNumber = true;
230        public Color centroidColor = Color.GREEN; // Color.decode("#ff5ebf");
231        public Color rotatedBoxColor = Color.BLUE;
232        public Color maximumBoxColor = Color.RED;
233        public Color shapeOutlineColour = Color.MAGENTA;
234        public Color textColor = Color.GREEN;
235        public Color circleColor = Color.RED;
236
237        public final boolean showMultipleTargets;
238        public final boolean shouldDraw;
239
240        public final FrameDivisor divisor;
241
242        public boolean shouldShowRotatedBox(CVShape shape) {
243            return showRotatedBox && (shape == null || shape.shape.equals(ContourShape.Quadrilateral));
244        }
245
246        public boolean shouldShowCircle(CVShape shape) {
247            return shape != null && shape.shape.equals(ContourShape.Circle);
248        }
249
250        public Draw2dTargetsParams(
251                boolean shouldDraw, boolean showMultipleTargets, FrameDivisor divisor) {
252            this.shouldDraw = shouldDraw;
253            this.showMultipleTargets = showMultipleTargets;
254            this.divisor = divisor;
255        }
256    }
257}