ZF/SmartJavaAI
1
0
Fork 0
mirror of https://github.com/geekwenjie/SmartJavaAI.git synced 2026-09-13 13:18:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

1、集成车牌识别模型，支持车牌检测与识别

Browse Source 
2、新增 Milvus 身份验证支持
3、目标检测功能升级：可指定类别及topk
4、支持自定义线程池线程数量
This commit is contained in:
dengwenjie
2025-07-28 12:04:02 +08:00
parent 1bd74d1bb8
commit 1d45bc597d
117 changed files with 3490 additions and 437 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
smartjavaai-ocr/pom.xml
View File

@@ -6,7 +6,7 @@
<parent>
<groupId>cn.smartjavaai</groupId>
<artifactId>smartjavaai-parent</artifactId>
<version>1.0.20</version>
<version>1.0.22</version>
</parent>
<artifactId>smartjavaai-ocr</artifactId>
@@ -42,7 +42,7 @@
</dependency>
</dependencies>
<version>1.0.20</version>
<version>1.0.22</version>
<name>smartjavaai-ocr</name>
<description>SmartJavaAI</description>
<url>https://github.com/geekwenjie/SmartJavaAI</url>

41
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/PlateDetModelConfig.java Normal file
View File

@@ -0,0 +1,41 @@
package cn.smartjavaai.ocr.config;
import cn.smartjavaai.common.config.ModelConfig;
import cn.smartjavaai.ocr.enums.CommonDetModelEnum;
import cn.smartjavaai.ocr.enums.PlateDetModelEnum;
import lombok.Data;
/**
* 车牌检测模型配置
* @author dwj
*/
@Data
public class PlateDetModelConfig extends ModelConfig {
/**
* 模型
*/
private PlateDetModelEnum modelEnum;
/**
* 检测模型路径
*/
private String modelPath;
/**
* 置信度阈值
*/
private float confidenceThreshold;
/**
* iou阈值
*/
private float iouThreshold;
/**
* 检测结果数量
*/
private int topK;
}

47
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/PlateRecModelConfig.java Normal file
View File

@@ -0,0 +1,47 @@
package cn.smartjavaai.ocr.config;
import cn.smartjavaai.common.config.ModelConfig;
import cn.smartjavaai.ocr.enums.PlateDetModelEnum;
import cn.smartjavaai.ocr.enums.PlateRecModelEnum;
import cn.smartjavaai.ocr.model.plate.PlateDetModel;
import lombok.Data;
/**
* 车牌识别模型配置
* @author dwj
*/
@Data
public class PlateRecModelConfig extends ModelConfig {
/**
* 模型
*/
private PlateRecModelEnum modelEnum;
/**
* 检测模型路径
*/
private String modelPath;
/**
* 车牌检测模型
*/
private PlateDetModel plateDetModel;
/**
* 置信度阈值
*/
private float confidenceThreshold;
/**
* iou阈值
*/
private float iouThreshold;
/**
* 检测结果数量
*/
private int topK;
}

45
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/PlateInfo.java Normal file
View File

@@ -0,0 +1,45 @@
package cn.smartjavaai.ocr.entity;
import cn.smartjavaai.common.entity.DetectionRectangle;
import cn.smartjavaai.ocr.enums.PlateType;
import lombok.Data;
/**
* 车牌识别信息
* @author dwj
*/
@Data
public class PlateInfo {
/**
* 车牌类型
*/
private PlateType plateType;
/**
* 车牌号码
*/
private String plateNumber;
/**
* 车牌颜色
*/
private String plateColor;
/**
* 检测位置信息
*/
private DetectionRectangle detectionRectangle;
/**
* 车牌4角坐标
*/
private OcrBox box;
/**
* 检测得分
*/
private float score;
}

35
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/PlateResult.java Normal file
View File

@@ -0,0 +1,35 @@
package cn.smartjavaai.ocr.entity;
import lombok.Data;
/**
* @author dwj
*/
@Data
public class PlateResult {
/**
* 车牌号码
*/
private String plateNo;
/**
* 车牌颜色
*/
private String plateColor;
public PlateResult(String plateNo, String plateColor) {
this.plateNo = plateNo;
this.plateColor = plateColor;
}
@Override
public String toString() {
return "PlateResult{" +
"plateNo='" + plateNo + '\'' +
", plateColor='" + plateColor + '\'' +
'}';
}
}

28
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/PlateDetModelEnum.java Normal file
View File

@@ -0,0 +1,28 @@
package cn.smartjavaai.ocr.enums;
/**
* 车牌检测模型枚举
* @author dwj
*/
public enum PlateDetModelEnum {
YOLOV5,
YOLOV7;
/**
* 根据名称获取枚举 (忽略大小写和下划线变体)
*/
public static PlateDetModelEnum fromName(String name) {
String formatted = name.trim().toUpperCase().replaceAll("[-_]", "");
for (PlateDetModelEnum model : values()) {
if (model.name().replaceAll("_", "").equals(formatted)) {
return model;
}
}
throw new IllegalArgumentException("未知模型名称: " + name);
}
}

26
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/PlateRecModelEnum.java Normal file
View File

@@ -0,0 +1,26 @@
package cn.smartjavaai.ocr.enums;
/**
* 车牌识别模型枚举
* @author dwj
*/
public enum PlateRecModelEnum {
PLATE_REC_CRNN;
/**
* 根据名称获取枚举 (忽略大小写和下划线变体)
*/
public static PlateRecModelEnum fromName(String name) {
String formatted = name.trim().toUpperCase().replaceAll("[-_]", "");
for (PlateRecModelEnum model : values()) {
if (model.name().replaceAll("_", "").equals(formatted)) {
return model;
}
}
throw new IllegalArgumentException("未知模型名称: " + name);
}
}

45
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/PlateType.java Normal file
View File

@@ -0,0 +1,45 @@
package cn.smartjavaai.ocr.enums;
/**
* @author dwj
*/
public enum PlateType {
SINGLE("single", "单层"),
DOUBLE("double", "双层"),
UNKNOWN("unknown", "未知");
private final String className;
private final String description;
PlateType(String className, String description) {
this.className = className;
this.description = description;
}
public String getClassName() {
return className;
}
public String getDescription() {
return description;
}
/**
* 根据value获取对应的PlateType
* @param className
* @return PlateType
*/
public static PlateType fromClassName(String className) {
for (PlateType type : values()) {
if (type.className.equals(className)) {
return type;
}
}
return null;
}
}

168
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/factory/PlateModelFactory.java Normal file
View File

@@ -0,0 +1,168 @@
package cn.smartjavaai.ocr.factory;
import cn.smartjavaai.common.config.Config;
import cn.smartjavaai.ocr.config.PlateDetModelConfig;
import cn.smartjavaai.ocr.config.PlateRecModelConfig;
import cn.smartjavaai.ocr.config.TableStructureConfig;
import cn.smartjavaai.ocr.enums.PlateDetModelEnum;
import cn.smartjavaai.ocr.enums.PlateRecModelEnum;
import cn.smartjavaai.ocr.enums.TableStructureModelEnum;
import cn.smartjavaai.ocr.exception.OcrException;
import cn.smartjavaai.ocr.model.plate.CRNNPlateRecModel;
import cn.smartjavaai.ocr.model.plate.PlateDetModel;
import cn.smartjavaai.ocr.model.plate.PlateRecModel;
import cn.smartjavaai.ocr.model.plate.Yolov5PlateDetModel;
import cn.smartjavaai.ocr.model.table.CommonTableStructureModel;
import cn.smartjavaai.ocr.model.table.TableStructureModel;
import lombok.extern.slf4j.Slf4j;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
/**
* 车牌识别模型工厂
* @author dwj
*/
@Slf4j
public class PlateModelFactory {
// 使用 volatile 和双重检查锁定来确保线程安全的单例模式
private static volatile PlateModelFactory instance;
/**
* 模型缓存
*/
private static final ConcurrentHashMap<PlateDetModelEnum, PlateDetModel> detModelMap = new ConcurrentHashMap<>();
/**
* 模型缓存
*/
private static final ConcurrentHashMap<PlateRecModelEnum, PlateRecModel> recModelMap = new ConcurrentHashMap<>();
/**
* 模型注册表
*/
private static final Map<PlateDetModelEnum, Class<? extends PlateDetModel>> detModelRegistry =
new ConcurrentHashMap<>();
/**
* 模型注册表
*/
private static final Map<PlateRecModelEnum, Class<? extends PlateRecModel>> recModelRegistry =
new ConcurrentHashMap<>();
public static PlateModelFactory getInstance() {
if (instance == null) {
synchronized (PlateModelFactory.class) {
if (instance == null) {
instance = new PlateModelFactory();
}
}
}
return instance;
}
/**
* 注册模型
* @param plateDetModelEnum
* @param clazz
*/
private static void registerDetModel(PlateDetModelEnum plateDetModelEnum, Class<? extends PlateDetModel> clazz) {
detModelRegistry.put(plateDetModelEnum, clazz);
}
/**
* 注册模型
* @param plateRecModelEnum
* @param clazz
*/
private static void registerRecModel(PlateRecModelEnum plateRecModelEnum, Class<? extends PlateRecModel> clazz) {
recModelRegistry.put(plateRecModelEnum, clazz);
}
/**
* 获取模型
* @param config
* @return
*/
public PlateDetModel getDetModel(PlateDetModelConfig config) {
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
throw new OcrException("未配置OCR模型");
}
return detModelMap.computeIfAbsent(config.getModelEnum(), k -> {
return createDetModel(config);
});
}
/**
* 获取模型
* @param config
* @return
*/
public PlateRecModel getRecModel(PlateRecModelConfig config) {
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
throw new OcrException("未配置OCR模型");
}
return recModelMap.computeIfAbsent(config.getModelEnum(), k -> {
return createRecModel(config);
});
}
/**
* 创建检测模型
* @param config
* @return
*/
private PlateDetModel createDetModel(PlateDetModelConfig config) {
Class<?> clazz = detModelRegistry.get(config.getModelEnum());
if(clazz == null){
throw new OcrException("Unsupported model");
}
PlateDetModel model = null;
try {
model = (PlateDetModel) clazz.newInstance();
} catch (InstantiationException | IllegalAccessException e) {
throw new OcrException(e);
}
model.loadModel(config);
return model;
}
/**
* 创建识别模型
* @param config
* @return
*/
private PlateRecModel createRecModel(PlateRecModelConfig config) {
Class<?> clazz = recModelRegistry.get(config.getModelEnum());
if(clazz == null){
throw new OcrException("Unsupported model");
}
PlateRecModel model = null;
try {
model = (PlateRecModel) clazz.newInstance();
} catch (InstantiationException | IllegalAccessException e) {
throw new OcrException(e);
}
model.loadModel(config);
return model;
}
// 初始化默认算法
static {
registerDetModel(PlateDetModelEnum.YOLOV5, Yolov5PlateDetModel.class);
registerDetModel(PlateDetModelEnum.YOLOV7, Yolov5PlateDetModel.class);
registerRecModel(PlateRecModelEnum.PLATE_REC_CRNN, CRNNPlateRecModel.class);
log.debug("缓存目录：{}", Config.getCachePath());
}
}

