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.jni;
019
020import java.awt.Color;
021import java.lang.ref.Cleaner;
022import java.util.List;
023import java.util.concurrent.atomic.AtomicBoolean;
024import org.opencv.core.Mat;
025import org.opencv.core.Size;
026import org.photonvision.common.logging.LogGroup;
027import org.photonvision.common.logging.Logger;
028import org.photonvision.common.util.ColorHelper;
029import org.photonvision.rknn.RknnJNI;
030import org.photonvision.vision.objects.Letterbox;
031import org.photonvision.vision.objects.ObjectDetector;
032import org.photonvision.vision.objects.RknnModel;
033import org.photonvision.vision.pipe.impl.NeuralNetworkPipeResult;
034
035/** Manages an object detector using the rknn backend. */
036public class RknnObjectDetector implements ObjectDetector {
037    private static final Logger logger = new Logger(RknnDetectorJNI.class, LogGroup.General);
038
039    /** Cleaner instance to release the detector when it goes out of scope */
040    private final Cleaner cleaner = Cleaner.create();
041
042    /** Atomic boolean to ensure that the native object can only be released once. */
043    private AtomicBoolean released = new AtomicBoolean(false);
044
045    /** Pointer to the native object */
046    private final long objPointer;
047
048    private final RknnModel model;
049
050    private final Size inputSize;
051
052    /** Returns the model in use by this detector. */
053    @Override
054    public RknnModel getModel() {
055        return model;
056    }
057
058    /**
059     * Creates a new RknnObjectDetector from the given model.
060     *
061     * @param model The model to create the detector from.
062     * @param inputSize The required image dimensions for the model. Images will be {@link
063     *     Letterbox}ed to this shape.
064     */
065    public RknnObjectDetector(RknnModel model, Size inputSize) {
066        this.model = model;
067        this.inputSize = inputSize;
068
069        // Create the detector
070        objPointer =
071                RknnJNI.create(model.modelFile.getPath(), model.labels.size(), model.version.ordinal(), -1);
072        if (objPointer <= 0) {
073            throw new RuntimeException(
074                    "Failed to create detector from path " + model.modelFile.getPath());
075        }
076
077        logger.debug("Created detector for model " + model.modelFile.getName());
078
079        // Register the cleaner to release the detector when it goes out of scope
080        cleaner.register(this, this::release);
081    }
082
083    /**
084     * Returns the classes that the detector can detect
085     *
086     * @return The classes
087     */
088    @Override
089    public List<String> getClasses() {
090        return model.labels;
091    }
092
093    /**
094     * Detects objects in the given input image using the RknnDetector.
095     *
096     * @param in The input image to perform object detection on.
097     * @param nmsThresh The threshold value for non-maximum suppression.
098     * @param boxThresh The threshold value for bounding box detection.
099     * @return A list of NeuralNetworkPipeResult objects representing the detected objects. Returns an
100     *     empty list if the detector is not initialized or if no objects are detected.
101     */
102    @Override
103    public List<NeuralNetworkPipeResult> detect(Mat in, double nmsThresh, double boxThresh) {
104        if (objPointer <= 0) {
105            // Report error and make sure to include the model name
106            logger.error("Detector is not initialized! Model: " + model.modelFile.getName());
107            return List.of();
108        }
109
110        // Resize the frame to the input size of the model
111        Mat letterboxed = new Mat();
112        Letterbox scale =
113                Letterbox.letterbox(in, letterboxed, this.inputSize, ColorHelper.colorToScalar(Color.GRAY));
114        if (!letterboxed.size().equals(this.inputSize)) {
115            letterboxed.release();
116            throw new RuntimeException("Letterboxed frame is not the right size!");
117        }
118
119        // Detect objects in the letterboxed frame
120        var results = RknnJNI.detect(objPointer, letterboxed.getNativeObjAddr(), nmsThresh, boxThresh);
121
122        letterboxed.release();
123
124        if (results == null) {
125            return List.of();
126        }
127
128        return scale.resizeDetections(
129                List.of(results).stream()
130                        .map(it -> new NeuralNetworkPipeResult(it.rect, it.class_id, it.conf))
131                        .toList());
132    }
133
134    /** Thread-safe method to release the detector. */
135    @Override
136    public void release() {
137        // Checks if the atomic is 'false', and if so, sets it to 'true'
138        if (released.compareAndSet(false, true)) {
139            if (objPointer <= 0) {
140                logger.error(
141                        "Detector is not initialized, and so it can't be released! Model: "
142                                + model.modelFile.getName());
143                return;
144            }
145
146            RknnJNI.destroy(objPointer);
147            logger.debug("Released detector for model " + model.modelFile.getName());
148        }
149    }
150}