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.util.List;
021import java.util.Optional;
022import org.opencv.core.Mat;
023import org.photonvision.common.configuration.NeuralNetworkModelManager;
024import org.photonvision.vision.objects.Model;
025import org.photonvision.vision.objects.NullModel;
026import org.photonvision.vision.objects.ObjectDetector;
027import org.photonvision.vision.opencv.CVMat;
028import org.photonvision.vision.opencv.Releasable;
029import org.photonvision.vision.pipe.CVPipe;
030
031public class ObjectDetectionPipe
032        extends CVPipe<
033                CVMat, List<NeuralNetworkPipeResult>, ObjectDetectionPipe.ObjectDetectionPipeParams>
034        implements Releasable {
035    private ObjectDetector detector;
036
037    public ObjectDetectionPipe() {
038        Optional<Model> defaultModel = NeuralNetworkModelManager.getInstance().getDefaultModel();
039        detector = defaultModel.map(Model::load).orElse(NullModel.getInstance());
040    }
041
042    @Override
043    protected List<NeuralNetworkPipeResult> process(CVMat in) {
044        // Check if the model has changed
045        if (detector.getModel() != params.model) {
046            detector.release();
047            detector = params.model.load();
048        }
049
050        Mat frame = in.getMat();
051        if (frame.empty()) {
052            return List.of();
053        }
054
055        return detector.detect(in.getMat(), params.nms, params.confidence);
056    }
057
058    public static class ObjectDetectionPipeParams {
059        public double confidence;
060        public double nms;
061        public int max_detections;
062        public Model model;
063
064        public ObjectDetectionPipeParams() {}
065    }
066
067    public List<String> getClassNames() {
068        return detector.getClasses();
069    }
070
071    @Override
072    public void release() {
073        detector.release();
074    }
075}