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.Color;
021import java.util.List;
022import org.apache.commons.lang3.tuple.Pair;
023import org.opencv.core.Mat;
024import org.opencv.core.Point;
025import org.opencv.core.Scalar;
026import org.opencv.imgproc.Imgproc;
027import org.photonvision.common.util.ColorHelper;
028import org.photonvision.vision.frame.FrameDivisor;
029import org.photonvision.vision.pipe.MutatingPipe;
030import org.photonvision.vision.target.TrackedTarget;
031
032public class DrawCalibrationPipe
033        extends MutatingPipe<
034                Pair<Mat, List<TrackedTarget>>, DrawCalibrationPipe.DrawCalibrationPipeParams> {
035    Scalar[] chessboardColors =
036            new Scalar[] {
037                ColorHelper.colorToScalar(Color.RED, 0.4),
038                ColorHelper.colorToScalar(Color.ORANGE, 0.4),
039                ColorHelper.colorToScalar(Color.GREEN, 0.4),
040                ColorHelper.colorToScalar(Color.BLUE, 0.4),
041                ColorHelper.colorToScalar(Color.MAGENTA, 0.4),
042            };
043
044    @Override
045    protected Void process(Pair<Mat, List<TrackedTarget>> in) {
046        if (!params.drawAllSnapshots) return null;
047
048        var image = in.getLeft();
049
050        var imgSz = image.size();
051        var diag = Math.hypot(imgSz.width, imgSz.height);
052
053        // heuristic: about 4px at a diagonal of 750px, or .5%, 'looks good'. keep it at least 3px at
054        // worst tho
055        int r = (int) Math.max(diag * 4.0 / 750.0, 3);
056        int thickness = (int) Math.max(diag * 1.0 / 600.0, 1);
057
058        int i = 0;
059        for (var target : in.getRight()) {
060            for (var c : target.getTargetCorners()) {
061                if (c.x < 0 || c.y < 0) {
062                    // Skip if the corner is less than zero
063                    continue;
064                }
065
066                c =
067                        new Point(
068                                c.x / params.divisor.value.doubleValue(), c.y / params.divisor.value.doubleValue());
069
070                var r2 = r / Math.sqrt(2);
071                var color = chessboardColors[i % chessboardColors.length];
072                Imgproc.circle(image, c, r, color, thickness);
073                Imgproc.line(
074                        image, new Point(c.x - r2, c.y - r2), new Point(c.x + r2, c.y + r2), color, thickness);
075                Imgproc.line(
076                        image, new Point(c.x + r2, c.y - r2), new Point(c.x - r2, c.y + r2), color, thickness);
077            }
078
079            i++;
080        }
081
082        return null;
083    }
084
085    public static class DrawCalibrationPipeParams {
086        private final FrameDivisor divisor;
087        public boolean drawAllSnapshots;
088
089        public DrawCalibrationPipeParams(FrameDivisor divisor, boolean drawSnapshots) {
090            this.divisor = divisor;
091            this.drawAllSnapshots = drawSnapshots;
092        }
093    }
094}