diff --git a/photon-core/src/main/java/org/photonvision/common/configuration/NeuralNetworkModelManager.java b/photon-core/src/main/java/org/photonvision/common/configuration/NeuralNetworkModelManager.java index ba6c2267b..6877c49d8 100644 --- a/photon-core/src/main/java/org/photonvision/common/configuration/NeuralNetworkModelManager.java +++ b/photon-core/src/main/java/org/photonvision/common/configuration/NeuralNetworkModelManager.java @@ -26,6 +26,7 @@ import java.nio.file.StandardCopyOption; import java.util.ArrayList; import java.util.List; +import org.opencv.core.Size; import org.photonvision.common.logging.LogGroup; import org.photonvision.common.logging.Logger; import org.photonvision.rknn.RknnJNI; @@ -97,15 +98,36 @@ public class Model { public final File modelFile; public final RknnJNI.ModelVersion version; public final List labels; + public final Size inputSize; + /** + * Model constructor. + * + * @param model format `name-width-height-model.format` + * @param labels + * @throws IllegalArgumentException + */ public Model(String model, String labels) throws IllegalArgumentException { - this.version = getModelVersion(model); + String[] parts = model.split("-"); + if (parts.length != 4) { + throw new IllegalArgumentException("Invalid model file name: " + model); + } + + // TODO: model 'version' need to be replaced the by the product of 'Version' x 'Format' + this.version = getModelVersion(parts[3]); + + int width = Integer.parseInt(parts[1]); + int height = Integer.parseInt(parts[2]); + this.inputSize = new Size(width, height); + this.modelFile = new File(model); try { this.labels = Files.readAllLines(Paths.get(labels)); } catch (IOException e) { - throw new IllegalArgumentException("Error reading labels file " + labels, e); + throw new IllegalArgumentException("Failed to read labels file " + labels, e); } + + logger.info("Loaded model " + modelFile.getName()); } public String getPath() { @@ -141,7 +163,7 @@ public List getModels() { /** * Returns the model with the given name. * - *

TODO: Java 17 This should return an Optional instead of null. + *

TODO: Java 17 This should return an Optional Model instead of null. * * @param modelName The model name * @return The model @@ -214,7 +236,7 @@ public void extractModels(File modelsFolder) { modelsFolder.mkdirs(); } - String resourcePath = "models"; // Adjust path if necessary + String resourcePath = "models"; try { URL resourceURL = NeuralNetworkModelManager.class.getClassLoader().getResource(resourcePath); if (resourceURL == null) { diff --git a/photon-core/src/main/java/org/photonvision/jni/RknnDetectorJNI.java b/photon-core/src/main/java/org/photonvision/jni/RknnDetectorJNI.java index a24781805..6ad21fd9c 100644 --- a/photon-core/src/main/java/org/photonvision/jni/RknnDetectorJNI.java +++ b/photon-core/src/main/java/org/photonvision/jni/RknnDetectorJNI.java @@ -18,21 +18,10 @@ package org.photonvision.jni; import java.io.IOException; -import java.util.Arrays; import java.util.List; -import java.util.concurrent.CopyOnWriteArrayList; -import java.util.stream.Collectors; -import org.opencv.core.Mat; -import org.photonvision.common.configuration.NeuralNetworkModelManager; -import org.photonvision.common.logging.LogGroup; -import org.photonvision.common.logging.Logger; import org.photonvision.common.util.TestUtils; -import org.photonvision.rknn.RknnJNI; -import org.photonvision.rknn.RknnJNI.RknnResult; -import org.photonvision.vision.pipe.impl.NeuralNetworkPipeResult; public class RknnDetectorJNI extends PhotonJNICommon { - private static final Logger logger = new Logger(RknnDetectorJNI.class, LogGroup.General); private boolean isLoaded; private static RknnDetectorJNI instance = null; @@ -61,94 +50,4 @@ public boolean isLoaded() { public void setLoaded(boolean state) { isLoaded = state; } - - public static class RknnObjectDetector { - long objPointer = -1; - private List labels; - private final Object lock = new Object(); - private static final CopyOnWriteArrayList detectors = - new CopyOnWriteArrayList<>(); - - static volatile boolean hook = false; - - public RknnObjectDetector(NeuralNetworkModelManager.Model model) { - this(model.getPath(), model.labels, model.version); - } - - public RknnObjectDetector(String modelPath, List labels, RknnJNI.ModelVersion version) { - synchronized (lock) { - objPointer = RknnJNI.create(modelPath, labels.size(), version.ordinal(), -1); - detectors.add(this); - logger.debug( - "Created detector " - + objPointer - + " from path " - + modelPath - + "! Detectors: " - + Arrays.toString(detectors.toArray())); - } - this.labels = labels; - - // the kernel should probably alredy deal with this for us, but I'm gunna be paranoid anyways. - if (!hook) { - Runtime.getRuntime() - .addShutdownHook( - new Thread( - () -> { - System.err.println("Shutdown hook rknn"); - for (var d : detectors) { - d.release(); - } - })); - hook = true; - } - } - - public List getClasses() { - return labels; - } - - /** - * Detect forwards using this model - * - * @param in The image to process - * @param nmsThresh Non-maximum supression threshold. Probably should not change - * @param boxThresh Minimum confidence for a box to be added. Basically just confidence - * threshold - */ - public List detect(Mat in, double nmsThresh, double boxThresh) { - RknnResult[] ret; - synchronized (lock) { - // We can technically be asked to detect and the lock might be acquired _after_ release has - // been called. This would mean objPointer would be invalid which would call everything to - // explode. - if (objPointer > 0) { - ret = RknnJNI.detect(objPointer, in.getNativeObjAddr(), nmsThresh, boxThresh); - } else { - logger.warn("Detect called after destroy -- giving up"); - return List.of(); - } - } - if (ret == null) { - return List.of(); - } - return List.of(ret).stream() - .map(it -> new NeuralNetworkPipeResult(it.rect, it.class_id, it.conf)) - .collect(Collectors.toList()); - } - - public void release() { - synchronized (lock) { - if (objPointer > 0) { - RknnJNI.destroy(objPointer); - detectors.remove(this); - System.out.println( - "Killed " + objPointer + "! Detectors: " + Arrays.toString(detectors.toArray())); - objPointer = -1; - } else { - logger.error("RKNN Detector has already been destroyed!"); - } - } - } - } } diff --git a/photon-core/src/main/java/org/photonvision/jni/RknnObjectDetector.java b/photon-core/src/main/java/org/photonvision/jni/RknnObjectDetector.java new file mode 100644 index 000000000..03a67045d --- /dev/null +++ b/photon-core/src/main/java/org/photonvision/jni/RknnObjectDetector.java @@ -0,0 +1,116 @@ +package org.photonvision.jni; + +import java.lang.ref.Cleaner; +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; +import org.opencv.core.Mat; +import org.photonvision.common.configuration.NeuralNetworkModelManager; +import org.photonvision.common.logging.LogGroup; +import org.photonvision.common.logging.Logger; +import org.photonvision.rknn.RknnJNI; +import org.photonvision.vision.opencv.Releasable; +import org.photonvision.vision.pipe.impl.NeuralNetworkPipeResult; + +/** + * A class to represent an object detector using the Rknn library. + * + *

TODO: When we start supporting more platforms, we should consider moving most of this code + * into a common "ObjectDetector" class to define the common interface for all object detectors. + */ +public class RknnObjectDetector implements Releasable { + /** logger for the RknnObjectDetector */ + private static final Logger logger = new Logger(RknnDetectorJNI.class, LogGroup.General); + + /** Cleaner instance to release the detector when it is no longer needed */ + private final Cleaner cleaner = Cleaner.create(); + + /** Pointer to the native object */ + private final long objPointer; + + /** Model configuration */ + private final NeuralNetworkModelManager.Model model; + + /** + * Returns the model used by the detector. + */ + public NeuralNetworkModelManager.Model getModel() { + return model; + } + + /** Atomic boolean to ensure that the detector is only released _once_. */ + private AtomicBoolean released = new AtomicBoolean(false); + + /** + * Creates a new RknnObjectDetector from the given model. + * + * @param model The model to create the detector from. + */ + public RknnObjectDetector(NeuralNetworkModelManager.Model model) { + this.model = model; + + // Create the detector + objPointer = RknnJNI.create(model.getPath(), model.labels.size(), model.version.ordinal(), -1); + if (objPointer <= 0) { + throw new RuntimeException("Failed to create detector from path " + model.getPath()); + } + + logger.debug("Created detector for model " + model.modelFile.getName()); + + // Register the cleaner to release the detector when it goes out of scope + cleaner.register(this, this::release); + + // Set the detector to be released when the JVM exits + Runtime.getRuntime().addShutdownHook(new Thread(this::release)); + } + + /** + * Returns the classes that the detector can detect + * + * @return The classes + */ + public List getClasses() { + return model.labels; + } + + /** + * Detects objects in the given input image using the RknnDetector. + * + * @param in The input image to perform object detection on. + * @param nmsThresh The threshold value for non-maximum suppression. + * @param boxThresh The threshold value for bounding box detection. + * @return A list of NeuralNetworkPipeResult objects representing the detected objects. Returns an + * empty list if the detector is not initialized or if no objects are detected. + */ + public List detect(Mat in, double nmsThresh, double boxThresh) { + if (objPointer <= 0) { + // Report error and make sure to include the model name + logger.error("Detector is not initialized! Model: " + model.modelFile.getName()); + return List.of(); + } + + var results = RknnJNI.detect(objPointer, in.getNativeObjAddr(), nmsThresh, boxThresh); + if (results == null) { + return List.of(); + } + + return List.of(results).stream() + .map(it -> new NeuralNetworkPipeResult(it.rect, it.class_id, it.conf)) + .toList(); + } + + /** Thread-safe method to release the detector. */ + @Override + public void release() { + if (released.compareAndSet(false, true)) { + if (objPointer <= 0) { + logger.error( + "Detector is not initialized, and so it can't be released! Model: " + + model.modelFile.getName()); + return; + } + + RknnJNI.destroy(objPointer); + logger.debug("Released detector for model " + model.modelFile.getName()); + } + } +} diff --git a/photon-core/src/main/java/org/photonvision/vision/pipe/impl/RknnDetectionPipe.java b/photon-core/src/main/java/org/photonvision/vision/pipe/impl/RknnDetectionPipe.java index 70d02e649..59284c5f3 100644 --- a/photon-core/src/main/java/org/photonvision/vision/pipe/impl/RknnDetectionPipe.java +++ b/photon-core/src/main/java/org/photonvision/vision/pipe/impl/RknnDetectionPipe.java @@ -27,23 +27,29 @@ import org.opencv.core.Size; import org.opencv.imgproc.Imgproc; import org.photonvision.common.configuration.NeuralNetworkModelManager; +import org.photonvision.common.configuration.NeuralNetworkModelManager.Model; import org.photonvision.common.util.ColorHelper; -import org.photonvision.jni.RknnDetectorJNI.RknnObjectDetector; +import org.photonvision.jni.RknnObjectDetector; import org.photonvision.vision.opencv.CVMat; import org.photonvision.vision.opencv.Releasable; import org.photonvision.vision.pipe.CVPipe; +/** + * A pipe that uses an rknn model to detect objects in an image. + * + *

TODO: This class should be refactored into a generic "ObjectDetectionPipe" that can use any + * "ObjectDetector" implementation. + */ public class RknnDetectionPipe extends CVPipe, RknnDetectionPipe.RknnDetectionPipeParams> implements Releasable { + private RknnObjectDetector detector; public RknnDetectionPipe() { - // For now this is hard-coded to defaults. Should be refactored into set pipe - // params, though. And ideally a little wrapper helper for only changing native stuff on content - // change created. - this.detector = - new RknnObjectDetector(NeuralNetworkModelManager.getInstance().getDefaultRknnModel()); + // Default model + Model model = NeuralNetworkModelManager.getInstance().getDefaultRknnModel(); + this.detector = new RknnObjectDetector(model); } private static class Letterbox { @@ -60,28 +66,43 @@ public Letterbox(double dx, double dy, double scale) { @Override protected List process(CVMat in) { - var frame = in.getMat(); + // Check that we're still using the same model, and if not, release the old one and load the new one + // + // Pointer equality is fine here because Model are immutable and the ModelManager will always return the same instance + // for the same model file. + if (detector.getModel() != params.model) { + detector.release(); + detector = new RknnObjectDetector(params.model); + } - // Make sure we don't get a weird empty frame + Mat frame = in.getMat(); if (frame.empty()) { return List.of(); } - // letterbox - var letterboxed = new Mat(); - var scale = - letterbox(frame, letterboxed, new Size(640, 640), ColorHelper.colorToScalar(Color.GRAY)); - - if (letterboxed.width() != 640 || letterboxed.height() != 640) { - // huh whack give up lol - throw new RuntimeException("RGA bugged but still wrong size"); + // Resize the frame to the input size of the model + Size shape = this.params.model.inputSize; + Mat letterboxed = new Mat(); + Letterbox scale = letterbox(frame, letterboxed, shape, ColorHelper.colorToScalar(Color.GRAY)); + if (!letterboxed.size().equals(shape)) { + throw new RuntimeException("Letterboxed frame is not the right size!"); } - var ret = detector.detect(letterboxed, params.nms, params.confidence); + + // Detect objects in the letterboxed frame + List ret = detector.detect(letterboxed, params.nms, params.confidence); letterboxed.release(); + // Resize the detections to the original frame size return resizeDetections(ret, scale); } + /** + * Resizes the detections to the original frame size. + * + * @param unscaled The detections to resize + * @param letterbox The letterbox information + * @return The resized detections + */ private List resizeDetections( List unscaled, Letterbox letterbox) { var ret = new ArrayList(); @@ -101,6 +122,18 @@ private List resizeDetections( return ret; } + /** + * Resize the frame to the new shape and "letterbox" it. + * + *

Letterboxing is the process of resizing an image to a new shape while maintaining the aspect + * ratio of the original image. The new image is padded with a color to fill the remaining space. + * + * @param frame + * @param letterboxed + * @param newShape + * @param color + * @return + */ private static Letterbox letterbox(Mat frame, Mat letterboxed, Size newShape, Scalar color) { // from https://github.com/ultralytics/yolov5/issues/8427#issuecomment-1172469631 var frameSize = frame.size(); @@ -134,6 +167,7 @@ public static class RknnDetectionPipeParams { public double confidence; public double nms; public int max_detections; + public NeuralNetworkModelManager.Model model; public RknnDetectionPipeParams() {} } diff --git a/photon-core/src/main/java/org/photonvision/vision/pipeline/ObjectDetectionPipeline.java b/photon-core/src/main/java/org/photonvision/vision/pipeline/ObjectDetectionPipeline.java index 4919c9151..6f35ec519 100644 --- a/photon-core/src/main/java/org/photonvision/vision/pipeline/ObjectDetectionPipeline.java +++ b/photon-core/src/main/java/org/photonvision/vision/pipeline/ObjectDetectionPipeline.java @@ -19,6 +19,7 @@ import java.util.List; import java.util.stream.Collectors; +import org.photonvision.common.configuration.NeuralNetworkModelManager; import org.photonvision.vision.frame.Frame; import org.photonvision.vision.frame.FrameThresholdType; import org.photonvision.vision.opencv.DualOffsetValues; @@ -56,6 +57,7 @@ protected void setPipeParamsImpl() { var params = new RknnDetectionPipeParams(); params.confidence = settings.confidence; params.nms = settings.nms; + params.model = NeuralNetworkModelManager.getInstance().getModel(settings.model); rknnPipe.setParams(params); DualOffsetValues dualOffsetValues = @@ -99,7 +101,6 @@ protected CVPipelineResult process(Frame frame, ObjectDetectionPipelineSettings CVPipeResult> rknnResult = rknnPipe.run(frame.colorImage); sumPipeNanosElapsed += rknnResult.nanosElapsed; - List targetList; var names = rknnPipe.getClassNames(); diff --git a/photon-core/src/main/java/org/photonvision/vision/pipeline/ObjectDetectionPipelineSettings.java b/photon-core/src/main/java/org/photonvision/vision/pipeline/ObjectDetectionPipelineSettings.java index f074bf1f5..d12b108c6 100644 --- a/photon-core/src/main/java/org/photonvision/vision/pipeline/ObjectDetectionPipelineSettings.java +++ b/photon-core/src/main/java/org/photonvision/vision/pipeline/ObjectDetectionPipelineSettings.java @@ -20,6 +20,7 @@ public class ObjectDetectionPipelineSettings extends AdvancedPipelineSettings { public double confidence; public double nms; // non maximal suppression + public String model; public ObjectDetectionPipelineSettings() { super();