7
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/OcrCommonDetModel.java
View File

@@ -1,8 +1,11 @@
package cn.smartjavaai.ocr.model.common.detect;
import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.ndarray.NDList;
import cn.smartjavaai.ocr.config.OcrDetModelConfig;
import cn.smartjavaai.ocr.entity.OcrBox;
import org.apache.commons.pool2.impl.GenericObjectPool;
import java.awt.image.BufferedImage;
import java.util.List;
@@ -93,5 +96,9 @@ public interface OcrCommonDetModel extends AutoCloseable{
throw new UnsupportedOperationException("默认不支持该功能");
}
default GenericObjectPool<Predictor<Image, NDList>> getPool(){
throw new UnsupportedOperationException("默认不支持该功能");
}
}

13
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/OcrCommonDetModelImpl.java
View File

@@ -43,7 +43,7 @@ import java.util.*;
@Slf4j
public class OcrCommonDetModelImpl implements OcrCommonDetModel{
private ObjectPool<Predictor<Image, NDList>> detPredictorPool;
private GenericObjectPool<Predictor<Image, NDList>> detPredictorPool;
private ZooModel<Image, NDList> detectionModel;
@@ -61,8 +61,14 @@ public class OcrCommonDetModelImpl implements OcrCommonDetModel{
detectionModel = ModelZoo.loadModel(detCriteria);
// 创建池子：每个线程独享 Predictor
this.detPredictorPool = new GenericObjectPool<>(new PredictorFactory<>(detectionModel));
int predictorPoolSize = config.getPredictorPoolSize();
if(config.getPredictorPoolSize() <= 0){
predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
}
detPredictorPool.setMaxTotal(predictorPoolSize);
log.debug("当前设备: " + detectionModel.getNDManager().getDevice());
log.debug("当前引擎: " + Engine.getInstance().getEngineName());
log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
throw new OcrException("检测模型加载失败", e);
}
@@ -209,6 +215,11 @@ public class OcrCommonDetModelImpl implements OcrCommonDetModel{
}
}
@Override
public GenericObjectPool<Predictor<Image, NDList>> getPool() {
return detPredictorPool;
}
@Override
public void close() throws Exception {
try {

2
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/criteria/OcrCommonDetCriterialFactory.java
View File

@@ -25,7 +25,7 @@ public class OcrCommonDetCriterialFactory {
public static Criteria<Image, NDList> createCriteria(OcrDetModelConfig config) {
Device device = null;
if(!Objects.isNull(config.getDevice())){
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu();
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
}
Criteria<Image, NDList> criteria = null;
ConcurrentHashMap params = new ConcurrentHashMap<String, String>();

5
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/OcrDirectionModel.java
View File

@@ -9,6 +9,7 @@ import cn.smartjavaai.ocr.entity.OcrBox;
import cn.smartjavaai.ocr.entity.OcrInfo;
import cn.smartjavaai.ocr.entity.OcrItem;
import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.opencv.core.Mat;
import java.awt.image.BufferedImage;
@@ -106,4 +107,8 @@ public interface OcrDirectionModel extends AutoCloseable{
throw new UnsupportedOperationException("默认不支持该功能");
}
default GenericObjectPool<Predictor<Image, DirectionInfo>> getPool() {
throw new UnsupportedOperationException("默认不支持该功能");
}
}

17
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/PPOCRMobileV2ClsModel.java
View File

@@ -55,7 +55,7 @@ import java.util.concurrent.ConcurrentHashMap;
public class PPOCRMobileV2ClsModel implements OcrDirectionModel {
private ObjectPool<Predictor<Image, DirectionInfo>> predictorPool;
private GenericObjectPool<Predictor<Image, DirectionInfo>> predictorPool;
private DirectionModelConfig config;
@@ -71,10 +71,6 @@ public class PPOCRMobileV2ClsModel implements OcrDirectionModel {
}
this.config = config;
this.textDetModel = config.getTextDetModel();
Device device = null;
if(!Objects.isNull(config.getDevice())){
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu();
}
ConcurrentHashMap params = new ConcurrentHashMap<String, String>();
if(StringUtils.isNotBlank(config.getBatchifier())){
params.put("batchifier", config.getBatchifier());
@@ -84,8 +80,14 @@ public class PPOCRMobileV2ClsModel implements OcrDirectionModel {
model = ModelZoo.loadModel(criteria);
// 创建池子：每个线程独享 Predictor
this.predictorPool = new GenericObjectPool<>(new PredictorFactory<>(model));
int predictorPoolSize = config.getPredictorPoolSize();
if(config.getPredictorPoolSize() <= 0){
predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
}
predictorPool.setMaxTotal(predictorPoolSize);
log.debug("当前设备: " + model.getNDManager().getDevice());
log.debug("当前引擎: " + Engine.getInstance().getEngineName());
log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
throw new OcrException("模型加载失败", e);
}
@@ -369,6 +371,11 @@ public class PPOCRMobileV2ClsModel implements OcrDirectionModel {
return textDetModel;
}
@Override
public GenericObjectPool<Predictor<Image, DirectionInfo>> getPool() {
return predictorPool;
}
@Override
public void close() throws Exception {
try {

2
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/criteria/DirectionCriteriaFactory.java
View File

@@ -28,7 +28,7 @@ public class DirectionCriteriaFactory {
public static Criteria<Image, DirectionInfo> createCriteria(DirectionModelConfig config) {
Device device = null;
if(!Objects.isNull(config.getDevice())){
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu();
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
}
Criteria<Image, DirectionInfo> criteria = null;
ConcurrentHashMap params = new ConcurrentHashMap<String, String>();

6
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/OcrCommonRecModel.java
View File

@@ -1,5 +1,6 @@
package cn.smartjavaai.ocr.model.common.recognize;
import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import cn.smartjavaai.ocr.config.OcrDetModelConfig;
import cn.smartjavaai.ocr.config.OcrRecModelConfig;
@@ -8,6 +9,7 @@ import cn.smartjavaai.ocr.entity.OcrBox;
import cn.smartjavaai.ocr.entity.OcrInfo;
import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel;
import cn.smartjavaai.ocr.model.common.direction.OcrDirectionModel;
import org.apache.commons.pool2.impl.GenericObjectPool;
import java.awt.image.BufferedImage;
import java.util.List;
@@ -106,4 +108,8 @@ public interface OcrCommonRecModel extends AutoCloseable{
throw new UnsupportedOperationException("默认不支持该功能");
}
default GenericObjectPool<Predictor<Image, String>> getPool() {
throw new UnsupportedOperationException("默认不支持该功能");
}
}

16
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/OcrCommonRecModelImpl.java
View File

@@ -24,7 +24,6 @@ import cn.smartjavaai.ocr.exception.OcrException;
import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel;
import cn.smartjavaai.ocr.model.common.direction.OcrDirectionModel;
import cn.smartjavaai.ocr.model.common.recognize.criteria.OcrCommonRecCriterialFactory;
import cn.smartjavaai.ocr.opencv.OcrOpenCVUtils;
import cn.smartjavaai.ocr.utils.OcrUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
@@ -49,7 +48,7 @@ import java.util.stream.Collectors;
@Slf4j
public class OcrCommonRecModelImpl implements OcrCommonRecModel {
private ObjectPool<Predictor<Image, String>> recPredictorPool;
private GenericObjectPool<Predictor<Image, String>> recPredictorPool;
private OcrRecModelConfig config;
@@ -72,8 +71,14 @@ public class OcrCommonRecModelImpl implements OcrCommonRecModel {
try{
recognitionModel = ModelZoo.loadModel(recCriteria);
this.recPredictorPool = new GenericObjectPool<>(new PredictorFactory<>(recognitionModel));
int predictorPoolSize = config.getPredictorPoolSize();
if(config.getPredictorPoolSize() <= 0){
predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
}
recPredictorPool.setMaxTotal(predictorPoolSize);
log.debug("当前设备: " + recognitionModel.getNDManager().getDevice());
log.debug("当前引擎: " + Engine.getInstance().getEngineName());
log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
throw new OcrException("识别模型加载失败", e);
}
@@ -238,7 +243,7 @@ public class OcrCommonRecModelImpl implements OcrCommonRecModel {
throw new OcrException("未检测到文字");
}
Mat wrappedImage = (Mat) img.getWrappedImage();
BufferedImage bufferedImage = OcrOpenCVUtils.mat2Image(wrappedImage);
BufferedImage bufferedImage = OpenCVUtils.mat2Image(wrappedImage);
OcrUtils.drawRectWithText(bufferedImage, ocrInfo, fontSize);
ImageUtils.saveImage(bufferedImage, outputPath);
wrappedImage.release();
@@ -443,6 +448,11 @@ public class OcrCommonRecModelImpl implements OcrCommonRecModel {
return directionModel;
}
public GenericObjectPool<Predictor<Image, String>> getRecPredictorPool() {
return recPredictorPool;
}
@Override
public void close() throws Exception {
try {

2
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/criteria/OcrCommonRecCriterialFactory.java
View File

@@ -24,7 +24,7 @@ public class OcrCommonRecCriterialFactory {
public static Criteria<Image, String> createCriteria(OcrRecModelConfig config) {
Device device = null;
if(!Objects.isNull(config.getDevice())){
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu();
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
}
Criteria<Image, String> criteria = null;
ConcurrentHashMap params = new ConcurrentHashMap<String, String>();

342
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/CRNNPlateRecModel.java Normal file
View File

@@ -0,0 +1,342 @@
package cn.smartjavaai.ocr.model.plate;
import ai.djl.MalformedModelException;
import ai.djl.engine.Engine;
import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.output.BoundingBox;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelNotFoundException;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;
import cn.hutool.core.lang.UUID;
import cn.hutool.core.lang.generator.UUIDGenerator;
import cn.smartjavaai.common.entity.DetectionRectangle;
import cn.smartjavaai.common.entity.R;
import cn.smartjavaai.common.pool.PredictorFactory;
import cn.smartjavaai.common.utils.Base64ImageUtils;
import cn.smartjavaai.common.utils.FileUtils;
import cn.smartjavaai.common.utils.ImageUtils;
import cn.smartjavaai.common.utils.OpenCVUtils;
import cn.smartjavaai.ocr.config.PlateDetModelConfig;
import cn.smartjavaai.ocr.config.PlateRecModelConfig;
import cn.smartjavaai.ocr.entity.PlateInfo;
import cn.smartjavaai.ocr.entity.PlateResult;
import cn.smartjavaai.ocr.enums.PlateType;
import cn.smartjavaai.ocr.exception.OcrException;
import cn.smartjavaai.ocr.model.plate.criteria.PlateDetCriterialFactory;
import cn.smartjavaai.ocr.model.plate.criteria.PlateRecCriterialFactory;
import cn.smartjavaai.ocr.utils.OcrUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.pool2.ObjectPool;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.opencv.core.Core;
import org.opencv.core.Mat;
import org.opencv.core.Rect;
import org.opencv.core.Size;
import org.opencv.imgproc.Imgproc;
import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.*;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Objects;
/**
* @author dwj
*/
@Slf4j
public class CRNNPlateRecModel implements PlateRecModel{
private GenericObjectPool<Predictor<Image, PlateResult>> recPredictorPool;
private ZooModel<Image, PlateResult> recModel;
private PlateRecModelConfig config;
@Override
public void loadModel(PlateRecModelConfig config) {
if(StringUtils.isBlank(config.getModelPath())){
throw new OcrException("modelPath is null");
}
this.config = config;
//初始化 检测Criteria
Criteria<Image, PlateResult> detCriteria = PlateRecCriterialFactory.createCriteria(config);
try{
recModel = ModelZoo.loadModel(detCriteria);
// 创建池子：每个线程独享 Predictor
this.recPredictorPool = new GenericObjectPool<>(new PredictorFactory<>(recModel));
int predictorPoolSize = config.getPredictorPoolSize();
if(config.getPredictorPoolSize() <= 0){
predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
}
recPredictorPool.setMaxTotal(predictorPoolSize);
log.debug("当前设备: " + recModel.getNDManager().getDevice());
log.debug("当前引擎: " + Engine.getInstance().getEngineName());
log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
throw new OcrException("检测模型加载失败", e);
}
}
@Override
public R<List<PlateInfo>> recognize(String imagePath) {
if(!FileUtils.isFileExists(imagePath)){
return R.fail(R.Status.FILE_NOT_FOUND);
}
Image img = null;
try {
img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
R<List<PlateInfo>> plateResult = recognize(img);
return plateResult;
} catch (IOException e) {
throw new OcrException("无效的图片", e);
} finally {
((Mat)img.getWrappedImage()).release();
}
}
@Override
public R<List<PlateInfo>> recognizeBase64(String base64Image) {
if(StringUtils.isBlank(base64Image)){
return R.fail(R.Status.INVALID_IMAGE);
}
byte[] imageData = Base64ImageUtils.base64ToImage(base64Image);
return recognize(imageData);
}
@Override
public R<List<PlateInfo>> recognize(BufferedImage image) {
if(!ImageUtils.isImageValid(image)){
return R.fail(R.Status.INVALID_IMAGE);
}
Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image));
R<List<PlateInfo>> plateResult = recognize(img);
((Mat)img.getWrappedImage()).release();
return plateResult;
}
@Override
public R<List<PlateInfo>> recognize(byte[] imageData) {
if(Objects.isNull(imageData)){
return R.fail(R.Status.INVALID_IMAGE);
}
return recognize(new ByteArrayInputStream(imageData));
}
@Override
public R<List<PlateInfo>> recognize(Image image) {
if(Objects.isNull(config.getPlateDetModel())){
return R.fail(R.Status.PARAM_ERROR.getCode(), "未指定车牌检测模型");
}
DetectedObjects detectedObjects = config.getPlateDetModel().detect(image);
if(Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
return R.fail(R.Status.NO_OBJECT_DETECTED);
}
List<PlateInfo> plateInfoList = OcrUtils.convertToPlateInfo(detectedObjects, image);
Predictor<Image, PlateResult> predictor = null;
try {
predictor = recPredictorPool.borrowObject();
for (PlateInfo plateInfo : plateInfoList){
DetectionRectangle detectionRectangle = plateInfo.getDetectionRectangle();
// Image subImage = image.getSubImage(detectionRectangle.getX(), detectionRectangle.getY(), detectionRectangle.getWidth(), detectionRectangle.getHeight());
//透视变换
Image subImage = OcrUtils.transformAndCrop((Mat)image.getWrappedImage(), plateInfo.getBox());
//双层车牌
if(plateInfo.getPlateType() == PlateType.DOUBLE){
Mat mergeImage = getSplitMerge((Mat)subImage.getWrappedImage());
subImage = ImageFactory.getInstance().fromImage(mergeImage);
}
PlateResult plateResult = predictor.predict(subImage);
if(Objects.nonNull(plateResult)){
plateInfo.setPlateNumber(plateResult.getPlateNo());
plateInfo.setPlateColor(plateResult.getPlateColor());
}
}
return R.ok(plateInfoList);
} catch (Exception e) {
throw new OcrException("车牌识别错误", e);
}finally {
if (predictor != null) {
try {
recPredictorPool.returnObject(predictor); //归还
} catch (Exception e) {
log.warn("归还Predictor失败", e);
try {
predictor.close(); // 归还失败才销毁
} catch (Exception ex) {
log.error("关闭Predictor失败", ex);
}
}
}
}
}
/**
* 双层车牌进行分割后识别
* @param img
* @return
*/
private Mat getSplitMerge(Mat img) {
int h = img.rows();
int w = img.cols();
// 上半部分：高度的前 5/12
Rect upperRect = new Rect(0, 0, w, (int)(5.0 / 12 * h));
Mat imgUpper = new Mat(img, upperRect);
// 下半部分：高度从 1/3 开始
Rect lowerRect = new Rect(0, (int)(1.0 / 3 * h), w, h - (int)(1.0 / 3 * h));
Mat imgLower = new Mat(img, lowerRect);
// 将上半部分 resize 到与下半部分相同大小
Mat resizedUpper = new Mat();
Size lowerSize = imgLower.size();
Imgproc.resize(imgUpper, resizedUpper, lowerSize);
// 水平拼接（将上下拼成左右）
List<Mat> mergeList = new ArrayList<>();
mergeList.add(resizedUpper);
mergeList.add(imgLower);
Mat merged = new Mat();
Core.hconcat(mergeList, merged);
return merged;
}
@Override
public PlateResult recognizeCropped(Image image) {
Predictor<Image, PlateResult> predictor = null;
try {
predictor = recPredictorPool.borrowObject();
return predictor.predict(image);
} catch (Exception e) {
throw new OcrException("车牌检测错误", e);
}finally {
if (predictor != null) {
try {
recPredictorPool.returnObject(predictor); //归还
} catch (Exception e) {
log.warn("归还Predictor失败", e);
try {
predictor.close(); // 归还失败才销毁
} catch (Exception ex) {
log.error("关闭Predictor失败", ex);
}
}
}
}
}
@Override
public R<List<PlateInfo>> recognize(InputStream inputStream) {
if(Objects.isNull(inputStream)){
return R.fail(R.Status.INVALID_IMAGE);
}
Image img = null;
try {
img = ImageFactory.getInstance().fromInputStream(inputStream);
return recognize(img);
} catch (IOException e) {
throw new OcrException("无效图片输入流", e);
} finally {
if (img != null){
((Mat)img.getWrappedImage()).release();
}
}
}
@Override
public R<Void> recognizeAndDraw(String imagePath, String outputPath) {
if(!FileUtils.isFileExists(imagePath)){
return R.fail(R.Status.FILE_NOT_FOUND);
}
Image img = null;
try {
img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
R<List<PlateInfo>> plateResult = recognize(img);
if(!plateResult.isSuccess()){
return R.fail(plateResult.getCode(), plateResult.getMessage());
}
if(CollectionUtils.isEmpty(plateResult.getData())){
return R.fail(R.Status.NO_OBJECT_DETECTED);
}
BufferedImage bufferedImage = OpenCVUtils.mat2Image((Mat)img.getWrappedImage());
OcrUtils.drawPlateInfo(bufferedImage, plateResult.getData());
ImageIO.write(bufferedImage, "jpg", new File(outputPath));
return R.ok();
} catch (IOException e) {
throw new OcrException(e);
} finally {
if (img != null){
((Mat)img.getWrappedImage()).release();
}
}
}
@Override
public R<BufferedImage> recognizeAndDraw(BufferedImage sourceImage) {
if(!ImageUtils.isImageValid(sourceImage)){
return R.fail(R.Status.INVALID_IMAGE);
}
Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(sourceImage));
try {
R<List<PlateInfo>> plateResult = recognize(img);
if(!plateResult.isSuccess()){
return R.fail(plateResult.getCode(), plateResult.getMessage());
}
if(CollectionUtils.isEmpty(plateResult.getData())){
return R.fail(R.Status.NO_OBJECT_DETECTED);
}
OcrUtils.drawPlateInfo((Mat)img.getWrappedImage(), plateResult.getData());
ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
// 调用 save 方法将 Image 写入字节流
img.save(outputStream, "png");
// 将字节流转换为 BufferedImage
byte[] imageBytes = outputStream.toByteArray();
return R.ok(ImageIO.read(new ByteArrayInputStream(imageBytes)));
} catch (IOException e) {
throw new OcrException("导出图片失败", e);
} finally {
if (img != null){
((Mat)img.getWrappedImage()).release();
}
}
}
@Override
public GenericObjectPool<Predictor<Image, PlateResult>> getPool() {
return recPredictorPool;
}
@Override
public void close() throws Exception {
try {
if (recPredictorPool != null) {
recPredictorPool.close();
}
} catch (Exception e) {
log.warn("关闭 predictorPool 失败", e);
}
try {
if (recModel != null) {
recModel.close();
}
} catch (Exception e) {
log.warn("关闭 model 失败", e);
}
}
}

115
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/PlateDetModel.java Normal file
View File

@@ -0,0 +1,115 @@
package cn.smartjavaai.ocr.model.plate;
import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.output.DetectedObjects;
import cn.smartjavaai.common.entity.DetectionResponse;
import cn.smartjavaai.ocr.config.OcrDetModelConfig;
import cn.smartjavaai.ocr.config.PlateDetModelConfig;
import cn.smartjavaai.ocr.entity.OcrBox;
import cn.smartjavaai.ocr.entity.PlateInfo;
import cn.smartjavaai.common.entity.R;
import org.apache.commons.pool2.impl.GenericObjectPool;
import java.awt.image.BufferedImage;
import java.io.InputStream;
import java.util.List;
/**
* 车牌检测模型
* @author dwj
*/
public interface PlateDetModel extends AutoCloseable{
/**
* 加载模型
* @param config
*/
void loadModel(PlateDetModelConfig config); // 加载模型
/**
* 车牌检测
* @param imagePath 图片路径
* @return
*/
default R<List<PlateInfo>> detect(String imagePath) {
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 车牌检测
* @param inputStream
* @return
*/
default R<List<PlateInfo>> detect(InputStream inputStream) {
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 车牌检测
* @param base64Image
* @return
*/
default R<List<PlateInfo>> detectBase64(String base64Image){
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 车牌检测
* @param image BufferedImage
* @return
*/
default R<List<PlateInfo>> detect(BufferedImage image) {
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 车牌检测
* @param imageData 图片字节数组
* @return
*/
default R<List<PlateInfo>> detect(byte[] imageData) {
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 车牌检测
* @param image DJL Image
* @return
*/
default DetectedObjects detect(Image image){
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 检测并绘制结果
* @param imagePath 图片输入路径（包含文件名称）
* @param outputPath 图片输出路径（包含文件名称）
*/
default R<Void> detectAndDraw(String imagePath, String outputPath) {
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 检测并绘制结果
* @param sourceImage
* @return
*/
default R<BufferedImage> detectAndDraw(BufferedImage sourceImage){
throw new UnsupportedOperationException("默认不支持该功能");
}
default GenericObjectPool<Predictor<Image, DetectedObjects>> getPool(){
throw new UnsupportedOperationException("默认不支持该功能");
}
}

118
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/PlateRecModel.java Normal file
View File

@@ -0,0 +1,118 @@
package cn.smartjavaai.ocr.model.plate;
import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.output.DetectedObjects;
import cn.smartjavaai.common.entity.R;
import cn.smartjavaai.ocr.config.PlateDetModelConfig;
import cn.smartjavaai.ocr.config.PlateRecModelConfig;
import cn.smartjavaai.ocr.entity.PlateInfo;
import cn.smartjavaai.ocr.entity.PlateResult;
import org.apache.commons.pool2.impl.GenericObjectPool;
import java.awt.image.BufferedImage;
import java.io.InputStream;
import java.util.List;
/**
* 车牌识别模型
* @author dwj
*/
public interface PlateRecModel extends AutoCloseable{
/**
* 加载模型
* @param config
*/
void loadModel(PlateRecModelConfig config); // 加载模型
/**
* 车牌识别
* @param imagePath 图片路径
* @return
*/
default R<List<PlateInfo>> recognize(String imagePath) {
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 车牌识别
* @param inputStream
* @return
*/
default R<List<PlateInfo>> recognize(InputStream inputStream) {
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 车牌识别
* @param base64Image
* @return
*/
default R<List<PlateInfo>> recognizeBase64(String base64Image){
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 车牌识别
* @param image BufferedImage
* @return
*/
default R<List<PlateInfo>> recognize(BufferedImage image) {
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 车牌识别
* @param imageData 图片字节数组
* @return
*/
default R<List<PlateInfo>> recognize(byte[] imageData) {
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 车牌识别
* @param image DJL Image
* @return
*/
default R<List<PlateInfo>> recognize(Image image){
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 识别裁剪后的图片
* @return
*/
default PlateResult recognizeCropped(Image image){
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 检测并绘制结果
* @param imagePath 图片输入路径（包含文件名称）
* @param outputPath 图片输出路径（包含文件名称）
*/
default R<Void> recognizeAndDraw(String imagePath, String outputPath) {
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 检测并绘制结果
* @param sourceImage
* @return
*/
default R<BufferedImage> recognizeAndDraw(BufferedImage sourceImage){
throw new UnsupportedOperationException("默认不支持该功能");
}
default GenericObjectPool<Predictor<Image, PlateResult>> getPool() {
throw new UnsupportedOperationException("默认不支持该功能");
}
}

244
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/Yolov5PlateDetModel.java Normal file
View File

@@ -0,0 +1,244 @@
package cn.smartjavaai.ocr.model.plate;
import ai.djl.MalformedModelException;
import ai.djl.engine.Engine;
import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelNotFoundException;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;
import cn.smartjavaai.common.entity.R;
import cn.smartjavaai.common.pool.PredictorFactory;
import cn.smartjavaai.common.utils.Base64ImageUtils;
import cn.smartjavaai.common.utils.FileUtils;
import cn.smartjavaai.common.utils.ImageUtils;
import cn.smartjavaai.common.utils.OpenCVUtils;
import cn.smartjavaai.ocr.config.OcrDetModelConfig;
import cn.smartjavaai.ocr.config.PlateDetModelConfig;
import cn.smartjavaai.ocr.entity.OcrBox;
import cn.smartjavaai.ocr.entity.PlateInfo;
import cn.smartjavaai.ocr.exception.OcrException;
import cn.smartjavaai.ocr.model.common.detect.criteria.OcrCommonDetCriterialFactory;
import cn.smartjavaai.ocr.model.plate.criteria.PlateDetCriterialFactory;
import cn.smartjavaai.ocr.utils.OcrUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.pool2.ObjectPool;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.opencv.core.Mat;
import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;
import java.util.Objects;
/**
* Yolov5 车牌检测模型
* @author dwj
* @date 2025/7/23
*/
@Slf4j
public class Yolov5PlateDetModel implements PlateDetModel{
private GenericObjectPool<Predictor<Image, DetectedObjects>> detPredictorPool;
private ZooModel<Image, DetectedObjects> detectionModel;
private PlateDetModelConfig config;
@Override
public void loadModel(PlateDetModelConfig config) {
if(StringUtils.isBlank(config.getModelPath())){
throw new OcrException("modelPath is null");
}
this.config = config;
//初始化 检测Criteria
Criteria<Image, DetectedObjects> detCriteria = PlateDetCriterialFactory.createCriteria(config);
try{
detectionModel = ModelZoo.loadModel(detCriteria);
// 创建池子：每个线程独享 Predictor
this.detPredictorPool = new GenericObjectPool<>(new PredictorFactory<>(detectionModel));
int predictorPoolSize = config.getPredictorPoolSize();
if(config.getPredictorPoolSize() <= 0){
predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
}
detPredictorPool.setMaxTotal(predictorPoolSize);
log.debug("当前设备: " + detectionModel.getNDManager().getDevice());
log.debug("当前引擎: " + Engine.getInstance().getEngineName());
log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
throw new OcrException("检测模型加载失败", e);
}
}
@Override
public R<List<PlateInfo>> detect(String imagePath) {
if(!FileUtils.isFileExists(imagePath)){
return R.fail(R.Status.FILE_NOT_FOUND);
}
Image img = null;
try {
img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
} catch (IOException e) {
throw new OcrException("无效的图片", e);
}
DetectedObjects detectedObjects = detect(img);
if (Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
return R.fail(R.Status.NO_OBJECT_DETECTED);
}
List<PlateInfo> plateInfoList = OcrUtils.convertToPlateInfo(detectedObjects, img);
((Mat)img.getWrappedImage()).release();
return R.ok(plateInfoList);
}
@Override
public R<List<PlateInfo>> detectBase64(String base64Image) {
if(StringUtils.isBlank(base64Image)){
return R.fail(R.Status.INVALID_IMAGE);
}
byte[] imageData = Base64ImageUtils.base64ToImage(base64Image);
return detect(imageData);
}
@Override
public R<List<PlateInfo>> detect(BufferedImage image) {
if(!ImageUtils.isImageValid(image)){
return R.fail(R.Status.INVALID_IMAGE);
}
Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image));
DetectedObjects detectedObjects = detect(img);
if (Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
return R.fail(R.Status.NO_OBJECT_DETECTED);
}
List<PlateInfo> plateInfoList = OcrUtils.convertToPlateInfo(detectedObjects, img);
((Mat)img.getWrappedImage()).release();
return R.ok(plateInfoList);
}
@Override
public R<List<PlateInfo>> detect(byte[] imageData) {
if(Objects.isNull(imageData)){
return R.fail(R.Status.INVALID_IMAGE);
}
return detect(new ByteArrayInputStream(imageData));
}
@Override
public DetectedObjects detect(Image image) {
Predictor<Image, DetectedObjects> predictor = null;
try {
predictor = detPredictorPool.borrowObject();
return predictor.predict(image);
} catch (Exception e) {
throw new OcrException("车牌检测错误", e);
}finally {
if (predictor != null) {
try {
detPredictorPool.returnObject(predictor); //归还
} catch (Exception e) {
log.warn("归还Predictor失败", e);
try {
predictor.close(); // 归还失败才销毁
} catch (Exception ex) {
log.error("关闭Predictor失败", ex);
}
}
}
}
}
@Override
public R<List<PlateInfo>> detect(InputStream inputStream) {
if(Objects.isNull(inputStream)){
return R.fail(R.Status.INVALID_IMAGE);
}
try {
Image img = ImageFactory.getInstance().fromInputStream(inputStream);
DetectedObjects detection = detect(img);
List<PlateInfo> plateInfoList = OcrUtils.convertToPlateInfo(detection, img);
((Mat)img.getWrappedImage()).release();
return R.ok(plateInfoList);
} catch (IOException e) {
throw new OcrException("无效图片输入流", e);
}
}
@Override
public R<Void> detectAndDraw(String imagePath, String outputPath) {
if(!FileUtils.isFileExists(imagePath)){
return R.fail(R.Status.FILE_NOT_FOUND);
}
try {
Image img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
DetectedObjects detectedObjects = detect(img);
if(Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
return R.fail(R.Status.NO_FACE_DETECTED);
}
img.drawBoundingBoxes(detectedObjects);
Path output = Paths.get(outputPath);
log.debug("Saving to {}", output.toAbsolutePath().toString());
img.save(Files.newOutputStream(output), "png");
return R.ok();
} catch (IOException e) {
throw new OcrException(e);
}
}
@Override
public R<BufferedImage> detectAndDraw(BufferedImage sourceImage) {
if(!ImageUtils.isImageValid(sourceImage)){
return R.fail(R.Status.INVALID_IMAGE);
}
Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(sourceImage));
DetectedObjects detectedObjects = detect(img);
if(Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
return R.fail(R.Status.NO_FACE_DETECTED);
}
img.drawBoundingBoxes(detectedObjects);
try {
ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
// 调用 save 方法将 Image 写入字节流
img.save(outputStream, "png");
// 将字节流转换为 BufferedImage
byte[] imageBytes = outputStream.toByteArray();
return R.ok(ImageIO.read(new ByteArrayInputStream(imageBytes)));
} catch (IOException e) {
throw new OcrException("导出图片失败", e);
}
}
@Override
public GenericObjectPool<Predictor<Image, DetectedObjects>> getPool() {
return detPredictorPool;
}
@Override
public void close() throws Exception {
try {
if (detPredictorPool != null) {
detPredictorPool.close();
}
} catch (Exception e) {
log.warn("关闭 predictorPool 失败", e);
}
try {
if (detectionModel != null) {
detectionModel.close();
}
} catch (Exception e) {
log.warn("关闭 model 失败", e);
}
}
}

60
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/criteria/PlateDetCriterialFactory.java Normal file
View File

@@ -0,0 +1,60 @@
package cn.smartjavaai.ocr.model.plate.criteria;
import ai.djl.Device;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.repository.zoo.Criteria;
import ai.djl.training.util.ProgressBar;
import cn.smartjavaai.common.enums.DeviceEnum;
import cn.smartjavaai.ocr.config.PlateDetModelConfig;
import cn.smartjavaai.ocr.enums.PlateDetModelEnum;
import cn.smartjavaai.ocr.model.plate.translator.Yolo5PlateDetectTranslator;
import cn.smartjavaai.ocr.model.plate.translator.Yolov7PlateDetectTranslator;
import org.apache.commons.lang3.StringUtils;
import java.nio.file.Paths;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
/**
* @author dwj
* @date 2025/7/8
*/
public class PlateDetCriterialFactory {
public static Criteria<Image, DetectedObjects> createCriteria(PlateDetModelConfig config) {
Device device = null;
if(!Objects.isNull(config.getDevice())){
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
}
Criteria<Image, DetectedObjects> criteria = null;
ConcurrentHashMap params = new ConcurrentHashMap<String, String>();
params.putAll(config.getCustomParams());
if(StringUtils.isNotBlank(config.getBatchifier())){
params.put("batchifier", config.getBatchifier());
}
if(config.getModelEnum() == PlateDetModelEnum.YOLOV5){
criteria =
Criteria.builder()
.optEngine("OnnxRuntime")
.setTypes(Image.class, DetectedObjects.class)
.optModelPath(Paths.get(config.getModelPath()))
.optTranslator(new Yolo5PlateDetectTranslator(params))
.optDevice(device)
.optProgress(new ProgressBar())
.build();
}else if (config.getModelEnum() == PlateDetModelEnum.YOLOV7){
criteria =
Criteria.builder()
.optEngine("OnnxRuntime")
.setTypes(Image.class, DetectedObjects.class)
.optModelPath(Paths.get(config.getModelPath()))
.optTranslator(new Yolov7PlateDetectTranslator(params))
.optDevice(device)
.optProgress(new ProgressBar())
.build();
}
return criteria;
}
}

42
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/criteria/PlateRecCriterialFactory.java Normal file
View File

@@ -0,0 +1,42 @@
package cn.smartjavaai.ocr.model.plate.criteria;
import ai.djl.Device;
import ai.djl.modality.cv.Image;
import ai.djl.repository.zoo.Criteria;
import ai.djl.training.util.ProgressBar;
import cn.smartjavaai.common.enums.DeviceEnum;
import cn.smartjavaai.ocr.config.PlateRecModelConfig;
import cn.smartjavaai.ocr.entity.PlateResult;
import cn.smartjavaai.ocr.enums.PlateRecModelEnum;
import cn.smartjavaai.ocr.model.plate.translator.CRNNPlateRecTranslator;
import java.nio.file.Paths;
import java.util.Objects;
/**
* @author dwj
* @date 2025/7/8
*/
public class PlateRecCriterialFactory {
public static Criteria<Image, PlateResult> createCriteria(PlateRecModelConfig config) {
Device device = null;
if(!Objects.isNull(config.getDevice())){
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
}
Criteria<Image, PlateResult> criteria = null;
if(config.getModelEnum() == PlateRecModelEnum.PLATE_REC_CRNN){
criteria =
Criteria.builder()
.optEngine("OnnxRuntime")
.setTypes(Image.class, PlateResult.class)
.optModelPath(Paths.get(config.getModelPath()))
.optTranslator(new CRNNPlateRecTranslator())
.optDevice(device)
.optProgress(new ProgressBar())
.build();
}
return criteria;
}
}

88
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/CRNNPlateRecTranslator.java Normal file
View File

@@ -0,0 +1,88 @@
package cn.smartjavaai.ocr.model.plate.translator;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.util.NDImageUtils;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.ndarray.types.DataType;
import ai.djl.translate.Batchifier;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;
import cn.smartjavaai.ocr.entity.PlateResult;
import java.util.ArrayList;
import java.util.List;
/**
* @author dwj
*/
public class CRNNPlateRecTranslator implements Translator<Image, PlateResult> {
private static final String plateName = "#京沪津渝冀晋蒙辽吉黑苏浙皖闽赣鲁豫鄂湘粤桂琼川贵云藏陕甘青宁新学警港澳挂使领民航危0123456789ABCDEFGHJKLMNPQRSTUVWXYZ险品";
private static final String[] plateColors = {"黑色", "蓝色", "绿色", "白色", "黄色"};
private static final float MEAN = 0.588f;
private static final float STD = 0.193f;
@Override
public NDList processInput(TranslatorContext ctx, Image input) {
NDManager manager = ctx.getNDManager();
// Resize to (168, 48)
NDArray array = input.toNDArray(manager, Image.Flag.COLOR);
array = NDImageUtils.resize(array, 168, 48);
// Normalize
array = array.toType(DataType.FLOAT32, false)
.div(255f)
.sub(MEAN)
.div(STD);
// HWC to CHW
array = array.transpose(2, 0, 1);
array = array.expandDims(0); // batch dimension
return new NDList(array);
}
@Override
public PlateResult processOutput(TranslatorContext ctx, NDList list) {
NDArray plateOutput = list.get(0); // shape: [1, T, num_classes]
NDArray colorOutput = list.get(1); // shape: [1, num_colors]
int[] plateIdx = plateOutput.argMax(-1)
.toType(DataType.INT32, false)
.toIntArray();
int colorIdx = colorOutput.argMax(1).toType(DataType.INT32, false).toIntArray()[0];
String plateNo = decodePlate(plateIdx);
String plateColor = plateColors[colorIdx];
return new PlateResult(plateNo, plateColor);
}
private String decodePlate(int[] preds) {
int pre = 0;
List<Integer> newPreds = new ArrayList<>();
for (int idx : preds) {
if (idx != 0 && idx != pre) {
newPreds.add(idx);
}
pre = idx;
}
StringBuilder sb = new StringBuilder();
for (int i : newPreds) {
if (i >= 0 && i < plateName.length()) {
sb.append(plateName.charAt(i));
}
}
return sb.toString();
}
@Override
public Batchifier getBatchifier() {
return null; // 非批量任务
}
}

185
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/Yolo5PlateDetectTranslator.java Normal file
View File

@@ -0,0 +1,185 @@
package cn.smartjavaai.ocr.model.plate.translator;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.output.BoundingBox;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.modality.cv.output.Landmark;
import ai.djl.modality.cv.output.Point;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDArrays;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.ndarray.types.DataType;
import ai.djl.translate.Batchifier;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;
import cn.smartjavaai.common.utils.LetterBoxUtils;
import cn.smartjavaai.common.utils.NMSUtils;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
/**
* @author dwj
*/
public class Yolo5PlateDetectTranslator implements Translator<Image, DetectedObjects> {
private int inputSize = 640;
private float minConfThreshold = 0.3f;
private float iouThreshold = 0.5f;
private float confThreshold = 0;
private int imageWidth;
private int imageHeight;
private int topK;
private LetterBoxUtils.ResizeResult letterBoxResult;
public Yolo5PlateDetectTranslator(Map<String, ?> arguments) {
confThreshold =
arguments.containsKey("confThreshold")
? Integer.parseInt(arguments.get("confThreshold").toString())
: 0.3f;
iouThreshold =
arguments.containsKey("iouThreshold")
? Integer.parseInt(arguments.get("iouThreshold").toString())
: 0.5f;
topK = arguments.containsKey("topk")
? Integer.parseInt(arguments.get("topk").toString())
: 100;
}
@Override
public NDList processInput(TranslatorContext ctx, Image input) {
NDManager manager = ctx.getNDManager();
NDArray array = input.toNDArray(manager, Image.Flag.COLOR);
imageWidth = (int) array.getShape().get(1);
imageHeight = (int) array.getShape().get(0);
//Letter box resize 640x640 with padding (保持比例，补边缘)
letterBoxResult = LetterBoxUtils.letterbox(manager, array, inputSize, inputSize, 114f, LetterBoxUtils.PaddingPosition.CENTER);
array = letterBoxResult.image;
// 转为 float32 且归一化到 0~1
array = array.toType(DataType.FLOAT32, false).div(255f); // HWC
// HWC -> CHW
array = array.transpose(2, 0, 1); // CHW
return new NDList(array.expandDims(0));
}
@Override
public DetectedObjects processOutput(TranslatorContext ctx, NDList list) {
NDManager manager = ctx.getNDManager();
//[x_center, y_center, w, h, obj_conf, 8个关键点, class1_conf, class2_conf]
//目标置信度 obj_conf 5:13 关键点 [13:15]分类得分：单层车牌 / 双层车牌
NDArray dets = list.singletonOrThrow();
//置信度过滤 (1,25200, 15)
NDArray dets0 = dets.get(0);
NDArray conf = dets0.get(":, 4"); // shape [N]
NDArray mask = conf.gt(minConfThreshold);
//筛选出符合条件的框(17,15)
NDArray detsFiltered = dets0.get(mask); // 筛掉低置信度
//把分类得分 [13:15] * 置信度 [4:5] 做联合概率
NDArray clsLogits = detsFiltered.get(":, 13:15"); // (N, 2)
NDArray confFiltered = detsFiltered.get(":, 4").reshape(-1, 1); // (N, 1)
clsLogits = clsLogits.mul(confFiltered); // (N, 2)，变成 obj_conf * class_conf
NDArray jointScore = clsLogits.max(new int[]{1}); // shape (N,)
// 联合过滤
NDArray jointMask = jointScore.gt(confThreshold);
detsFiltered = detsFiltered.get(jointMask);
clsLogits = clsLogits.get(jointMask);
//中心点框 [x,y,w,h] ➔ 左上右下 [x1,y1,x2,y2]
NDArray xywh = detsFiltered.get(":, 0:4"); // (N, 4)
NDArray halfWH = xywh.get(":, 2:4").div(2); // (N, 2)
NDArray xy1 = xywh.get(":, 0:2").sub(halfWH); // (N, 2)
NDArray xy2 = xywh.get(":, 0:2").add(halfWH); // (N, 2)
NDArray boxes = NDArrays.concat(new NDList(xy1, xy2), 1); // (N, 4)
// 分类得分最大值：score (N, 1)，对应类别 index (N, 1)
NDArray scores = clsLogits.max(new int[]{1}, true); // (N, 1)
NDArray indices = clsLogits.argMax(1).reshape(-1, 1).toType(DataType.FLOAT32, false); // (N, 1)
// 关键点坐标 [5:13]
NDArray keyPoints = detsFiltered.get(":, 5:13"); // (N, 8)
// 拼成最终结果：(x1, y1, x2, y2, score, 8关键点, index)
NDArray output = NDArrays.concat(new NDList(boxes, scores, keyPoints, indices), 1); // (N, 14)
// NMS 过滤掉重叠框
int[] keepIndices = NMSUtils.nms(boxes, scores.squeeze(), iouThreshold); // scores.squeeze() ➝ (N,)
NDArray kept = output.get(manager.create(keepIndices));
// 如果超过 topK，则截断
if (keepIndices.length > topK) {
int[] topkIndices = new int[topK];
System.arraycopy(keepIndices, 0, topkIndices, 0, topK);
keepIndices = topkIndices;
}
//恢复原图坐标（除回比例，减掉 padding）
NDArray restored = LetterBoxUtils.restoreBox(kept, letterBoxResult.r, letterBoxResult.left, letterBoxResult.top, 5,8);
List<String> classNames = new ArrayList<>();
List<Double> probabilities = new ArrayList<>();
List<BoundingBox> boundingBoxes = new ArrayList<>();
float[] flatData = restored.toFloatArray();
long[] shape = restored.getShape().getShape(); // 比如 (N, 14)
int rows = (int) shape[0];
int cols = (int) shape[1];
// 把一维数组重组为二维数组
float[][] data = new float[rows][cols];
for (int i = 0; i < rows; i++) {
System.arraycopy(flatData, i * cols, data[i], 0, cols);
}
for (float[] row : data) {
// row结构：(x1, y1, x2, y2, score, kp1,..., kp8, classIndex)
float x1 = row[0];
float y1 = row[1];
float x2 = row[2];
float y2 = row[3];
float score = row[4];
int classIndex = (int) row[13];
double prob = score;
String className = classIndex == 0 ? "single" : "double";
// 转相对坐标，DJL的Rectangle用比例坐标（0~1）
double rectX = x1 / imageWidth;
double rectY = y1 / imageHeight;
double rectW = (x2 - x1) / imageWidth;
double rectH = (y2 - y1) / imageHeight;
// 构建 Polygon 四个角点
List<Point> pointsSrc = new ArrayList<>();
pointsSrc.add(new Point(row[5], row[6]));
pointsSrc.add(new Point(row[7], row[8]));
pointsSrc.add(new Point(row[9], row[10]));
pointsSrc.add(new Point(row[11], row[12]));
Landmark box = new Landmark(rectX, rectY, rectW, rectH, pointsSrc);
classNames.add(className);
probabilities.add(prob);
boundingBoxes.add(box);
}
DetectedObjects detectedObjects = new DetectedObjects(classNames, probabilities, boundingBoxes);
return detectedObjects;
}
@Override
public Batchifier getBatchifier() {
return null;
}
}

189
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/Yolov7PlateDetectTranslator.java Normal file
View File

@@ -0,0 +1,189 @@
package cn.smartjavaai.ocr.model.plate.translator;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.output.*;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDArrays;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.ndarray.types.DataType;
import ai.djl.translate.Batchifier;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;
import cn.smartjavaai.common.utils.LetterBoxUtils;
import cn.smartjavaai.common.utils.NMSUtils;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
/**
* @author dwj
*/
public class Yolov7PlateDetectTranslator implements Translator<Image, DetectedObjects> {
private int inputSize = 640;
private float minConfThreshold = 0.3f;
private float iouThreshold = 0.5f;
private float confThreshold = 0;
private int imageWidth;
private int imageHeight;
private int topK;
private LetterBoxUtils.ResizeResult letterBoxResult;
public Yolov7PlateDetectTranslator(Map<String, ?> arguments) {
confThreshold =
arguments.containsKey("confThreshold")
? Integer.parseInt(arguments.get("confThreshold").toString())
: 0.3f;
iouThreshold =
arguments.containsKey("iouThreshold")
? Integer.parseInt(arguments.get("iouThreshold").toString())
: 0.5f;
topK = arguments.containsKey("topk")
? Integer.parseInt(arguments.get("topk").toString())
: 100;
}
@Override
public NDList processInput(TranslatorContext ctx, Image input) {
NDManager manager = ctx.getNDManager();
NDArray array = input.toNDArray(manager, Image.Flag.COLOR);
imageWidth = (int) array.getShape().get(1);
imageHeight = (int) array.getShape().get(0);
//Letter box resize 640x640 with padding (保持比例，补边缘)
letterBoxResult = LetterBoxUtils.letterbox(manager, array, inputSize, inputSize, 114f, LetterBoxUtils.PaddingPosition.CENTER);
array = letterBoxResult.image;
// 转为 float32 且归一化到 0~1
array = array.toType(DataType.FLOAT32, false).div(255f); // HWC
// HWC -> CHW
array = array.transpose(2, 0, 1); // CHW
return new NDList(array.expandDims(0));
}
@Override
public DetectedObjects processOutput(TranslatorContext ctx, NDList list) {
NDManager manager = ctx.getNDManager();
int num_cls = 2;
//[x_center, y_center, w, h, obj_conf, class1_conf, class2_conf,8个关键点]
//目标置信度 obj_conf 5:13 关键点 [13:15]分类得分：单层车牌 / 双层车牌
NDArray dets = list.singletonOrThrow();
//置信度过滤 (1,25200, 15)
NDArray dets0 = dets.get(0);
NDArray conf = dets0.get(":, 4"); // shape [N]
NDArray mask = conf.gt(minConfThreshold);
//筛选出符合条件的框(17,15)
NDArray detsFiltered = dets0.get(mask); // 筛掉低置信度
//把分类得分 [5:7] * 置信度 [4:5] 做联合概率
NDArray clsLogits = detsFiltered.get(":, 5:7"); // (N, 2)
NDArray confFiltered = detsFiltered.get(":, 4").reshape(-1, 1); // (N, 1)
clsLogits = clsLogits.mul(confFiltered); // (N, 2)，变成 obj_conf * class_conf
NDArray jointScore = clsLogits.max(new int[]{1}); // shape (N,)
// 联合过滤
NDArray jointMask = jointScore.gt(confThreshold);
detsFiltered = detsFiltered.get(jointMask);
clsLogits = clsLogits.get(jointMask);
//中心点框 [x,y,w,h] ➔ 左上右下 [x1,y1,x2,y2]
NDArray xywh = detsFiltered.get(":, 0:4"); // (N, 4)
NDArray halfWH = xywh.get(":, 2:4").div(2); // (N, 2)
NDArray xy1 = xywh.get(":, 0:2").sub(halfWH); // (N, 2)
NDArray xy2 = xywh.get(":, 0:2").add(halfWH); // (N, 2)
NDArray boxes = NDArrays.concat(new NDList(xy1, xy2), 1); // (N, 4)
// 分类得分最大值：score (N, 1)，对应类别 index (N, 1)
NDArray scores = clsLogits.max(new int[]{1}, true); // (N, 1)
NDArray indices = clsLogits.argMax(1).reshape(-1, 1).toType(DataType.FLOAT32, false); // (N, 1)
// 关键点坐标 [7,8,10,11,13,14,16,17]
NDArray keyPoints = NDArrays.concat(new NDList(
detsFiltered.get(":, 7:8"),
detsFiltered.get(":, 8:9"),
detsFiltered.get(":, 10:11"),
detsFiltered.get(":, 11:12"),
detsFiltered.get(":, 13:14"),
detsFiltered.get(":, 14:15"),
detsFiltered.get(":, 16:17"),
detsFiltered.get(":, 17:18")
), 1); // 拼成 (N, 8)
// 拼成最终结果：(x1, y1, x2, y2, score, 8关键点, index)
NDArray output = NDArrays.concat(new NDList(boxes, scores, keyPoints, indices), 1); // (N, 14)
// NMS 过滤掉重叠框
int[] keepIndices = NMSUtils.nms(boxes, scores.squeeze(), iouThreshold); // scores.squeeze() ➝ (N,)
NDArray kept = output.get(manager.create(keepIndices));
// 如果超过 topK，则截断
if (keepIndices.length > topK) {
int[] topkIndices = new int[topK];
System.arraycopy(keepIndices, 0, topkIndices, 0, topK);
keepIndices = topkIndices;
}
//恢复原图坐标（除回比例，减掉 padding）
NDArray restored = LetterBoxUtils.restoreBox(kept, letterBoxResult.r, letterBoxResult.left, letterBoxResult.top, 5,8);
List<String> classNames = new ArrayList<>();
List<Double> probabilities = new ArrayList<>();
List<BoundingBox> boundingBoxes = new ArrayList<>();
float[] flatData = restored.toFloatArray();
long[] shape = restored.getShape().getShape(); // 比如 (N, 14)
int rows = (int) shape[0];
int cols = (int) shape[1];
// 把一维数组重组为二维数组
float[][] data = new float[rows][cols];
for (int i = 0; i < rows; i++) {
System.arraycopy(flatData, i * cols, data[i], 0, cols);
}
for (float[] row : data) {
// row结构：(x1, y1, x2, y2, score, kp1,..., kp8, classIndex)
float x1 = row[0];
float y1 = row[1];
float x2 = row[2];
float y2 = row[3];
float score = row[4];
int classIndex = (int) row[13];
double prob = score;
String className = classIndex == 0 ? "single" : "double";
// 转相对坐标，DJL的Rectangle用比例坐标（0~1）
double rectX = x1 / imageWidth;
double rectY = y1 / imageHeight;
double rectW = (x2 - x1) / imageWidth;
double rectH = (y2 - y1) / imageHeight;
// 构建 Polygon 四个角点
List<Point> pointsSrc = new ArrayList<>();
pointsSrc.add(new Point(row[5], row[6]));
pointsSrc.add(new Point(row[7], row[8]));
pointsSrc.add(new Point(row[9], row[10]));
pointsSrc.add(new Point(row[11], row[12]));
Landmark box = new Landmark(rectX, rectY, rectW, rectH, pointsSrc);
classNames.add(className);
probabilities.add(prob);
boundingBoxes.add(box);
}
DetectedObjects detectedObjects = new DetectedObjects(classNames, probabilities, boundingBoxes);
return detectedObjects;
}
@Override
public Batchifier getBatchifier() {
return null;
}
}

186
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/Yolov8PlateDetectTranslator.java Normal file
View File

@@ -0,0 +1,186 @@
package cn.smartjavaai.ocr.model.plate.translator;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.output.*;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDArrays;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.ndarray.types.DataType;
import ai.djl.translate.Batchifier;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;
import cn.smartjavaai.common.utils.LetterBoxUtils;
import cn.smartjavaai.common.utils.NMSUtils;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
/**
* @author dwj
*/
public class Yolov8PlateDetectTranslator implements Translator<Image, DetectedObjects> {
private int inputSize = 640;
private float minConfThreshold = 0.3f;
private float iouThreshold = 0.5f;
private float confThreshold = 0;
private int imageWidth;
private int imageHeight;
private int topK;
private LetterBoxUtils.ResizeResult letterBoxResult;
public Yolov8PlateDetectTranslator(Map<String, ?> arguments) {
confThreshold =
arguments.containsKey("confThreshold")
? Integer.parseInt(arguments.get("confThreshold").toString())
: 0.3f;
iouThreshold =
arguments.containsKey("iouThreshold")
? Integer.parseInt(arguments.get("iouThreshold").toString())
: 0.5f;
topK = arguments.containsKey("topk")
? Integer.parseInt(arguments.get("topk").toString())
: 100;
}
@Override
public NDList processInput(TranslatorContext ctx, Image input) {
NDManager manager = ctx.getNDManager();
NDArray array = input.toNDArray(manager, Image.Flag.COLOR);
imageWidth = (int) array.getShape().get(1);
imageHeight = (int) array.getShape().get(0);
//Letter box resize 640x640 with padding (保持比例，补边缘)
letterBoxResult = LetterBoxUtils.letterbox(manager, array, inputSize, inputSize, 114f, LetterBoxUtils.PaddingPosition.CENTER);
array = letterBoxResult.image;
// 转为 float32 且归一化到 0~1
array = array.toType(DataType.FLOAT32, false).div(255f); // HWC
// HWC -> CHW
array = array.transpose(2, 0, 1); // CHW
return new NDList(array.expandDims(0));
}
@Override
public DetectedObjects processOutput(TranslatorContext ctx, NDList list) {
NDManager manager = ctx.getNDManager();
NDArray preds = list.get(0); // shape: (1, 6, 8400)
preds = preds.squeeze(0).transpose(1, 0); // shape: (8400, 6)
// preds shape: (8400, 6)
NDArray classScores = preds.get(":, 4:6"); // shape: (8400, 2)
// 获取每行最大值（对应 Python 的 .amax(1)）
NDArray maxScores = classScores.max(new int[]{1}); // shape: (8400,)
// 构造 mask：score > conf
NDArray confMask = maxScores.gt(minConfThreshold); // shape: (8400,)
// 应用 mask 筛选
preds = preds.get(confMask); // shape: (N_filtered, 6)
if (preds.isEmpty()) {
return null;
}
// 提取 box (xywh)，转换为 xyxy
NDArray boxes = preds.get(":, 0:4"); // shape: (N, 4)
boxes = xywh2xyxy(boxes); // 自定义函数：center xywh -> xyxy
// 1. 得分和类别索引
NDArray scoresAndClasses = preds.get(":, 4:6"); // shape (num, 2)
NDArray scores = scoresAndClasses.max(new int[]{1}, true); // keepDim = true
NDArray index = scoresAndClasses.argMax(1).expandDims(1); // 最大值索引，类别，shape (num, 1)
// 4. 拼接
NDArray result = NDArrays.concat(new NDList(boxes, scores, index), 1); // 在列方向拼接
// NMS 过滤掉重叠框
int[] keepIndices = NMSUtils.nms(boxes, scores.squeeze(), iouThreshold); // scores.squeeze() ➝ (N,)
NDArray kept = result.get(manager.create(keepIndices));
// 如果超过 topK，则截断
if (keepIndices.length > topK) {
int[] topkIndices = new int[topK];
System.arraycopy(keepIndices, 0, topkIndices, 0, topK);
keepIndices = topkIndices;
}
//恢复原图坐标（除回比例，减掉 padding）
NDArray restored = LetterBoxUtils.restoreBox(kept, letterBoxResult.r, letterBoxResult.left, letterBoxResult.top, 5,0);
List<String> classNames = new ArrayList<>();
List<Double> probabilities = new ArrayList<>();
List<BoundingBox> boundingBoxes = new ArrayList<>();
float[] flatData = restored.toFloatArray();
long[] shape = restored.getShape().getShape(); // 比如 (N, 14)
int rows = (int) shape[0];
int cols = (int) shape[1];
// 把一维数组重组为二维数组
float[][] data = new float[rows][cols];
for (int i = 0; i < rows; i++) {
System.arraycopy(flatData, i * cols, data[i], 0, cols);
}
for (float[] row : data) {
// row结构：(x1, y1, x2, y2, score, classIndex)
float x1 = row[0];
float y1 = row[1];
float x2 = row[2];
float y2 = row[3];
float score = row[4];
int classIndex = (int) row[5];
double prob = score;
String className = classIndex == 0 ? "single" : "double";
// 转相对坐标，DJL的Rectangle用比例坐标（0~1）
double rectX = x1 / imageWidth;
double rectY = y1 / imageHeight;
double rectW = (x2 - x1) / imageWidth;
double rectH = (y2 - y1) / imageHeight;
// 构建 Polygon 四个角点
// List<Point> pointsSrc = new ArrayList<>();
// pointsSrc.add(new Point(row[5], row[6]));
// pointsSrc.add(new Point(row[7], row[8]));
// pointsSrc.add(new Point(row[9], row[10]));
// pointsSrc.add(new Point(row[11], row[12]));
Rectangle rectangle = new Rectangle(rectX, rectY, rectW, rectH);
classNames.add(className);
probabilities.add(prob);
boundingBoxes.add(rectangle);
}
DetectedObjects detectedObjects = new DetectedObjects(classNames, probabilities, boundingBoxes);
return detectedObjects;
}
@Override
public Batchifier getBatchifier() {
return null;
}
public static NDArray xywh2xyxy(NDArray xywh) {
NDArray x = xywh.get(":, 0");
NDArray y = xywh.get(":, 1");
NDArray w = xywh.get(":, 2").div(2);
NDArray h = xywh.get(":, 3").div(2);
NDArray x1 = x.sub(w);
NDArray y1 = y.sub(h);
NDArray x2 = x.add(w);
NDArray y2 = y.add(h);
return NDArrays.stack(new NDList(x1, y1, x2, y2), 1);
}
}

13
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/CommonTableStructureModel.java
View File

@@ -47,7 +47,7 @@ public class CommonTableStructureModel implements TableStructureModel{
private ZooModel<Image, TableStructureResult> model;
private ObjectPool<Predictor<Image, TableStructureResult>> predictorPool;
private GenericObjectPool<Predictor<Image, TableStructureResult>> predictorPool;
@Override
public void loadModel(TableStructureConfig config) {
@@ -59,8 +59,14 @@ public class CommonTableStructureModel implements TableStructureModel{
model = ModelZoo.loadModel(criteria);
// 创建池子：每个线程独享 Predictor
this.predictorPool = new GenericObjectPool<>(new PredictorFactory<>(model));
int predictorPoolSize = config.getPredictorPoolSize();
if(config.getPredictorPoolSize() <= 0){
predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
}
predictorPool.setMaxTotal(predictorPoolSize);
log.debug("当前设备: " + model.getNDManager().getDevice());
log.debug("当前引擎: " + Engine.getInstance().getEngineName());
log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
throw new OcrException("表格结构识别模型加载失败", e);
}
@@ -140,6 +146,11 @@ public class CommonTableStructureModel implements TableStructureModel{
}
}
@Override
public GenericObjectPool<Predictor<Image, TableStructureResult>> getPool() {
return predictorPool;
}
@Override
public void close() throws Exception {
try {

6
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/TableStructureModel.java
View File

@@ -1,5 +1,6 @@
package cn.smartjavaai.ocr.model.table;
import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import cn.smartjavaai.common.entity.R;
import cn.smartjavaai.ocr.config.OcrDetModelConfig;
@@ -7,6 +8,7 @@ import cn.smartjavaai.ocr.config.TableStructureConfig;
import cn.smartjavaai.ocr.entity.OcrBox;
import cn.smartjavaai.ocr.entity.OcrItem;
import cn.smartjavaai.ocr.entity.TableStructureResult;
import org.apache.commons.pool2.impl.GenericObjectPool;
import java.awt.image.BufferedImage;
import java.util.List;
@@ -60,4 +62,8 @@ public interface TableStructureModel extends AutoCloseable{
default R<TableStructureResult> detect(Image image){
throw new UnsupportedOperationException("默认不支持该功能");
}
default GenericObjectPool<Predictor<Image, TableStructureResult>> getPool() {
throw new UnsupportedOperationException("默认不支持该功能");
}
}

64
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/opencv/OcrOpenCVUtils.java
View File

@@ -1,64 +0,0 @@
package cn.smartjavaai.ocr.opencv;
import org.opencv.core.CvType;
import org.opencv.core.Mat;
import org.opencv.imgproc.Imgproc;
import java.awt.image.BufferedImage;
import java.awt.image.DataBufferByte;
/**
* OpenCV Utils
*
*/
public class OcrOpenCVUtils {
/**
* 透视变换
*
* @param src
* @param srcPoints
* @param dstPoints
* @return
*/
public static Mat perspectiveTransform(Mat src, Mat srcPoints, Mat dstPoints) {
Mat dst = src.clone();
Mat warp_mat = Imgproc.getPerspectiveTransform(srcPoints, dstPoints);
Imgproc.warpPerspective(src, dst, warp_mat, dst.size());
warp_mat.release();
return dst;
}
/**
* Mat to BufferedImage
*
* @param mat
* @return
*/
public static BufferedImage mat2Image(Mat mat) {
int width = mat.width();
int height = mat.height();
byte[] data = new byte[width * height * (int) mat.elemSize()];
Imgproc.cvtColor(mat, mat, 4);
mat.get(0, 0, data);
BufferedImage ret = new BufferedImage(width, height, 5);
ret.getRaster().setDataElements(0, 0, width, height, data);
return ret;
}
/**
* BufferedImage to Mat
*
* @param img
* @return
*/
public static Mat image2Mat(BufferedImage img) {
int width = img.getWidth();
int height = img.getHeight();
byte[] data = ((DataBufferByte) img.getRaster().getDataBuffer()).getData();
Mat mat = new Mat(height, width, CvType.CV_8UC3);
mat.put(0, 0, data);
return mat;
}
}

166
smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/utils/OcrUtils.java
View File

@@ -11,10 +11,13 @@ import ai.djl.ndarray.NDManager;
import ai.djl.opencv.OpenCVImageFactory;
import cn.smartjavaai.common.entity.*;
import cn.smartjavaai.common.entity.Point;
import cn.smartjavaai.common.utils.ImageUtils;
import cn.smartjavaai.common.utils.OpenCVUtils;
import cn.smartjavaai.common.utils.PointUtils;
import cn.smartjavaai.ocr.entity.*;
import cn.smartjavaai.ocr.enums.AngleEnum;
import cn.smartjavaai.ocr.enums.PlateType;
import cn.smartjavaai.ocr.opencv.OcrNDArrayUtils;
import cn.smartjavaai.ocr.opencv.OcrOpenCVUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.opencv.core.Mat;
@@ -73,19 +76,6 @@ public class OcrUtils {
/**
* 欧式距离计算
*
* @param point1
* @param point2
* @return
*/
public static float distance(float[] point1, float[] point2) {
float disX = point1[0] - point2[0];
float disY = point1[1] - point2[1];
float dis = (float) Math.sqrt(disX * disX + disY * disY);
return dis;
}
/**
* 图片旋转
@@ -204,9 +194,64 @@ public class OcrUtils {
}
/**
* 透视变换 + 裁剪
* @param srcMat
* @param landMarks
* @return
*/
public static Image transformAndCrop(Mat srcMat, List<ai.djl.modality.cv.output.Point> landMarks){
if (landMarks == null || landMarks.size() != 4) {
throw new IllegalArgumentException("必须提供4个关键点");
}
// 步骤 1：排序为 左上、右上、右下、左下
List<ai.djl.modality.cv.output.Point> ordered = PointUtils.orderPoints(landMarks);
ai.djl.modality.cv.output.Point lt = ordered.get(0);
ai.djl.modality.cv.output.Point rt = ordered.get(1);
ai.djl.modality.cv.output.Point rb = ordered.get(2);
ai.djl.modality.cv.output.Point lb = ordered.get(3);
// 步骤 2：计算目标图像尺寸（宽、高）
int img_crop_width = (int) Math.max(
PointUtils.distance(lt, rt),
PointUtils.distance(rb, lb)
);
int img_crop_height = (int) Math.max(
PointUtils.distance(lt, lb),
PointUtils.distance(rt, rb)
);
// 步骤 3：构造目标坐标点
List<ai.djl.modality.cv.output.Point> dstPoints = Arrays.asList(
new ai.djl.modality.cv.output.Point(0, 0),
new ai.djl.modality.cv.output.Point(img_crop_width, 0),
new ai.djl.modality.cv.output.Point(img_crop_width, img_crop_height),
new ai.djl.modality.cv.output.Point(0, img_crop_height)
);
// 步骤 4：透视变换
Mat srcPoint2f = OcrNDArrayUtils.toMat(ordered);
Mat dstPoint2f = OcrNDArrayUtils.toMat(dstPoints);
Mat cvMat = OpenCVUtils.perspectiveTransform(srcMat, srcPoint2f, dstPoint2f);
// 步骤 5：转为 DJL Image + 裁剪
Image subImg = OpenCVImageFactory.getInstance().fromImage(cvMat);
subImg = subImg.getSubImage(0, 0, img_crop_width, img_crop_height);
// 释放资源
cvMat.release();
srcPoint2f.release();
dstPoint2f.release();
return subImg;
}
/**
* 放射变换+裁剪
* 透视变换+裁剪
* @param srcMat
* @param box
* @return
@@ -217,8 +262,8 @@ public class OcrUtils {
float[] rt = java.util.Arrays.copyOfRange(pointsArr, 2, 4);
float[] rb = java.util.Arrays.copyOfRange(pointsArr, 4, 6);
float[] lb = java.util.Arrays.copyOfRange(pointsArr, 6, 8);
int img_crop_width = (int) Math.max(OcrUtils.distance(lt, rt), OcrUtils.distance(rb, lb));
int img_crop_height = (int) Math.max(OcrUtils.distance(lt, lb), OcrUtils.distance(rt, rb));
int img_crop_width = (int) Math.max(PointUtils.distance(lt, rt), PointUtils.distance(rb, lb));
int img_crop_height = (int) Math.max(PointUtils.distance(lt, lb), PointUtils.distance(rt, rb));
List<ai.djl.modality.cv.output.Point> srcPoints = new ArrayList<>();
srcPoints.add(new ai.djl.modality.cv.output.Point(lt[0], lt[1]));
srcPoints.add(new ai.djl.modality.cv.output.Point(rt[0], rt[1]));
@@ -232,7 +277,7 @@ public class OcrUtils {
Mat srcPoint2f = OcrNDArrayUtils.toMat(srcPoints);
Mat dstPoint2f = OcrNDArrayUtils.toMat(dstPoints);
//透视变换
Mat cvMat = OcrOpenCVUtils.perspectiveTransform(srcMat, srcPoint2f, dstPoint2f);
Mat cvMat = OpenCVUtils.perspectiveTransform(srcMat, srcPoint2f, dstPoint2f);
Image subImg = OpenCVImageFactory.getInstance().fromImage(cvMat);
//ImageUtils.saveImage(subImg, i + ".png", "build/output");
//变换后裁剪
@@ -320,5 +365,90 @@ public class OcrUtils {
}
}
public static List<PlateInfo> convertToPlateInfo(DetectedObjects detectedObjects, Image image) {
List<PlateInfo> plateInfoList = new ArrayList<>();
Iterator iterator = detectedObjects.items().iterator();
int index = 0;
while(iterator.hasNext()) {
DetectedObjects.DetectedObject result = (DetectedObjects.DetectedObject)iterator.next();
BoundingBox box = result.getBoundingBox();
List<Point> keyPoints = new ArrayList<Point>();
box.getBounds().getPath().forEach(point -> {
keyPoints.add(new Point(point.getX(), point.getY()));
});
int x = (int)(box.getBounds().getX() * image.getWidth());
int y = (int)(box.getBounds().getY() * image.getHeight());
int width = (int)(box.getBounds().getWidth() * image.getWidth());
int height = (int)(box.getBounds().getHeight() * image.getHeight());
// 修正边界，防止越界
if (x < 0) x = 0;
if (y < 0) y = 0;
if (x + width > image.getWidth()) width = image.getWidth() - x;
if (y + height > image.getHeight()) height = image.getHeight() - y;
PlateInfo plateInfo = new PlateInfo();
plateInfo.setPlateType(PlateType.fromClassName(detectedObjects.getClassNames().get(index)));
plateInfo.setScore(detectedObjects.getProbabilities().get(index).floatValue());
plateInfo.setDetectionRectangle(new DetectionRectangle(x, y, width, height));
OcrBox ocrBox = new OcrBox(keyPoints.get(0), keyPoints.get(1), keyPoints.get(2), keyPoints.get(3));
plateInfo.setBox(ocrBox);
plateInfoList.add(plateInfo);
index++;
}
return plateInfoList;
}
/**
* 绘制车牌信息
* @param srcMat
* @param plateInfoList
*/
public static void drawPlateInfo(Mat srcMat, List<PlateInfo> plateInfoList) {
for(PlateInfo plateInfo : plateInfoList){
OcrBox ocrBox = plateInfo.getBox();
Imgproc.line(srcMat, ocrBox.getTopLeft().toCvPoint(), ocrBox.getTopRight().toCvPoint(), new Scalar(0, 0, 255), 1);
Imgproc.line(srcMat, ocrBox.getTopRight().toCvPoint(), ocrBox.getBottomRight().toCvPoint(), new Scalar(0, 0, 255),1);
Imgproc.line(srcMat, ocrBox.getBottomRight().toCvPoint(), ocrBox.getBottomLeft().toCvPoint(), new Scalar(0, 0, 255),1);
Imgproc.line(srcMat, ocrBox.getBottomLeft().toCvPoint(), ocrBox.getTopLeft().toCvPoint(), new Scalar(0, 0, 255), 1);
// 中文乱码
ImageUtils.putTextWithBackground(srcMat, plateInfo.getPlateNumber() + " " + plateInfo.getPlateColor(), ocrBox.getTopLeft().toCvPoint(), new Scalar(255, 255, 255), new Scalar(0, 0, 0), 1);
}
}
/**
* 在图像上绘制带白色背景、黑色文字的文本
*/
public static void drawPlateInfo(BufferedImage image, List<PlateInfo> plateInfoList) {
// 将绘制图像转换为Graphics2D
Graphics2D graphics = (Graphics2D) image.getGraphics();
try {
graphics.setColor(Color.RED);// 边框颜色
graphics.setStroke(new BasicStroke(2)); // 线宽2像素
graphics.setRenderingHint(RenderingHints.KEY_ANTIALIASING,
RenderingHints.VALUE_ANTIALIAS_ON); // 抗锯齿
int stroke = 2;
for(PlateInfo plateInfo : plateInfoList){
DetectionRectangle rectangle = plateInfo.getDetectionRectangle();
graphics.setColor(Color.RED);// 边框颜色
//绘制车牌框
graphics.drawRect(rectangle.getX(), rectangle.getY(), rectangle.getWidth(), rectangle.getHeight());
graphics.setColor(Color.BLACK);// 字体颜色
ImageUtils.drawText(graphics, plateInfo.getPlateNumber() + " " + plateInfo.getPlateColor(), rectangle.getX(), rectangle.getY(), stroke, 4);
OcrBox ocrBox = plateInfo.getBox();
//绘制关键点
graphics.setColor(Color.BLUE);
graphics.drawRect((int)ocrBox.getTopLeft().getX(), (int)ocrBox.getTopLeft().getY(), 2, 2);
graphics.setColor(Color.GREEN);
graphics.drawRect((int)ocrBox.getTopRight().getX(), (int)ocrBox.getTopRight().getY(), 2, 2);
graphics.setColor(Color.RED);
graphics.drawRect((int)ocrBox.getBottomLeft().getX(), (int)ocrBox.getBottomLeft().getY(), 2, 2);
graphics.setColor(Color.CYAN);
graphics.drawRect((int)ocrBox.getBottomRight().getX(), (int)ocrBox.getBottomRight().getY(), 2, 2);
}
} finally {
graphics.dispose();
}
}
}