init

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);
             //转换数据为张量

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<String, String> getImagePathMap(String imagePath){
+        Map<String, String> 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;
+    }
+
+}

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<String, String> 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<PlateInfo> 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();
+        }
+    }
+}