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.List;
021import java.util.Optional;
022import java.util.stream.Collectors;
023import org.photonvision.common.configuration.NeuralNetworkModelManager;
024import org.photonvision.vision.frame.Frame;
025import org.photonvision.vision.frame.FrameThresholdType;
026import org.photonvision.vision.objects.Model;
027import org.photonvision.vision.objects.NullModel;
028import org.photonvision.vision.opencv.DualOffsetValues;
029import org.photonvision.vision.pipe.CVPipe.CVPipeResult;
030import org.photonvision.vision.pipe.impl.*;
031import org.photonvision.vision.pipe.impl.ObjectDetectionPipe.ObjectDetectionPipeParams;
032import org.photonvision.vision.pipeline.result.CVPipelineResult;
033import org.photonvision.vision.target.PotentialTarget;
034import org.photonvision.vision.target.TargetOrientation;
035import org.photonvision.vision.target.TrackedTarget;
036
037public class ObjectDetectionPipeline
038        extends CVPipeline<CVPipelineResult, ObjectDetectionPipelineSettings> {
039    private final CalculateFPSPipe calculateFPSPipe = new CalculateFPSPipe();
040    private final ObjectDetectionPipe objectDetectorPipe = new ObjectDetectionPipe();
041    private final SortContoursPipe sortContoursPipe = new SortContoursPipe();
042    private final Collect2dTargetsPipe collect2dTargetsPipe = new Collect2dTargetsPipe();
043    private final FilterObjectDetectionsPipe filterContoursPipe = new FilterObjectDetectionsPipe();
044
045    private static final FrameThresholdType PROCESSING_TYPE = FrameThresholdType.NONE;
046
047    public ObjectDetectionPipeline() {
048        super(PROCESSING_TYPE);
049        settings = new ObjectDetectionPipelineSettings();
050    }
051
052    public ObjectDetectionPipeline(ObjectDetectionPipelineSettings settings) {
053        super(PROCESSING_TYPE);
054        this.settings = settings;
055    }
056
057    @Override
058    protected void setPipeParamsImpl() {
059        var params = new ObjectDetectionPipeParams();
060        params.confidence = settings.confidence;
061        params.nms = settings.nms;
062        Optional<Model> selectedModel =
063                NeuralNetworkModelManager.getInstance().getModel(settings.model);
064
065        // If the desired model couldn't be found, log an error and try to use the default model
066        if (selectedModel.isEmpty()) {
067            selectedModel = NeuralNetworkModelManager.getInstance().getDefaultModel();
068        }
069
070        // If the model remains empty, use the NullModel
071        if (selectedModel.isEmpty()) {
072            selectedModel = Optional.of(NullModel.getInstance());
073        }
074
075        params.model = selectedModel.get();
076
077        objectDetectorPipe.setParams(params);
078
079        DualOffsetValues dualOffsetValues =
080                new DualOffsetValues(
081                        settings.offsetDualPointA,
082                        settings.offsetDualPointAArea,
083                        settings.offsetDualPointB,
084                        settings.offsetDualPointBArea);
085
086        SortContoursPipe.SortContoursParams sortContoursParams =
087                new SortContoursPipe.SortContoursParams(
088                        settings.contourSortMode,
089                        settings.outputShowMultipleTargets ? MAX_MULTI_TARGET_RESULTS : 1,
090                        frameStaticProperties);
091        sortContoursPipe.setParams(sortContoursParams);
092
093        var filterContoursParams =
094                new FilterObjectDetectionsPipe.FilterContoursParams(
095                        settings.contourArea,
096                        settings.contourRatio,
097                        frameStaticProperties,
098                        settings.contourTargetOrientation == TargetOrientation.Landscape);
099        filterContoursPipe.setParams(filterContoursParams);
100
101        Collect2dTargetsPipe.Collect2dTargetsParams collect2dTargetsParams =
102                new Collect2dTargetsPipe.Collect2dTargetsParams(
103                        settings.offsetRobotOffsetMode,
104                        settings.offsetSinglePoint,
105                        dualOffsetValues,
106                        settings.contourTargetOffsetPointEdge,
107                        settings.contourTargetOrientation,
108                        frameStaticProperties);
109        collect2dTargetsPipe.setParams(collect2dTargetsParams);
110    }
111
112    @Override
113    protected CVPipelineResult process(Frame frame, ObjectDetectionPipelineSettings settings) {
114        long sumPipeNanosElapsed = 0;
115
116        // ***************** change based on backend ***********************
117
118        CVPipeResult<List<NeuralNetworkPipeResult>> rknnResult =
119                objectDetectorPipe.run(frame.colorImage);
120        sumPipeNanosElapsed += rknnResult.nanosElapsed;
121
122        var names = objectDetectorPipe.getClassNames();
123
124        frame.colorImage.getMat().copyTo(frame.processedImage.getMat());
125
126        // ***************** change based on backend ***********************
127
128        var filterContoursResult = filterContoursPipe.run(rknnResult.output);
129        sumPipeNanosElapsed += filterContoursResult.nanosElapsed;
130
131        CVPipeResult<List<PotentialTarget>> sortContoursResult =
132                sortContoursPipe.run(
133                        filterContoursResult.output.stream()
134                                .map(shape -> new PotentialTarget(shape))
135                                .collect(Collectors.toList()));
136        sumPipeNanosElapsed += sortContoursResult.nanosElapsed;
137
138        CVPipeResult<List<TrackedTarget>> collect2dTargetsResult =
139                collect2dTargetsPipe.run(sortContoursResult.output);
140        sumPipeNanosElapsed += collect2dTargetsResult.nanosElapsed;
141
142        var fpsResult = calculateFPSPipe.run(null);
143        var fps = fpsResult.output;
144
145        return new CVPipelineResult(
146                frame.sequenceID, sumPipeNanosElapsed, fps, collect2dTargetsResult.output, frame, names);
147    }
148
149    @Override
150    public void release() {
151        objectDetectorPipe.release();
152        super.release();
153    }
154}