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 edu.wpi.first.math.util.Units;
021import java.nio.file.Path;
022import java.util.ArrayList;
023import java.util.List;
024import java.util.stream.Collectors;
025import org.apache.commons.lang3.tuple.Pair;
026import org.opencv.core.Mat;
027import org.opencv.core.Point;
028import org.photonvision.common.dataflow.DataChangeService;
029import org.photonvision.common.dataflow.events.OutgoingUIEvent;
030import org.photonvision.common.logging.LogGroup;
031import org.photonvision.common.logging.Logger;
032import org.photonvision.common.util.SerializationUtils;
033import org.photonvision.vision.calibration.BoardObservation;
034import org.photonvision.vision.calibration.CameraCalibrationCoefficients;
035import org.photonvision.vision.frame.Frame;
036import org.photonvision.vision.frame.FrameThresholdType;
037import org.photonvision.vision.opencv.CVMat;
038import org.photonvision.vision.opencv.ImageRotationMode;
039import org.photonvision.vision.pipe.CVPipe.CVPipeResult;
040import org.photonvision.vision.pipe.impl.CalculateFPSPipe;
041import org.photonvision.vision.pipe.impl.Calibrate3dPipe;
042import org.photonvision.vision.pipe.impl.Calibrate3dPipe.CalibrationInput;
043import org.photonvision.vision.pipe.impl.FindBoardCornersPipe;
044import org.photonvision.vision.pipe.impl.FindBoardCornersPipe.FindBoardCornersPipeResult;
045import org.photonvision.vision.pipeline.result.CVPipelineResult;
046import org.photonvision.vision.pipeline.result.CalibrationPipelineResult;
047
048public class Calibrate3dPipeline
049        extends CVPipeline<CVPipelineResult, Calibration3dPipelineSettings> {
050    // For logging
051    private static final Logger logger = new Logger(Calibrate3dPipeline.class, LogGroup.General);
052
053    // Find board corners decides internally between opencv and mrgingham
054    private final FindBoardCornersPipe findBoardCornersPipe = new FindBoardCornersPipe();
055    private final Calibrate3dPipe calibrate3dPipe = new Calibrate3dPipe();
056    private final CalculateFPSPipe calculateFPSPipe = new CalculateFPSPipe();
057
058    // Getter methods have been set for calibrate and takeSnapshot
059    private boolean takeSnapshot = false;
060
061    // Output of the corners
062    public final List<FindBoardCornersPipeResult> foundCornersList;
063
064    /// Output of the calibration, getter method is set for this.
065    private CVPipeResult<CameraCalibrationCoefficients> calibrationOutput;
066
067    private final int minSnapshots;
068
069    private boolean calibrating = false;
070
071    private static final FrameThresholdType PROCESSING_TYPE = FrameThresholdType.NONE;
072
073    public Calibrate3dPipeline() {
074        this(12);
075    }
076
077    public Calibrate3dPipeline(int minSnapshots) {
078        super(PROCESSING_TYPE);
079        this.settings = new Calibration3dPipelineSettings();
080        this.foundCornersList = new ArrayList<>();
081        this.minSnapshots = minSnapshots;
082    }
083
084    @Override
085    protected void setPipeParamsImpl() {
086        FindBoardCornersPipe.FindCornersPipeParams findCornersPipeParams =
087                new FindBoardCornersPipe.FindCornersPipeParams(
088                        settings.boardHeight,
089                        settings.boardWidth,
090                        settings.boardType,
091                        settings.tagFamily,
092                        settings.gridSize,
093                        settings.markerSize,
094                        settings.streamingFrameDivisor,
095                        settings.useOldPattern);
096        findBoardCornersPipe.setParams(findCornersPipeParams);
097
098        Calibrate3dPipe.CalibratePipeParams calibratePipeParams =
099                new Calibrate3dPipe.CalibratePipeParams(
100                        settings.boardHeight, settings.boardWidth, settings.gridSize, settings.useMrCal);
101        calibrate3dPipe.setParams(calibratePipeParams);
102    }
103
104    @Override
105    protected CVPipelineResult process(Frame frame, Calibration3dPipelineSettings settings) {
106        Mat inputColorMat = frame.colorImage.getMat();
107
108        if (this.calibrating || inputColorMat.empty()) {
109            return new CVPipelineResult(frame.sequenceID, 0, 0, null, frame);
110        }
111
112        if (getSettings().inputImageRotationMode != ImageRotationMode.DEG_0) {
113            // All this calibration assumes zero rotation. If we want a rotation, it should
114            // be applied at
115            // the output
116            logger.error(
117                    "Input image rotation was non-zero! Calibration wasn't designed to deal with this. Attempting to manually change back to zero");
118            getSettings().inputImageRotationMode = ImageRotationMode.DEG_0;
119            return new CVPipelineResult(frame.sequenceID, 0, 0, List.of(), frame);
120        }
121
122        long sumPipeNanosElapsed = 0L;
123
124        // Check if the frame has chessboard corners
125        var outputColorCVMat = new CVMat();
126        inputColorMat.copyTo(outputColorCVMat.getMat());
127
128        FindBoardCornersPipeResult findBoardResult;
129
130        findBoardResult =
131                findBoardCornersPipe.run(Pair.of(inputColorMat, outputColorCVMat.getMat())).output;
132
133        if (takeSnapshot) {
134            // Set snapshot to false even if we don't find a board
135            takeSnapshot = false;
136
137            if (findBoardResult != null) {
138                // Only copy the image into the result when we absolutely must
139                findBoardResult.inputImage = inputColorMat.clone();
140
141                foundCornersList.add(findBoardResult);
142
143                // update the UI
144                broadcastState();
145            }
146        }
147
148        var fpsResult = calculateFPSPipe.run(null);
149        var fps = fpsResult.output;
150
151        frame.release();
152
153        // Return the drawn chessboard if corners are found, if not, then return the
154        // input image.
155        return new CalibrationPipelineResult(
156                frame.sequenceID,
157                sumPipeNanosElapsed,
158                fps, // Unused but here in case
159                new Frame(
160                        frame.sequenceID,
161                        new CVMat(),
162                        outputColorCVMat,
163                        FrameThresholdType.NONE,
164                        frame.frameStaticProperties),
165                getCornersList());
166    }
167
168    List<List<Point>> getCornersList() {
169        return foundCornersList.stream()
170                .map(it -> it.imagePoints.toList())
171                .collect(Collectors.toList());
172    }
173
174    public boolean hasEnough() {
175        return foundCornersList.size() >= minSnapshots;
176    }
177
178    public CameraCalibrationCoefficients tryCalibration(Path imageSavePath) {
179        if (!hasEnough()) {
180            logger.info(
181                    "Not enough snapshots! Only got "
182                            + foundCornersList.size()
183                            + " of "
184                            + minSnapshots
185                            + " -- returning null..");
186            return null;
187        }
188
189        this.calibrating = true;
190
191        /*
192         * Pass the board corners to the pipe, which will check again to see if all
193         * boards are valid
194         * and returns the corresponding image and object points
195         */
196        calibrationOutput =
197                calibrate3dPipe.run(
198                        new CalibrationInput(foundCornersList, frameStaticProperties, imageSavePath));
199
200        this.calibrating = false;
201
202        return calibrationOutput.output;
203    }
204
205    public void takeSnapshot() {
206        takeSnapshot = true;
207    }
208
209    public List<BoardObservation> perViewErrors() {
210        return calibrationOutput.output.observations;
211    }
212
213    public void finishCalibration() {
214        foundCornersList.forEach(it -> it.release());
215        foundCornersList.clear();
216
217        broadcastState();
218    }
219
220    private void broadcastState() {
221        var state =
222                SerializationUtils.objectToHashMap(
223                        new UICalibrationData(
224                                foundCornersList.size(),
225                                settings.cameraVideoModeIndex,
226                                minSnapshots,
227                                hasEnough(),
228                                Units.metersToInches(settings.gridSize),
229                                Units.metersToInches(settings.markerSize),
230                                settings.boardWidth,
231                                settings.boardHeight,
232                                settings.boardType,
233                                settings.useOldPattern,
234                                settings.tagFamily));
235
236        DataChangeService.getInstance()
237                .publishEvent(OutgoingUIEvent.wrappedOf("calibrationData", state));
238    }
239
240    public boolean removeSnapshot(int index) {
241        try {
242            foundCornersList.remove(index);
243            return true;
244        } catch (ArrayIndexOutOfBoundsException e) {
245            logger.error("Could not remove snapshot at index " + index, e);
246            return false;
247        }
248    }
249
250    public CameraCalibrationCoefficients cameraCalibrationCoefficients() {
251        return calibrationOutput.output;
252    }
253
254    @Override
255    public void release() {
256        // we never actually need to give resources up since pipelinemanager only makes
257        // one of us
258    }
259}