From 0ec7da30dfd75f7b8da8f27755621c7a1d73bafc Mon Sep 17 00:00:00 2001 From: diven Date: Tue, 17 Jan 2023 15:05:30 +0800 Subject: [PATCH] init --- .../open/anpr/core/base/PlateRecognition.java | 10 ++ .../open/anpr/core/domain/PlateInfo.java | 17 +- .../core/models/TorchPlateRecognition.java | 145 ++++++++++++++++++ .../models/TorchPlateRecognitionTest.java | 29 ++++ 4 files changed, 200 insertions(+), 1 deletion(-) create mode 100644 open-anpr-core/src/main/java/com/visual/open/anpr/core/base/PlateRecognition.java create mode 100644 open-anpr-core/src/main/java/com/visual/open/anpr/core/models/TorchPlateRecognition.java create mode 100644 open-anpr-core/src/test/java/com/visual/open/anpr/core/models/TorchPlateRecognitionTest.java diff --git 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 new file mode 100644 index 0000000..729e1ed --- /dev/null +++ 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 params); +} diff --git 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 index 6be6bfa..fa3d7ba 100755 --- 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, 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, 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); + } + } } diff --git 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 new file mode 100644 index 0000000..346bc8f --- /dev/null +++ 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 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}; + } +} diff --git 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 new file mode 100644 index 0000000..440ac4a --- /dev/null +++ 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 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); + } + } +}