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.calib3d.Calib3d;
024import org.opencv.core.*;
025import org.opencv.core.Point;
026import org.opencv.imgproc.Imgproc;
027import org.photonvision.common.logging.LogGroup;
028import org.photonvision.common.logging.Logger;
029import org.photonvision.common.util.ColorHelper;
030import org.photonvision.estimation.OpenCVHelp;
031import org.photonvision.vision.calibration.CameraCalibrationCoefficients;
032import org.photonvision.vision.frame.FrameDivisor;
033import org.photonvision.vision.pipe.MutatingPipe;
034import org.photonvision.vision.target.TargetModel;
035import org.photonvision.vision.target.TrackedTarget;
036
037public class Draw3dTargetsPipe
038        extends MutatingPipe<Pair<Mat, List<TrackedTarget>>, Draw3dTargetsPipe.Draw3dContoursParams> {
039    Logger logger = new Logger(Draw3dTargetsPipe.class, LogGroup.VisionModule);
040
041    @Override
042    protected Void process(Pair<Mat, List<TrackedTarget>> in) {
043        if (!params.shouldDraw) return null;
044        if (params.cameraCalibrationCoefficients == null
045                || params.cameraCalibrationCoefficients.getCameraIntrinsicsMat() == null
046                || params.cameraCalibrationCoefficients.getDistCoeffsMat() == null) {
047            return null;
048        }
049
050        for (var target : in.getRight()) {
051            // draw convex hull
052            if (params.shouldDrawHull(target)) {
053                var pointMat = new MatOfPoint();
054                divideMat2f(target.m_mainContour.getConvexHull(), pointMat);
055                if (pointMat.size().empty()) {
056                    logger.error("Convex hull is empty?");
057                    logger.debug(
058                            "Orig. Convex Hull: " + target.m_mainContour.getConvexHull().size().toString());
059                    continue;
060                }
061                Imgproc.drawContours(
062                        in.getLeft(), List.of(pointMat), -1, ColorHelper.colorToScalar(Color.green), 1);
063
064                // draw approximate polygon
065                var poly = target.getApproximateBoundingPolygon();
066                if (poly != null) {
067                    divideMat2f(poly, pointMat);
068                    Imgproc.drawContours(
069                            in.getLeft(), List.of(pointMat), -1, ColorHelper.colorToScalar(Color.blue), 2);
070                }
071                pointMat.release();
072            }
073
074            // Draw floor and top
075            if (target.getCameraRelativeRvec() != null
076                    && target.getCameraRelativeTvec() != null
077                    && params.shouldDrawBox) {
078                var tempMat = new MatOfPoint2f();
079
080                var jac = new Mat();
081                var bottomModel = params.targetModel.getVisualizationBoxBottom();
082                var topModel = params.targetModel.getVisualizationBoxTop();
083
084                Calib3d.projectPoints(
085                        bottomModel,
086                        target.getCameraRelativeRvec(),
087                        target.getCameraRelativeTvec(),
088                        params.cameraCalibrationCoefficients.getCameraIntrinsicsMat(),
089                        params.cameraCalibrationCoefficients.getDistCoeffsMat(),
090                        tempMat,
091                        jac);
092
093                if (params.redistortPoints) {
094                    // Distort the points, so they match the image they're being overlaid on
095                    tempMat.fromList(
096                            OpenCVHelp.distortPoints(
097                                    tempMat.toList(),
098                                    params.cameraCalibrationCoefficients.getCameraIntrinsicsMat(),
099                                    params.cameraCalibrationCoefficients.getDistCoeffsMat()));
100                }
101
102                var bottomPoints = tempMat.toList();
103
104                Calib3d.projectPoints(
105                        topModel,
106                        target.getCameraRelativeRvec(),
107                        target.getCameraRelativeTvec(),
108                        params.cameraCalibrationCoefficients.getCameraIntrinsicsMat(),
109                        params.cameraCalibrationCoefficients.getDistCoeffsMat(),
110                        tempMat,
111                        jac);
112
113                if (params.redistortPoints) {
114                    // Distort the points, so they match the image they're being overlaid on
115                    tempMat.fromList(
116                            OpenCVHelp.distortPoints(
117                                    tempMat.toList(),
118                                    params.cameraCalibrationCoefficients.getCameraIntrinsicsMat(),
119                                    params.cameraCalibrationCoefficients.getDistCoeffsMat()));
120                }
121                var topPoints = tempMat.toList();
122
123                dividePointList(bottomPoints);
124                dividePointList(topPoints);
125
126                // floor, then pillars, then top
127                for (int i = 0; i < bottomPoints.size(); i++) {
128                    Imgproc.line(
129                            in.getLeft(),
130                            bottomPoints.get(i),
131                            bottomPoints.get((i + 1) % (bottomPoints.size())),
132                            ColorHelper.colorToScalar(Color.green),
133                            3);
134                }
135
136                // Draw X, Y and Z axis
137                MatOfPoint3f pointMat = new MatOfPoint3f();
138                // OpenCV expects coordinates in EDN, but we want to visualize in NWU
139                // NWU | EDN
140                // X: Z
141                // Y: -X
142                // Z: -Y
143                final double AXIS_LEN = 0.2;
144                var list =
145                        List.of(
146                                new Point3(0, 0, 0),
147                                new Point3(0, 0, AXIS_LEN), // x-axis
148                                new Point3(-AXIS_LEN, 0, 0), // y-axis
149                                new Point3(0, -AXIS_LEN, 0)); // z-axis
150                pointMat.fromList(list);
151
152                // The detected target's rvec and tvec perform a rotation-translation transformation which
153                // converts points in the target's coordinate system to the camera's. This means applying
154                // the transformation to the target point (0,0,0) for example would give the target's center
155                // relative to the camera.
156                Calib3d.projectPoints(
157                        pointMat,
158                        target.getCameraRelativeRvec(),
159                        target.getCameraRelativeTvec(),
160                        params.cameraCalibrationCoefficients.getCameraIntrinsicsMat(),
161                        params.cameraCalibrationCoefficients.getDistCoeffsMat(),
162                        tempMat,
163                        jac);
164                var axisPoints = tempMat.toList();
165                dividePointList(axisPoints);
166
167                // XYZ is RGB
168                // y-axis = green
169                Imgproc.line(
170                        in.getLeft(),
171                        axisPoints.get(0),
172                        axisPoints.get(2),
173                        ColorHelper.colorToScalar(Color.GREEN),
174                        3);
175                // z-axis = blue
176                Imgproc.line(
177                        in.getLeft(),
178                        axisPoints.get(0),
179                        axisPoints.get(3),
180                        ColorHelper.colorToScalar(Color.BLUE),
181                        3);
182                // x-axis = red
183                Imgproc.line(
184                        in.getLeft(),
185                        axisPoints.get(0),
186                        axisPoints.get(1),
187                        ColorHelper.colorToScalar(Color.RED),
188                        3);
189
190                // box edges perpendicular to tag
191                for (int i = 0; i < bottomPoints.size(); i++) {
192                    Imgproc.line(
193                            in.getLeft(),
194                            bottomPoints.get(i),
195                            topPoints.get(i),
196                            ColorHelper.colorToScalar(Color.blue),
197                            3);
198                }
199                // box edges parallel to tag
200                for (int i = 0; i < topPoints.size(); i++) {
201                    Imgproc.line(
202                            in.getLeft(),
203                            topPoints.get(i),
204                            topPoints.get((i + 1) % (bottomPoints.size())),
205                            ColorHelper.colorToScalar(Color.orange),
206                            3);
207                }
208
209                tempMat.release();
210                jac.release();
211                pointMat.release();
212            }
213
214            // draw corners
215            var corners = target.getTargetCorners();
216            if (corners != null && !corners.isEmpty()) {
217                for (var corner : corners) {
218                    var x = corner.x / (double) params.divisor.value;
219                    var y = corner.y / (double) params.divisor.value;
220
221                    Imgproc.circle(
222                            in.getLeft(),
223                            new Point(x, y),
224                            params.radius,
225                            ColorHelper.colorToScalar(params.color),
226                            params.radius);
227                }
228            }
229        }
230
231        return null;
232    }
233
234    private void divideMat2f(MatOfPoint2f src, MatOfPoint dst) {
235        var hull = src.toArray();
236        var pointArray = new Point[hull.length];
237        for (int i = 0; i < hull.length; i++) {
238            var hullAtI = hull[i];
239            pointArray[i] =
240                    new Point(
241                            hullAtI.x / (double) params.divisor.value, hullAtI.y / (double) params.divisor.value);
242        }
243        dst.fromArray(pointArray);
244    }
245
246    /** Scale a given point list by the current frame divisor. the point list is mutated! */
247    private void dividePointList(List<Point> points) {
248        for (var p : points) {
249            p.x = p.x / (double) params.divisor.value;
250            p.y = p.y / (double) params.divisor.value;
251        }
252    }
253
254    public static class Draw3dContoursParams {
255        public int radius = 2;
256        public Color color = Color.RED;
257
258        public final boolean shouldDraw;
259        public boolean shouldDrawHull = true;
260        public boolean shouldDrawBox = true;
261        public final TargetModel targetModel;
262        public final CameraCalibrationCoefficients cameraCalibrationCoefficients;
263        public final FrameDivisor divisor;
264
265        public boolean redistortPoints = false;
266
267        public Draw3dContoursParams(
268                boolean shouldDraw,
269                CameraCalibrationCoefficients cameraCalibrationCoefficients,
270                TargetModel targetModel,
271                FrameDivisor divisor) {
272            this.shouldDraw = shouldDraw;
273            this.cameraCalibrationCoefficients = cameraCalibrationCoefficients;
274            this.targetModel = targetModel;
275            this.divisor = divisor;
276        }
277
278        public boolean shouldDrawHull(TrackedTarget t) {
279            return !t.isFiducial() && this.shouldDrawHull;
280        }
281    }
282}