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 38b4340889..c282ec7a79 100644 --- a/photon-core/src/main/java/org/photonvision/jni/RknnDetectorJNI.java +++ b/photon-core/src/main/java/org/photonvision/jni/RknnDetectorJNI.java @@ -22,12 +22,12 @@ import java.util.List; import java.util.concurrent.CopyOnWriteArrayList; import java.util.stream.Collectors; +import org.opencv.core.Mat; 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.opencv.CVMat; import org.photonvision.vision.pipe.impl.NeuralNetworkPipeResult; public class RknnDetectorJNI extends PhotonJNICommon { @@ -65,16 +65,38 @@ public static class RknnObjectDetector { long objPointer = -1; private List labels; private final Object lock = new Object(); - private static final CopyOnWriteArrayList detectors = new CopyOnWriteArrayList<>(); + private static final CopyOnWriteArrayList detectors = + new CopyOnWriteArrayList<>(); + + static volatile boolean hook = false; public RknnObjectDetector(String modelPath, List labels, RknnJNI.ModelVersion version) { synchronized (lock) { objPointer = RknnJNI.create(modelPath, labels.size(), version.ordinal(), -1); - detectors.add(objPointer); - System.out.println( - "Created " + objPointer + "! Detectors: " + Arrays.toString(detectors.toArray())); + 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() { @@ -89,14 +111,14 @@ public List getClasses() { * @param boxThresh Minimum confidence for a box to be added. Basically just confidence * threshold */ - public List detect(CVMat in, double nmsThresh, double boxThresh) { + 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.getMat().getNativeObjAddr(), nmsThresh, boxThresh); + ret = RknnJNI.detect(objPointer, in.getNativeObjAddr(), nmsThresh, boxThresh); } else { logger.warn("Detect called after destroy -- giving up"); return List.of(); @@ -114,7 +136,7 @@ public void release() { synchronized (lock) { if (objPointer > 0) { RknnJNI.destroy(objPointer); - detectors.remove(objPointer); + detectors.remove(this); System.out.println( "Killed " + objPointer + "! Detectors: " + Arrays.toString(detectors.toArray())); objPointer = -1; @@ -124,14 +146,4 @@ public void release() { } } } - - // public static void createRknnDetector() { - // objPointer = - // RknnJNI.create( - // NeuralNetworkModelManager.getInstance() - // .getDefaultRknnModel() - // .getAbsolutePath() - // .toString(), - // NeuralNetworkModelManager.getInstance().getLabels().size()); - // } } diff --git a/photon-core/src/main/java/org/photonvision/vision/pipe/impl/FilterObjectDetectionsPipe.java b/photon-core/src/main/java/org/photonvision/vision/pipe/impl/FilterObjectDetectionsPipe.java index c2fff471e8..d22556668a 100644 --- a/photon-core/src/main/java/org/photonvision/vision/pipe/impl/FilterObjectDetectionsPipe.java +++ b/photon-core/src/main/java/org/photonvision/vision/pipe/impl/FilterObjectDetectionsPipe.java @@ -41,7 +41,7 @@ protected List process(List in } private void filterContour(NeuralNetworkPipeResult contour) { - var boc = contour.box; + var boc = contour.bbox; // Area filtering double areaPercentage = boc.area() / params.getFrameStaticProperties().imageArea * 100.0; diff --git a/photon-core/src/main/java/org/photonvision/vision/pipe/impl/NeuralNetworkPipeResult.java b/photon-core/src/main/java/org/photonvision/vision/pipe/impl/NeuralNetworkPipeResult.java index 575c64e0e3..4fce74955e 100644 --- a/photon-core/src/main/java/org/photonvision/vision/pipe/impl/NeuralNetworkPipeResult.java +++ b/photon-core/src/main/java/org/photonvision/vision/pipe/impl/NeuralNetworkPipeResult.java @@ -20,13 +20,13 @@ import org.opencv.core.Rect2d; public class NeuralNetworkPipeResult { - public NeuralNetworkPipeResult(Rect2d box2, Integer classIdx, Float confidence) { - box = box2; + public NeuralNetworkPipeResult(Rect2d boundingBox, int classIdx, double confidence) { + bbox = boundingBox; this.classIdx = classIdx; this.confidence = confidence; } public final int classIdx; - public final Rect2d box; + public final Rect2d bbox; public final double confidence; } 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 81dc6bf9c9..b037bdd6f5 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 @@ -17,8 +17,17 @@ package org.photonvision.vision.pipe.impl; +import java.awt.Color; +import java.util.ArrayList; import java.util.List; +import org.opencv.core.Core; +import org.opencv.core.Mat; +import org.opencv.core.Rect2d; +import org.opencv.core.Scalar; +import org.opencv.core.Size; +import org.opencv.imgproc.Imgproc; import org.photonvision.common.configuration.NeuralNetworkModelManager; +import org.photonvision.common.util.ColorHelper; import org.photonvision.jni.RknnDetectorJNI.RknnObjectDetector; import org.photonvision.vision.opencv.CVMat; import org.photonvision.vision.opencv.Releasable; @@ -30,8 +39,10 @@ public class RknnDetectionPipe 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. + // 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().getAbsolutePath(), @@ -39,6 +50,18 @@ public RknnDetectionPipe() { NeuralNetworkModelManager.getInstance().getModelVersion()); } + private static class Letterbox { + double dx; + double dy; + double scale; + + public Letterbox(double dx, double dy, double scale) { + this.dx = dx; + this.dy = dy; + this.scale = scale; + } + } + @Override protected List process(CVMat in) { var frame = in.getMat(); @@ -48,7 +71,66 @@ protected List process(CVMat in) { return List.of(); } - return detector.detect(in, params.nms, params.confidence); + // 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"); + } + var ret = detector.detect(letterboxed, params.nms, params.confidence); + + return resizeDetections(ret, scale); + } + + private List resizeDetections( + List unscaled, Letterbox letterbox) { + var ret = new ArrayList(); + + for (var t : unscaled) { + var scale = 1.0 / letterbox.scale; + var boundingBox = t.bbox; + double x = (boundingBox.x - letterbox.dx) * scale; + double y = (boundingBox.y - letterbox.dy) * scale; + double width = boundingBox.width * scale; + double height = boundingBox.height * scale; + + ret.add( + new NeuralNetworkPipeResult(new Rect2d(x, y, width, height), t.classIdx, t.confidence)); + } + + return ret; + } + + 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(); + var r = Math.min(newShape.height / frameSize.height, newShape.width / frameSize.width); + + var newUnpad = new Size(Math.round(frameSize.width * r), Math.round(frameSize.height * r)); + + if (!(frameSize.equals(newUnpad))) { + Imgproc.resize(frame, letterboxed, newUnpad, Imgproc.INTER_LINEAR); + } else { + frame.copyTo(letterboxed); + } + + var dw = newShape.width - newUnpad.width; + var dh = newShape.height - newUnpad.height; + + dw /= 2; + dh /= 2; + + int top = (int) (Math.round(dh - 0.1f)); + int bottom = (int) (Math.round(dh + 0.1f)); + int left = (int) (Math.round(dw - 0.1f)); + int right = (int) (Math.round(dw + 0.1f)); + Core.copyMakeBorder( + letterboxed, letterboxed, top, bottom, left, right, Core.BORDER_CONSTANT, color); + + return new Letterbox(dw, dh, r); } public static class RknnDetectionPipeParams { diff --git a/photon-core/src/main/java/org/photonvision/vision/target/PotentialTarget.java b/photon-core/src/main/java/org/photonvision/vision/target/PotentialTarget.java index aba0dc5ccf..01f8dddfc3 100644 --- a/photon-core/src/main/java/org/photonvision/vision/target/PotentialTarget.java +++ b/photon-core/src/main/java/org/photonvision/vision/target/PotentialTarget.java @@ -56,7 +56,7 @@ public PotentialTarget(Contour inputContour, CVShape shape) { } public PotentialTarget(NeuralNetworkPipeResult det) { - this.shape = new CVShape(new Contour(det.box), ContourShape.Quadrilateral); + this.shape = new CVShape(new Contour(det.bbox), ContourShape.Quadrilateral); this.m_mainContour = this.shape.getContour(); m_subContours = List.of(); this.clsId = det.classIdx;