init

3 years ago · 0ec7da30df
4 changed files with 200 additions and 1 deletions
--- a/open-anpr-core/src/main/java/com/visual/open/anpr/core/base/PlateRecognition.java
+++ b/open-anpr-core/src/main/java/com/visual/open/anpr/core/base/PlateRecognition.java
@ -0,0 +1,10 @@
+package com.visual.open.anpr.core.base;
+
+import java.util.Map;
+import com.visual.open.anpr.core.domain.ImageMat;
+import com.visual.open.anpr.core.domain.PlateInfo.ParseInfo;
+
+public interface PlateRecognition {
+
+    ParseInfo inference(ImageMat image, boolean single, Map<String, Object> params);
+}
--- a/open-anpr-core/src/main/java/com/visual/open/anpr/core/domain/PlateInfo.java
+++ b/open-anpr-core/src/main/java/com/visual/open/anpr/core/domain/PlateInfo.java
@ -138,7 +138,7 @@ public class PlateInfo implements Comparable<PlateInfo>, Serializable {
         * @param rightBottom   右下角坐标
         * @param leftBottom    左下角坐标
         */
-        public PlateBox(Point leftTop, Point rightTop, Point rightBottom, Point leftBottom) {
+        private PlateBox(Point leftTop, Point rightTop, Point rightBottom, Point leftBottom) {
            this.leftTop = leftTop;
            this.rightTop = rightTop;
            this.rightBottom = rightBottom;
@ -311,4 +311,19 @@ public class PlateInfo implements Comparable<PlateInfo>, Serializable {

    }

+    public static class ParseInfo implements Serializable {
+        public String plateNo;
+        public String plateColor;
+        public float colorScore;
+
+        private ParseInfo(String plateNo, String plateColor, float colorScore) {
+            this.plateNo = plateNo;
+            this.plateColor = plateColor;
+            this.colorScore = colorScore;
+        }
+
+        public static ParseInfo build(String plateNo, String plateColor, float colorScore){
+            return new ParseInfo(plateNo, plateColor, colorScore);
+        }
+    }
 }
--- a/open-anpr-core/src/main/java/com/visual/open/anpr/core/models/TorchPlateRecognition.java
+++ b/open-anpr-core/src/main/java/com/visual/open/anpr/core/models/TorchPlateRecognition.java
@ -0,0 +1,145 @@
+package com.visual.open.anpr.core.models;
+
+import ai.onnxruntime.OnnxTensor;
+import ai.onnxruntime.OrtSession;
+import com.visual.open.anpr.core.base.BaseOnnxInfer;
+import com.visual.open.anpr.core.base.PlateRecognition;
+import com.visual.open.anpr.core.domain.ImageMat;
+import com.visual.open.anpr.core.domain.PlateInfo.ParseInfo;
+import com.visual.open.anpr.core.utils.ArrayUtil;
+import com.visual.open.anpr.core.utils.MatUtil;
+import com.visual.open.anpr.core.utils.ReleaseUtil;
+import com.visual.open.anpr.core.utils.SoftMaxUtil;
+import org.opencv.core.Mat;
+import org.opencv.core.Rect;
+import org.opencv.core.Scalar;
+import org.opencv.core.Size;
+import org.opencv.imgproc.Imgproc;
+import java.util.Collections;
+import java.util.Map;
+
+public class TorchPlateRecognition  extends BaseOnnxInfer implements PlateRecognition {
+    private static float STD_VALUE = 0.193f;
+    private static float MEAN_VALUE = 0.588f;
+    private final static String[] PLATE_COLOR = new String[]{"黑色", "蓝色", "绿色", "白色", "黄色"};
+    private final static String PLATE_NAME=
+            "#京沪津渝冀晋蒙辽吉黑苏浙皖闽赣鲁豫鄂湘粤桂琼川贵云藏陕甘青宁新" +
+            "学警港澳挂使领民航危0123456789ABCDEFGHJKLMNPQRSTUVWXYZ险品";
+
+    public TorchPlateRecognition(String modelPath, int threads) {
+        super(modelPath, threads);
+    }
+
+    @Override
+    public ParseInfo inference(ImageMat image, boolean single, Map<String, Object> params) {
+        OnnxTensor tensor = null;
+        OrtSession.Result output = null;
+        ImageMat imageMat = null;
+        try {
+            //图片处理
+            imageMat = single ? image.clone() : splitAndMergePlate(image);
+            //转换数据为张量
+            tensor = imageMat.resizeAndNoReleaseMat(168, 48)
+                .blobFromImageAndDoReleaseMat(1.0/255, new Scalar(MEAN_VALUE, MEAN_VALUE, MEAN_VALUE), true)
+                .to4dFloatOnnxTensorAndNoReleaseMat(new float[]{STD_VALUE, STD_VALUE, STD_VALUE}, true);
+            //ONNX推理
+            output = getSession().run(Collections.singletonMap(getInputName(), tensor));
+            //车牌识别
+            float[][][] result = (float[][][]) output.get(0).getValue();
+            String plateNo = decodePlate(maxScoreIndex(result[0]));
+            //车牌颜色识别
+            float[][] color = (float[][]) output.get(1).getValue();
+            double[] colorSoftMax = SoftMaxUtil.softMax(ArrayUtil.floatToDouble(color[0]));
+            Double[] colorRResult = decodeColor(colorSoftMax);
+            //返回解析到的数据
+            return ParseInfo.build(plateNo, PLATE_COLOR[colorRResult[0].intValue()], colorRResult[1].floatValue());
+        }catch (Exception e){
+            //抛出异常
+            throw new RuntimeException(e);
+        }finally {
+            //释放资源
+            if(null != tensor){
+                ReleaseUtil.release(tensor);
+            }
+            if(null != output){
+                ReleaseUtil.release(output);
+            }
+            if(null != imageMat){
+                ReleaseUtil.release(imageMat);
+            }
+        }
+    }
+
+    private static ImageMat splitAndMergePlate(ImageMat plate){
+        Mat upperImageTemp = null, upperImage = null;
+        Mat lowerImageTemp = null, lowerImage = null;
+        Mat upperReSize = null;
+        try {
+            int width = plate.getWidth();
+            int height = plate.getHeight();
+            //上半部分
+            int upperSplit = Double.valueOf(5.0f / 12 * height).intValue();
+            Rect upperRect = new Rect(0,0, width, upperSplit);
+            upperImageTemp = new Mat(plate.toCvMat(), upperRect);
+            upperImage = new Mat();
+            upperImageTemp.copyTo(upperImage);
+            //下半部分
+            int lowerSplit = Double.valueOf(1.0f / 3 * height).intValue();
+            Rect lowerRect = new Rect(0, lowerSplit, width, height-lowerSplit);
+            lowerImageTemp = new Mat(plate.toCvMat(), lowerRect);
+            lowerImage = new Mat();
+            lowerImageTemp.copyTo(lowerImage);
+            //合并图像
+            upperReSize = new Mat();
+            int heightSize = height-lowerSplit;
+            Imgproc.resize(upperImage, upperReSize, new Size(width, heightSize), 0, 0, Imgproc.INTER_AREA);
+            Mat concatMat = MatUtil.concat(upperReSize, lowerImage);
+            //返回合并后的图像
+            return ImageMat.fromCVMat(concatMat);
+        }finally {
+            ReleaseUtil.release(upperImageTemp, upperImage);
+            ReleaseUtil.release(lowerImageTemp, lowerImage);
+            ReleaseUtil.release(upperReSize);
+        }
+    }
+
+    private static int[] maxScoreIndex(float[][] result){
+        int[] indexes = new int[result.length];
+        for (int i = 0; i < result.length; i++){
+            int index = -1;
+            float max = Float.MIN_VALUE;
+            for (int j = 0; j < result[i].length; j++) {
+                if (max < result[i][j]){
+                    max = result[i][j];
+                    index = j;
+                }
+            }
+            indexes[i] = index;
+        }
+        return indexes;
+    }
+
+    private static String decodePlate(int[] indexes){
+        int pre = 0;
+        StringBuffer sb = new StringBuffer();
+        for(int index : indexes){
+            if(index != 0 && pre != index){
+                sb.append(PLATE_NAME.charAt(index));
+            }
+            pre = index;
+        }
+        return sb.toString();
+    }
+
+    private static Double[] decodeColor(double[] indexes){
+        double index = -1;
+        double max = Double.MIN_VALUE;
+        for (int i = 0; i < indexes.length; i++) {
+            if (max < indexes[i]){
+                max = indexes[i];
+                index = i;
+            }
+        }
+        return new Double[]{index, max};
+    }
+}
--- a/open-anpr-core/src/test/java/com/visual/open/anpr/core/models/TorchPlateRecognitionTest.java
+++ b/open-anpr-core/src/test/java/com/visual/open/anpr/core/models/TorchPlateRecognitionTest.java
@ -0,0 +1,29 @@
+package com.visual.open.anpr.core.models;
+
+import com.visual.open.anpr.core.domain.ImageMat;
+import com.visual.open.anpr.core.domain.PlateInfo;
+import com.visual.open.anpr.core.utils.CropUtil;
+import org.opencv.core.Mat;
+
+import java.util.HashMap;
+import java.util.List;
+
+public class TorchPlateRecognitionTest {
+
+    public static void main(String[] args) {
+        TorchPlateDetection torchPlateDetection = new TorchPlateDetection("open-anpr-core/src/main/resources/models/plate_detect.onnx", 1);
+        TorchPlateRecognition torchPlateRecognition = new TorchPlateRecognition("open-anpr-core/src/main/resources/models/plate_rec_color.onnx", 1);
+
+//        String imagePath = "/Users/diven/workspace/pycharm/github/Chinese_license_plate_detection_recognition/imgs3/double_yellow.jpg";
+        String imagePath = "/Users/diven/workspace/pycharm/github/Chinese_license_plate_detection_recognition/imgs2/image007.jpg";
+        ImageMat imageMat = ImageMat.fromImage(imagePath);
+        List<PlateInfo> plateInfos = torchPlateDetection.inference(imageMat, 0.3f,0.5f, new HashMap<>());
+        System.out.println(plateInfos);
+        for(PlateInfo plateInfo : plateInfos){
+            Mat crop = CropUtil.crop(imageMat.toCvMat(), plateInfo.box);
+//            ImageMat.fromCVMat(crop).imShow();
+            PlateInfo.ParseInfo parseInfo = torchPlateRecognition.inference(ImageMat.fromCVMat(crop), plateInfo.single, new HashMap<>());
+            System.out.println(parseInfo);
+        }
+    }
+}