From f3105976a928d7940f31ee57e165d67ae367ab5d Mon Sep 17 00:00:00 2001 From: diven Date: Tue, 17 Jan 2023 16:24:26 +0800 Subject: [PATCH] init --- .../anpr/core/models/TorchPlateDetection.java | 13 +++- .../visual/open/anpr/core/base/BaseTest.java | 28 +++++++ .../anpr/core/extract/PlateExtractorTest.java | 78 +++++++++++++++++++ 3 files changed, 116 insertions(+), 3 deletions(-) create mode 100644 open-anpr-core/src/test/java/com/visual/open/anpr/core/base/BaseTest.java create mode 100644 open-anpr-core/src/test/java/com/visual/open/anpr/core/extract/PlateExtractorTest.java diff --git a/open-anpr-core/src/main/java/com/visual/open/anpr/core/models/TorchPlateDetection.java b/open-anpr-core/src/main/java/com/visual/open/anpr/core/models/TorchPlateDetection.java index 46c666e..242aff7 100644 --- a/open-anpr-core/src/main/java/com/visual/open/anpr/core/models/TorchPlateDetection.java +++ b/open-anpr-core/src/main/java/com/visual/open/anpr/core/models/TorchPlateDetection.java @@ -19,9 +19,13 @@ import java.util.stream.Collectors; * git：https://github.com/we0091234/Chinese_license_plate_detection_recognition */ public class TorchPlateDetection extends BaseOnnxInfer implements PlateDetection { - private static int imageWidth = 640; - private static int imageHeight= 640; - private static Scalar border = new Scalar(114, 114, 114); + private final static int imageWidth = 640; + private final static int imageHeight= 640; + private final static Scalar border = new Scalar(114, 114, 114); + //人脸预测分数阈值 + public final static float defScoreTh = 0.3f; + //人脸重叠iou阈值 + public final static float defIouTh = 0.5f; public TorchPlateDetection(String modelPath, int threads) { super(modelPath, threads); @@ -34,6 +38,9 @@ public class TorchPlateDetection extends BaseOnnxInfer implements PlateDetection BorderMat makeBorderMat = null; ImageMat imageMat = image.clone(); try { + //清理参数 + iouTh = iouTh <= 0 ? defIouTh : iouTh; + scoreTh = scoreTh <= 0 ? defScoreTh : scoreTh; //对图像进行标准宽高的处理 makeBorderMat = resizeAndMakeBorderMat(imageMat.toCvMat(), imageWidth, imageHeight); //转换数据为张量 diff --git a/open-anpr-core/src/test/java/com/visual/open/anpr/core/base/BaseTest.java b/open-anpr-core/src/test/java/com/visual/open/anpr/core/base/BaseTest.java new file mode 100644 index 0000000..b67847c --- /dev/null +++ b/open-anpr-core/src/test/java/com/visual/open/anpr/core/base/BaseTest.java @@ -0,0 +1,28 @@ +package com.visual.open.anpr.core.base; + +import ai.onnxruntime.OrtEnvironment; + +import java.io.File; +import java.util.Map; +import java.util.TreeMap; + +public abstract class BaseTest { + + //静态加载动态链接库 + static{ nu.pattern.OpenCV.loadShared(); } + private OrtEnvironment env = OrtEnvironment.getEnvironment(); + + public static Map getImagePathMap(String imagePath){ + Map map = new TreeMap<>(); + File file = new File(imagePath); + if(file.isFile()){ + map.put(file.getName(), file.getAbsolutePath()); + }else if(file.isDirectory()){ + for(File tmpFile : file.listFiles()){ + map.putAll(getImagePathMap(tmpFile.getPath())); + } + } + return map; + } + +} diff --git a/open-anpr-core/src/test/java/com/visual/open/anpr/core/extract/PlateExtractorTest.java b/open-anpr-core/src/test/java/com/visual/open/anpr/core/extract/PlateExtractorTest.java new file mode 100644 index 0000000..ff22c20 --- /dev/null +++ b/open-anpr-core/src/test/java/com/visual/open/anpr/core/extract/PlateExtractorTest.java @@ -0,0 +1,78 @@ +package com.visual.open.anpr.core.extract; + +import com.visual.open.anpr.core.base.BaseTest; +import com.visual.open.anpr.core.domain.*; +import com.visual.open.anpr.core.models.TorchPlateDetection; +import com.visual.open.anpr.core.models.TorchPlateRecognition; +import org.opencv.core.Mat; +import org.opencv.core.Point; +import org.opencv.core.Scalar; +import org.opencv.highgui.HighGui; +import org.opencv.imgcodecs.Imgcodecs; +import org.opencv.imgproc.Imgproc; + +import java.awt.*; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +public class PlateExtractorTest extends BaseTest { + + private static String plateDetectionPath = "open-anpr-core/src/main/resources/models/plate_detect.onnx"; + private static String plateRecognitionPath = "open-anpr-core/src/main/resources/models/plate_rec_color.onnx"; + + public static void main(String[] args) { + TorchPlateDetection torchPlateDetection = new TorchPlateDetection(plateDetectionPath, 1); + TorchPlateRecognition torchPlateRecognition = new TorchPlateRecognition(plateRecognitionPath, 1); + PlateExtractor extractor = new PlateExtractorImpl(torchPlateDetection, torchPlateRecognition); + + String imagePath = "open-anpr-core/src/test/resources/images"; + Map map = getImagePathMap(imagePath); + for(String fileName : map.keySet()) { + String imageFilePath = map.get(fileName); + System.out.println(imageFilePath); + Mat image = Imgcodecs.imread(imageFilePath); + long s = System.currentTimeMillis(); + ExtParam extParam = ExtParam.build() + .setTopK(20) + .setScoreTh(0.3f) + .setIouTh(0.5f); + + PlateImage plateImage = extractor.extract(ImageMat.fromCVMat(image), extParam, new HashMap<>()); + List plateInfos = plateImage.PlateInfos(); + + DrawImage drawImage = DrawImage.build(imageFilePath); + for(PlateInfo plateInfo: plateInfos){ + //画框 + PlateInfo.Point [] points = plateInfo.box.toArray(); + for(int i =0; i< points.length; i++){ + if(i+1 == points.length){ + drawImage.drawLine( + new DrawImage.Point((int)points[i].x, (int)points[i].y), + new DrawImage.Point((int)points[0].x, (int)points[0].y), + 2, Color.RED + ); + }else{ + drawImage.drawLine( + new DrawImage.Point((int)points[i].x, (int)points[i].y), + new DrawImage.Point((int)points[i+1].x, (int)points[i+1].y), + 2, Color.RED + ); + } + } + //添加文本 + PlateInfo.ParseInfo parseInfo = plateInfo.parseInfo; + int fonSize = Float.valueOf(plateInfo.box.width() / parseInfo.plateNo.length() * 1.4f).intValue(); + System.out.println(fonSize); + + drawImage.drawText(parseInfo.plateNo, + new DrawImage.Point((int)points[0].x, (int)points[0].y-fonSize*2), fonSize, Color.RED); + drawImage.drawText((plateInfo.single ? "单排" : "双排") + ":" + parseInfo.plateColor, + new DrawImage.Point((int)points[0].x, (int)points[0].y-fonSize), fonSize, Color.RED); + } + //show + ImageMat.fromCVMat(drawImage.toMat()).imShow(); + image.release(); + } + } +}