diff --git a/README.md b/README.md index a59d57c..37ca1ed 100644 --- a/README.md +++ b/README.md @@ -192,6 +192,25 @@ SmartJavaAI是专为JAVA 开发者打造的一个功能丰富、开箱即用的 + + +
+

车牌识别

+ - 单层/双层检测
+ - 车牌颜色识别
+
+ + +
+ +
+ + +
+ +
+ +
@@ -251,21 +270,22 @@ SmartJavaAI是专为JAVA 开发者打造的一个功能丰富、开箱即用的 - 支持任意角度识别，方向校准 - 支持通用文字识别，通用手写字识别 - 支持表格识别 + - 支持中文车牌识别：单层/双层检测，颜色识别，支持12种中文车牌 - **机器翻译** - 集成NLLB-200模型：支持200+语言互相翻译 ## 🌟 AI集成方式对比 -| 方案 | 技术特点 | 优点 | 缺点 | -|---------------------|--------------------------------------------------------------------------|---------------------------------------------------------------------|---------------------------------------------------------------------| -| **OpenCV** | 传统图像处理方案 | ✅ 提供java接口
✅ 轻量级部署
✅ 社区资源丰富 | ❌ 基于传统算法精度低(60%-75%)
❌ 需本地安装环境 | -| ‌**商业闭源SDK（如虹软等）**‌ | 商业级闭源解决方案 | ✅ 开箱即用
✅ 提供完整文档和SDK
✅ 支持离线活体检测 | ❌ 免费版需年度授权更新
❌ 商业授权费用高
❌ 代码不可控 | -| ‌**云API(阿里云)**‌ | SaaS化云端服务 | ✅ 零部署成本
✅ 支持高并发
✅ 自带模型迭代 | ❌ 网络延迟风险(200-800ms)
❌ 按调用量计费
❌ 有数据安全风险 | +| 方案 | 技术特点 | 优点 | 缺点 | +|---------------------|--------------------------------------------------------------------------|---------------------------------------------------------------------|-----------------------------------------------| +| **OpenCV** | 传统图像处理方案 | ✅ 提供java接口
✅ 轻量级部署
✅ 社区资源丰富 | ❌ 基于传统算法精度低(60%-75%)
❌ 需本地安装环境 | +| ‌**商业闭源SDK（如虹软等）**‌ | 商业级闭源解决方案 | ✅ 开箱即用
✅ 提供完整文档和SDK
✅ 支持离线活体检测 | ❌ 免费版需年度授权更新
❌ 商业授权费用高
❌ 代码不可控 | +| ‌**云API(阿里云)**‌ | SaaS化云端服务 | ✅ 零部署成本
✅ 支持高并发
✅ 自带模型迭代 | ❌ 网络延迟风险(200-800ms)
❌ 按调用量计费
❌ 有数据安全风险 | | ‌**Python混合调用**‌ | 跨语言调用方案 | ✅ 可集成PyTorch/TF等框架
✅ 支持自定义算法
✅ 识别精度高 | ❌ 需维护双语言环境
❌ 进程通信性能损耗(30%+)
❌ 异常处理复杂度翻倍 | -| ‌**JNI/JNA**‌ | 跨语言底层调用方案 | ✅ 直接调用 C/C++ 高性能算法库✅ 支持调用各种原生成熟库✅ 可封装成通用工具Jar |❌ 开发成本高，JNI更复杂❌ 跨平台兼容性差 | -| ‌**DJL框架**‌ | 深度学习框架 | ✅ 纯Java实现
✅ 支持主流深度学习框架
✅ 可加载预训练模型(99%+) | ❌ 需掌握DL知识
❌ 需处理模型加载、预处理、后处理等复杂技术细节 | -| ‌**SmartJavaAI**‌ | java深度学习工具包 | ✅ 支持主流深度学习框架
✅ 提供丰富、开箱即用API
✅ 上手简单，单一Jar包集成 | ❌要求JDK版本11及以上 | +| ‌**JNI/JNA**‌ | 跨语言底层调用方案 | ✅ 直接调用 C/C++ 高性能算法库✅ 支持调用各种原生成熟库✅ 可封装成通用工具Jar | ❌ 开发成本高，JNI更复杂❌ 跨平台兼容性差 | +| ‌**DJL框架**‌ | 深度学习框架 | ✅ 纯Java实现
✅ 支持主流深度学习框架
✅ 可加载预训练模型(99%+) | ❌ 需掌握DL知识
❌ 需处理模型加载、预处理、后处理等复杂技术细节 | +| ‌**SmartJavaAI**‌ | java深度学习工具包 | ✅ 支持主流深度学习框架
✅ 提供丰富、开箱即用API
✅ 上手简单，单一Jar包集成 | 无 | @@ -300,7 +320,7 @@ SmartJavaAI是专为JAVA 开发者打造的一个功能丰富、开箱即用的 cn.smartjavaai smartjavaai-all - 1.0.20 + 1.0.22 ``` ### 3、完整示例代码 @@ -406,16 +426,6 @@ SmartJavaAI是专为JAVA 开发者打造的一个功能丰富、开箱即用的 - 自训练模型推理 - yolov3~yolov12 系列 -**SSD 系列** - -| 模型名称 | 引擎 | 骨干网络 | 输入尺寸 |
训练数据集
| 精度（mAP） |
推理速度
|
适用场景
| -| :-------- |------------------| ------------- | ----------- |--------------------------------------| -------------|-------------------------------------| -------------| -|SSD_300_RESNET5| PyTorch | ResNet‑50 | 300×300 | COCO | 中等 | 快 | 精度需求一般| -|SSD_512_RESNET50_V1_VOC| PyTorch | ResNet‑50 | 512×512 | Pascal VOC | 稍高 | 中等 | 精度优先、可接受略低速度的场景| -|SSD_512_VGG16_ATROUS_COCO| MXNet | VGG‑16 | 512×512 | COCO | 较高 | 中等 | 通用场景；对小目标有一定提升| -|SSD_300_VGG16_ATROUS_VOC| MXNet | VGG‑16 | 300×300 | Pascal VOC | 中等偏上 | 快 | VOC 数据集同类任务；资源受限时使用| -|SSD_512_MOBILENET1_VOC| MXNet | MobileNet‑1.0 | 512×512 | Pascal VOC | 中等 | 快 | 嵌入式/移动端设备；算力和内存都很有限| - **YOLO 系列** |模型名称 | 引擎 | 版本 | 大小（Backbone） |
数据集
|
精度
|
速度
|
适用场景
| @@ -427,6 +437,18 @@ SmartJavaAI是专为JAVA 开发者打造的一个功能丰富、开箱即用的 |YOLOV5S_ONNXRUNTIME| OnnxRuntime | v5 | 小型 | COCO | 较高 | 加速（需 ONNX 支持） | Windows/Linux 通用加速部署 | |YOLO (MXNet / 通用模型) | MXNet | v3 | DarkNet‑53 | COCO | 较高 | 快 | 需要 MXNet 生态或复现老项目时使用 | + +**SSD 系列** + +| 模型名称 | 引擎 | 骨干网络 | 输入尺寸 |
训练数据集
| 精度（mAP） |
推理速度
|
适用场景
| +| :-------- |------------------| ------------- | ----------- |--------------------------------------| -------------|-------------------------------------| -------------| +|SSD_300_RESNET5| PyTorch | ResNet‑50 | 300×300 | COCO | 中等 | 快 | 精度需求一般| +|SSD_512_RESNET50_V1_VOC| PyTorch | ResNet‑50 | 512×512 | Pascal VOC | 稍高 | 中等 | 精度优先、可接受略低速度的场景| +|SSD_512_VGG16_ATROUS_COCO| MXNet | VGG‑16 | 512×512 | COCO | 较高 | 中等 | 通用场景；对小目标有一定提升| +|SSD_300_VGG16_ATROUS_VOC| MXNet | VGG‑16 | 300×300 | Pascal VOC | 中等偏上 | 快 | VOC 数据集同类任务；资源受限时使用| +|SSD_512_MOBILENET1_VOC| MXNet | MobileNet‑1.0 | 512×512 | Pascal VOC | 中等 | 快 | 嵌入式/移动端设备；算力和内存都很有限| + + **YOLOv3 变体系列** |模型名称| 引擎 | 骨干网络 |
数据集
| 输入尺寸 |
精度
|
速度
|
适用场景
| @@ -486,6 +508,21 @@ SmartJavaAI是专为JAVA 开发者打造的一个功能丰富、开箱即用的 | SLANet_plus | （增强版）该模型通过轻量级骨干 PP-LCNet、CSP-PAN 融合与 SLA Head 解码，有效提升表格结构识别的精度与速度。 | [Github](https://paddlepaddle.github.io/PaddleOCR/v3.1.0/version3.x/module_usage/table_structure_recognition.html#_3) | +**车牌检测模型(License Plate Detection)** + +| 模型名称 | 模型简介 | 模型开源网站 | +|--------|----------------------|------------| +| YOLOV5 | 基于YOLOV5训练，支持12种中文车牌 | [Github](https://github.com/we0091234/Chinese_license_plate_detection_recognition) | +| YOLOV7 | 基于YOLOV7训练，支持12种中文车牌 | [Github](https://github.com/we0091234/Chinese_license_plate_detection_recognition) | + + +**车牌识别模型(License Plate Recognition)** + +| 模型名称 | 模型简介 | 模型开源网站 | +|--------|------------|------------| +| PLATE_REC_CRNN | CRNN中文字符识别 | [Github](https://github.com/Sierkinhane/CRNN_Chinese_Characters_Rec) | + + --- #### 机器翻译模型 @@ -519,7 +556,22 @@ SmartJavaAI是专为JAVA 开发者打造的一个功能丰富、开箱即用的 ## 近期更新日志 -## [v1.0.20] - 2025-07-06 +## [v1.0.22] - 2025-07-28 +- 新增 Milvus 身份验证支持 +- 集成车牌识别模型，支持车牌检测与识别 +- 目标检测功能升级：可指定类别及topk +- 支持自定义线程池线程数量 + + +## [v1.0.20] - 2025-07-18 +- OCR：新增表格识别模型 +- OCR：新增9个通用模型 +- OCR：支持批量检测识别 +- OCR：新增更多参数，使用更加灵活 +- 人脸识别：支持ID查询及分页获取人脸信息 +- 活体检测：视频检测支持设置最大帧数 + +## [v1.0.19] - 2025-07-06 - 人脸模块：新增小视科技（MiniVision）活体检测模型 - 人脸模块：新增阿里通义工作室活体检测模型 - 人脸模块：新增 2 个表情识别模型 diff --git a/examples/face-example/pom.xml b/examples/face-example/pom.xml index ff2d93e..40ed531 100644 --- a/examples/face-example/pom.xml +++ b/examples/face-example/pom.xml @@ -12,7 +12,7 @@ 11 11 UTF-8 - 1.0.20 + 1.0.22 smartai.examples.face.facedet.FaceDetDemo @@ -95,6 +95,13 @@ cn.smartjavaai smartjavaai-face + 1.0.22 + + + + cn.smartjavaai + smartjavaai-all + 1.0.22 diff --git a/examples/face-example/src/main/java/smartai/examples/face/VideoDemo.java b/examples/face-example/src/main/java/smartai/examples/face/VideoDemo.java new file mode 100644 index 0000000..64c94ff --- /dev/null +++ b/examples/face-example/src/main/java/smartai/examples/face/VideoDemo.java @@ -0,0 +1,25 @@ +package smartai.examples.face; + +import cn.smartjavaai.common.utils.VideoUtils; +import org.bytedeco.ffmpeg.global.avcodec; +import org.bytedeco.javacv.FFmpegFrameGrabber; +import org.bytedeco.javacv.FFmpegFrameRecorder; + +/** + * 视频预处理 + * @author dwj + * @date 2025/7/17 + */ +public class VideoDemo { + + public static void main(String[] args) { + try { + VideoUtils.rotateVideo("/Users/wenjie/Downloads/girl.mp4", "/Users/wenjie/Downloads/girl_rotate.mp4", 180,"mp4", avcodec.AV_CODEC_ID_H264); + } catch (FFmpegFrameRecorder.Exception e) { + throw new RuntimeException(e); + } catch (FFmpegFrameGrabber.Exception e) { + throw new RuntimeException(e); + } + } + +} diff --git a/examples/face-example/src/main/java/smartai/examples/face/attribute/FaceAttributeDetDemo.java b/examples/face-example/src/main/java/smartai/examples/face/attribute/FaceAttributeDetDemo.java index bb664aa..48f0c8e 100644 --- a/examples/face-example/src/main/java/smartai/examples/face/attribute/FaceAttributeDetDemo.java +++ b/examples/face-example/src/main/java/smartai/examples/face/attribute/FaceAttributeDetDemo.java @@ -57,7 +57,8 @@ public class FaceAttributeDetDemo { */ @Test public void testFaceAttributeDetect(){ - try (FaceAttributeModel faceAttributeModel = getFaceAttributeModel()){ + try { + FaceAttributeModel faceAttributeModel = getFaceAttributeModel(); DetectionResponse detectionResponse = faceAttributeModel.detect("src/main/resources/iu_1.jpg"); //绘制并导出人脸属性图片，小人脸仅有人脸框 BufferedImage image = ImageIO.read(new File(Paths.get("src/main/resources/iu_1.jpg").toAbsolutePath().toString())); @@ -73,7 +74,8 @@ public class FaceAttributeDetDemo { */ @Test public void testFaceAttributeDetect2(){ - try (FaceAttributeModel faceAttributeModel = getFaceAttributeModel()){ + try { + FaceAttributeModel faceAttributeModel = getFaceAttributeModel(); FaceAttribute faceAttribute = faceAttributeModel.detectTopFace("src/main/resources/iu_1.jpg"); log.info("人脸属性检测结果：{}", JSONObject.toJSONString(faceAttribute)); } catch (Exception e) { @@ -86,7 +88,8 @@ public class FaceAttributeDetDemo { */ @Test public void testFaceAttributeDetect3(){ - try (FaceAttributeModel faceAttributeModel = getFaceAttributeModel()){ + try { + FaceAttributeModel faceAttributeModel = getFaceAttributeModel(); FaceAttribute faceAttribute = faceAttributeModel.detectTopFace("src/main/resources/iu_1.jpg"); log.info("人脸属性检测结果：{}", JSONObject.toJSONString(faceAttribute)); } catch (Exception e) { @@ -103,8 +106,9 @@ public class FaceAttributeDetDemo { */ @Test public void testFaceAttributeDetect4(){ - try (FaceDetModel faceDetModel = getFaceDetModel(); - FaceAttributeModel faceAttributeModel = getFaceAttributeModel()){ + try { + FaceDetModel faceDetModel = getFaceDetModel(); + FaceAttributeModel faceAttributeModel = getFaceAttributeModel(); //人脸检测 BufferedImage image = ImageIO.read(new File(Paths.get("src/main/resources/iu_1.jpg").toAbsolutePath().toString())); R detectionResponse = faceDetModel.detect(image); diff --git a/examples/face-example/src/main/java/smartai/examples/face/expression/ExpressionRecDemo.java b/examples/face-example/src/main/java/smartai/examples/face/expression/ExpressionRecDemo.java index cabe47d..a187ea9 100644 --- a/examples/face-example/src/main/java/smartai/examples/face/expression/ExpressionRecDemo.java +++ b/examples/face-example/src/main/java/smartai/examples/face/expression/ExpressionRecDemo.java @@ -88,7 +88,8 @@ public class ExpressionRecDemo { */ @Test public void testExpressionDetect() { - try (ExpressionModel model = getExpressionModel()){ + try { + ExpressionModel model = getExpressionModel(); R result = model.detectTopFace("src/main/resources/emotion/happy.png"); if(result.isSuccess()){ log.info("识别结果：{}", JSONObject.toJSONString(result.getData().getExpression().getDescription())); @@ -106,7 +107,8 @@ public class ExpressionRecDemo { */ @Test public void testExpressionDetect2() { - try (ExpressionModel model = getExpressionModel()){ + try { + ExpressionModel model = getExpressionModel(); R result = model.detect("src/main/resources/emotion/happy.png"); if(result.isSuccess()){ //log.info("识别结果：{}", JSONObject.toJSONString(result.getData())); @@ -128,8 +130,9 @@ public class ExpressionRecDemo { */ @Test public void testExpressionDetect3() { - try (FaceDetModel faceDetModel = getFaceDetModel(); - ExpressionModel model = getExpressionModel()){ + try { + FaceDetModel faceDetModel = getFaceDetModel(); + ExpressionModel model = getExpressionModel(); BufferedImage image = ImageIO.read(new File(Paths.get("src/main/resources/emotion/happy.png").toAbsolutePath().toString())); R detResult = faceDetModel.detect(image); if(detResult.isSuccess()){ @@ -156,8 +159,9 @@ public class ExpressionRecDemo { */ @Test public void testExpressionDetect4() { - try (FaceDetModel faceDetModel = getFaceDetModel(); - ExpressionModel model = getExpressionModel()){ + try { + FaceDetModel faceDetModel = getFaceDetModel(); + ExpressionModel model = getExpressionModel(); BufferedImage image = ImageIO.read(new File(Paths.get("src/main/resources/emotion/happy.png").toAbsolutePath().toString())); R detResult = faceDetModel.detect(image); if(detResult.isSuccess()){ @@ -182,7 +186,8 @@ public class ExpressionRecDemo { */ @Test public void testExpressionDetectAndDraw(){ - try (ExpressionModel model = getExpressionModel()){ + try { + ExpressionModel model = getExpressionModel(); BufferedImage image = ImageIO.read(new File(Paths.get("src/main/resources/emotion/surprise.png").toAbsolutePath().toString())); R result = model.detect(image); if(result.isSuccess()){ @@ -206,7 +211,8 @@ public class ExpressionRecDemo { */ @Test public void testExpressionDetectCamera(){ - try (ExpressionModel expressionModel = getExpressionModel()){ + try { + ExpressionModel expressionModel = getExpressionModel(); OpenCV.loadShared(); VideoCapture capture = new VideoCapture(0); if (!capture.isOpened()) { diff --git a/examples/face-example/src/main/java/smartai/examples/face/facedet/FaceDetDemo.java b/examples/face-example/src/main/java/smartai/examples/face/facedet/FaceDetDemo.java index d03286e..2e084a0 100644 --- a/examples/face-example/src/main/java/smartai/examples/face/facedet/FaceDetDemo.java +++ b/examples/face-example/src/main/java/smartai/examples/face/facedet/FaceDetDemo.java @@ -82,7 +82,8 @@ public class FaceDetDemo { */ @Test public void testFaceDetect(){ - try (FaceDetModel faceModel = FaceDetModelFactory.getInstance().getModel()) { + try { + FaceDetModel faceModel = FaceDetModelFactory.getInstance().getModel(); R detectedResult = faceModel.detect(imgPath); if(detectedResult.isSuccess()){ log.info("人脸检测结果：{}", JSONObject.toJSONString(detectedResult.getData())); @@ -100,7 +101,8 @@ public class FaceDetDemo { */ @Test public void testFaceDetectCustomConfig(){ - try (FaceDetModel faceModel = getFaceDetModel()){ + try { + FaceDetModel faceModel = getFaceDetModel(); R detectedResult = faceModel.detect(imgPath); if(detectedResult.isSuccess()){ log.info("人脸检测结果：{}", JSONObject.toJSONString(detectedResult.getData())); @@ -118,7 +120,8 @@ public class FaceDetDemo { */ @Test public void testFaceDetectAndDraw(){ - try (FaceDetModel faceModel = getFaceDetModel()){ + try { + FaceDetModel faceModel = getFaceDetModel(); faceModel.detectAndDraw("src/main/resources/largest_selfie.jpg","output/largest_selfie_detected.png"); } catch (Exception e) { throw new RuntimeException(e); @@ -131,7 +134,8 @@ public class FaceDetDemo { */ @Test public void testFaceDetectAndDraw2(){ - try (FaceDetModel faceModel = getFaceDetModel()){ + try { + FaceDetModel faceModel = getFaceDetModel(); BufferedImage image = null; String imagePath = "src/main/resources/largest_selfie.jpg"; image = ImageIO.read(new File(Paths.get(imagePath).toAbsolutePath().toString())); @@ -154,11 +158,12 @@ public class FaceDetDemo { */ @Test public void testDetectFaceOffine(){ - FaceDetConfig config = new FaceDetConfig(); - config.setModelEnum(FaceDetModelEnum.RETINA_FACE);//人脸模型 - //模型路径,不同模型下载路径请参看文档 - config.setModelPath("/Users/xxx/Documents/develop/face_model/retinaface.pt"); - try (FaceDetModel faceModel = FaceDetModelFactory.getInstance().getModel(config)) { + try { + FaceDetConfig config = new FaceDetConfig(); + config.setModelEnum(FaceDetModelEnum.RETINA_FACE);//人脸模型 + //模型路径,不同模型下载路径请参看文档 + config.setModelPath("/Users/xxx/Documents/develop/face_model/retinaface.pt"); + FaceDetModel faceModel = FaceDetModelFactory.getInstance().getModel(config); R detectedResult = faceModel.detect(imgPath); if(detectedResult.isSuccess()){ log.info("人脸检测结果：{}", JSONObject.toJSONString(detectedResult.getData())); @@ -175,10 +180,11 @@ public class FaceDetDemo { */ @Test public void testDetectFaceGPU(){ - FaceDetConfig config = new FaceDetConfig(); - config.setModelEnum(FaceDetModelEnum.RETINA_FACE);//人脸模型 - config.setDevice(DeviceEnum.GPU); - try (FaceDetModel faceModel = FaceDetModelFactory.getInstance().getModel(config)) { + try { + FaceDetConfig config = new FaceDetConfig(); + config.setModelEnum(FaceDetModelEnum.RETINA_FACE);//人脸模型 + config.setDevice(DeviceEnum.GPU); + FaceDetModel faceModel = FaceDetModelFactory.getInstance().getModel(config); R detectedResult = faceModel.detect(imgPath); if(detectedResult.isSuccess()){ log.info("人脸检测结果：{}", JSONObject.toJSONString(detectedResult.getData())); @@ -196,7 +202,8 @@ public class FaceDetDemo { */ @Test public void testFaceDetectSeetaface6(){ - try (FaceDetModel faceModel = getSeetaface6DetModel()){ + try { + FaceDetModel faceModel = getSeetaface6DetModel(); R detectedResult = faceModel.detect(imgPath); if(detectedResult.isSuccess()){ log.info("人脸检测结果：{}", JSONObject.toJSONString(detectedResult.getData())); @@ -215,7 +222,8 @@ public class FaceDetDemo { */ @Test public void testDetectCamera(){ - try (FaceDetModel faceModel = getFaceDetModel()){ + try { + FaceDetModel faceModel = getFaceDetModel(); OpenCV.loadShared(); VideoCapture capture = new VideoCapture(0); if (!capture.isOpened()) { diff --git a/examples/face-example/src/main/java/smartai/examples/face/facerec/FaceRecDemo.java b/examples/face-example/src/main/java/smartai/examples/face/facerec/FaceRecDemo.java index 031c1e5..3cee699 100644 --- a/examples/face-example/src/main/java/smartai/examples/face/facerec/FaceRecDemo.java +++ b/examples/face-example/src/main/java/smartai/examples/face/facerec/FaceRecDemo.java @@ -95,6 +95,8 @@ public class FaceRecDemo { MilvusConfig vectorDBConfig = new MilvusConfig(); vectorDBConfig.setHost("127.0.0.1"); vectorDBConfig.setPort(19530); + //vectorDBConfig.setUsername("root"); + //vectorDBConfig.setPassword("Milvus"); //vectorDBConfig.setCollectionName("face5"); //ID策略：自动生成 vectorDBConfig.setIdStrategy(IdStrategy.AUTO); @@ -135,7 +137,8 @@ public class FaceRecDemo { */ @Test public void testExtractFeatures(){ - try (FaceRecModel faceRecModel = getFaceRecModel()){ + try { + FaceRecModel faceRecModel = getFaceRecModel(); //提取图片中所有人脸特征 R faceResult = faceRecModel.extractFeatures("src/main/resources/iu_1.jpg"); if(faceResult.isSuccess()){ @@ -158,7 +161,8 @@ public class FaceRecDemo { */ @Test public void featureComparison(){ - try (FaceRecModel faceRecModel = getFaceRecModel()){ + try { + FaceRecModel faceRecModel = getFaceRecModel(); //基于图像直接比对人脸特征 R similarResult = faceRecModel.featureComparison("src/main/resources/iu_1.jpg","src/main/resources/iu_2.jpg"); if(similarResult.isSuccess()){ @@ -183,7 +187,8 @@ public class FaceRecDemo { */ @Test public void featureComparison2(){ - try (FaceRecModel faceRecModel = getFaceRecModel()){ + try { + FaceRecModel faceRecModel = getFaceRecModel(); //特征提取（提取分数最高人脸特征）,适用于单人脸场景 R featureResult1 = faceRecModel.extractTopFaceFeature("src/main/resources/iu_1.jpg"); if(featureResult1.isSuccess()){ @@ -220,7 +225,8 @@ public class FaceRecDemo { */ @Test public void searchFace(){ - try (FaceRecModel faceRecModel = getFaceRecModelWithDbConfig()){ + try { + FaceRecModel faceRecModel = getFaceRecModelWithDbConfig(); //等待加载人脸库结束 while (!faceRecModel.isLoadFaceCompleted()){ Thread.sleep(100); @@ -296,7 +302,8 @@ public class FaceRecDemo { */ @Test public void searchFace2(){ - try (FaceRecModel faceRecModel = getFaceRecModelWithSQLiteConfig()){ + try { + FaceRecModel faceRecModel = getFaceRecModelWithSQLiteConfig(); //等待加载人脸库结束 while (!faceRecModel.isLoadFaceCompleted()){ Thread.sleep(100); @@ -367,7 +374,8 @@ public class FaceRecDemo { @Test public void getFaceInfo(){ //使用ID获取人脸信息 - try (FaceRecModel faceRecModel = getFaceRecModelWithSQLiteConfig()){ + try { + FaceRecModel faceRecModel = getFaceRecModelWithSQLiteConfig(); //等待加载人脸库结束 while (!faceRecModel.isLoadFaceCompleted()){ Thread.sleep(100); @@ -390,7 +398,8 @@ public class FaceRecDemo { @Test public void listFaces(){ //使用ID获取人脸信息 - try (FaceRecModel faceRecModel = getFaceRecModelWithDbConfig()){ + try { + FaceRecModel faceRecModel = getFaceRecModelWithDbConfig(); //等待加载人脸库结束 while (!faceRecModel.isLoadFaceCompleted()){ Thread.sleep(100); diff --git a/examples/face-example/src/main/java/smartai/examples/face/liveness/LivenessDetDemo.java b/examples/face-example/src/main/java/smartai/examples/face/liveness/LivenessDetDemo.java index bf1f966..b55ea7e 100644 --- a/examples/face-example/src/main/java/smartai/examples/face/liveness/LivenessDetDemo.java +++ b/examples/face-example/src/main/java/smartai/examples/face/liveness/LivenessDetDemo.java @@ -133,7 +133,8 @@ public class LivenessDetDemo { */ @Test public void testLivenessDetect(){ - try (LivenessDetModel livenessDetModel = getLivenessDetModel()){ + try { + LivenessDetModel livenessDetModel = getLivenessDetModel(); R response = livenessDetModel.detect("src/main/resources/liveness/1.jpg"); if(response.isSuccess()){ for (DetectionInfo detectionInfo : response.getData().getDetectionInfoList()){ @@ -152,7 +153,8 @@ public class LivenessDetDemo { */ @Test public void testLivenessDetectAndDraw(){ - try (LivenessDetModel livenessDetModel = getLivenessDetModel()){ + try { + LivenessDetModel livenessDetModel = getLivenessDetModel(); BufferedImage image = ImageIO.read(new File(Paths.get("src/main/resources/liveness/1.jpg").toAbsolutePath().toString())); R response = livenessDetModel.detect(image); if(response.isSuccess()){ @@ -175,7 +177,8 @@ public class LivenessDetDemo { */ @Test public void testLivenessDetect2(){ - try (LivenessDetModel livenessDetModel = getLivenessDetModel()){ + try { + LivenessDetModel livenessDetModel = getLivenessDetModel(); //指定文件夹路径 File dir = new File("face-example/src/main/resources/liveness"); File[] files = dir.listFiles(); @@ -198,8 +201,9 @@ public class LivenessDetDemo { */ @Test public void testLivenessDetect3(){ - try (FaceDetModel faceDetectModel = getFaceDetModel(); - LivenessDetModel livenessDetModel = getLivenessDetModel()){ + try { + FaceDetModel faceDetectModel = getFaceDetModel(); + LivenessDetModel livenessDetModel = getLivenessDetModel(); // 将图片路径转换为 BufferedImage BufferedImage image = ImageIO.read(new File(Paths.get("src/main/resources/liveness/1.jpg").toAbsolutePath().toString())); //人脸检测 @@ -230,8 +234,9 @@ public class LivenessDetDemo { */ @Test public void testLivenessDetect4(){ - try (FaceDetModel faceDetModel = getFaceDetModel(); - LivenessDetModel livenessDetModel = getMiniVisionLivenessDetModel()){ + try { + FaceDetModel faceDetModel = getFaceDetModel(); + LivenessDetModel livenessDetModel = getMiniVisionLivenessDetModel(); // 将图片路径转换为 BufferedImage BufferedImage image = ImageIO.read(new File(Paths.get("src/main/resources/liveness/1.jpg").toAbsolutePath().toString())); R detResult = faceDetModel.detect(image); @@ -259,7 +264,8 @@ public class LivenessDetDemo { */ @Test public void testLivenessDetectVideo(){ - try (LivenessDetModel livenessDetModel = getLivenessDetModel()){ + try { + LivenessDetModel livenessDetModel = getLivenessDetModel(); //视频路径 R livenessStatus = livenessDetModel.detectVideo("video.mp4"); if (livenessStatus.isSuccess()){ @@ -278,7 +284,8 @@ public class LivenessDetDemo { */ @Test public void testLivenessDetectCamera(){ - try (LivenessDetModel livenessDetModel = getLivenessDetModel()){ + try { + LivenessDetModel livenessDetModel = getLivenessDetModel(); OpenCV.loadShared(); VideoCapture capture = new VideoCapture(0); if (!capture.isOpened()) { diff --git a/examples/face-example/src/main/java/smartai/examples/face/quality/FaceQualityDetDemo.java b/examples/face-example/src/main/java/smartai/examples/face/quality/FaceQualityDetDemo.java index 01a7351..714b59e 100644 --- a/examples/face-example/src/main/java/smartai/examples/face/quality/FaceQualityDetDemo.java +++ b/examples/face-example/src/main/java/smartai/examples/face/quality/FaceQualityDetDemo.java @@ -78,8 +78,9 @@ public class FaceQualityDetDemo { */ @Test public void evaluateBrightness(){ - try (FaceQualityModel faceQualityModel = getFaceQualityModel(); - FaceDetModel faceDetModel = getFaceDetModel()){ + try { + FaceQualityModel faceQualityModel = getFaceQualityModel(); + FaceDetModel faceDetModel = getFaceDetModel(); //人脸检测 BufferedImage image = ImageIO.read(new File(Paths.get("src/main/resources/iu_1.jpg").toAbsolutePath().toString())); R detectionResponse = faceDetModel.detect(image); @@ -110,8 +111,9 @@ public class FaceQualityDetDemo { */ @Test public void evaluateCompleteness(){ - try (FaceQualityModel faceQualityModel = getFaceQualityModel(); - FaceDetModel faceDetModel = getFaceDetModel()){ + try { + FaceQualityModel faceQualityModel = getFaceQualityModel(); + FaceDetModel faceDetModel = getFaceDetModel(); //人脸检测 BufferedImage image = ImageIO.read(new File(Paths.get("src/main/resources/iu_1.jpg").toAbsolutePath().toString())); R detectionResponse = faceDetModel.detect(image); @@ -142,8 +144,9 @@ public class FaceQualityDetDemo { */ @Test public void evaluateClarity(){ - try (FaceQualityModel faceQualityModel = getFaceQualityModel(); - FaceDetModel faceDetModel = getFaceDetModel()){ + try { + FaceQualityModel faceQualityModel = getFaceQualityModel(); + FaceDetModel faceDetModel = getFaceDetModel(); //人脸检测 BufferedImage image = ImageIO.read(new File(Paths.get("src/main/resources/iu_1.jpg").toAbsolutePath().toString())); R detectionResponse = faceDetModel.detect(image); @@ -174,8 +177,9 @@ public class FaceQualityDetDemo { */ @Test public void evaluatePose(){ - try (FaceQualityModel faceQualityModel = getFaceQualityModel(); - FaceDetModel faceDetModel = getFaceDetModel()){ + try { + FaceQualityModel faceQualityModel = getFaceQualityModel(); + FaceDetModel faceDetModel = getFaceDetModel(); //人脸检测 BufferedImage image = ImageIO.read(new File(Paths.get("src/main/resources/iu_1.jpg").toAbsolutePath().toString())); R detectionResponse = faceDetModel.detect(image); @@ -207,8 +211,9 @@ public class FaceQualityDetDemo { */ @Test public void evaluateResolution(){ - try (FaceQualityModel faceQualityModel = getFaceQualityModel(); - FaceDetModel faceDetModel = getFaceDetModel()){ + try { + FaceQualityModel faceQualityModel = getFaceQualityModel(); + FaceDetModel faceDetModel = getFaceDetModel(); //人脸检测 BufferedImage image = ImageIO.read(new File(Paths.get("src/main/resources/iu_1.jpg").toAbsolutePath().toString())); R detectionResponse = faceDetModel.detect(image); @@ -240,8 +245,9 @@ public class FaceQualityDetDemo { */ @Test public void evaluateAll(){ - try (FaceQualityModel faceQualityModel = getFaceQualityModel(); - FaceDetModel faceDetModel = getFaceDetModel()){ + try { + FaceQualityModel faceQualityModel = getFaceQualityModel(); + FaceDetModel faceDetModel = getFaceDetModel(); //人脸检测 BufferedImage image = ImageIO.read(new File(Paths.get("src/main/resources/iu_1.jpg").toAbsolutePath().toString())); R detectionResponse = faceDetModel.detect(image); diff --git a/examples/objectdetection-example/pom.xml b/examples/objectdetection-example/pom.xml index 7d3cdce..b879a77 100644 --- a/examples/objectdetection-example/pom.xml +++ b/examples/objectdetection-example/pom.xml @@ -12,7 +12,7 @@ 11 11 UTF-8 - 1.0.20 + 1.0.22 smartai.examples.objectdetection.ObjectDetection diff --git a/examples/objectdetection-example/src/main/java/smartai/examples/objectdetection/ObjectDetection.java b/examples/objectdetection-example/src/main/java/smartai/examples/objectdetection/ObjectDetection.java index 11a295b..14be538 100644 --- a/examples/objectdetection-example/src/main/java/smartai/examples/objectdetection/ObjectDetection.java +++ b/examples/objectdetection-example/src/main/java/smartai/examples/objectdetection/ObjectDetection.java @@ -41,10 +41,8 @@ import java.awt.image.BufferedImage; import java.io.File; import java.io.IOException; import java.nio.file.Paths; -import java.util.ArrayList; -import java.util.Iterator; +import java.util.*; import java.util.List; -import java.util.Objects; import java.util.concurrent.Callable; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -71,7 +69,8 @@ public class ObjectDetection { @Test public void objectDetection(){ //默认cpu - try (DetectorModel detectorModel = ObjectDetectionModelFactory.getInstance().getModel()){ + try { + DetectorModel detectorModel = ObjectDetectionModelFactory.getInstance().getModel(); DetectionResponse detectionResponse = detectorModel.detect("src/main/resources/object_detection.jpg"); log.info("目标检测结果：{}", JSONObject.toJSONString(detectionResponse)); } catch (Exception e) { @@ -84,10 +83,15 @@ public class ObjectDetection { */ @Test public void objectDetection2(){ - DetectorModelConfig config = new DetectorModelConfig(); - config.setModelEnum(DetectorModelEnum.SSD_300_RESNET50);//检测模型，目前支持19种预置模型 - config.setDevice(device); - try (DetectorModel detectorModel = ObjectDetectionModelFactory.getInstance().getModel(config)){ + try { + DetectorModelConfig config = new DetectorModelConfig(); + config.setModelEnum(DetectorModelEnum.SSD_300_RESNET50);//检测模型，目前支持19种预置模型 + // 指定允许的类别 +// config.setAllowedClasses(Arrays.asList("person")); + //指定返回检测数量 + config.setTopK(100); + config.setDevice(device); + DetectorModel detectorModel = ObjectDetectionModelFactory.getInstance().getModel(config); DetectionResponse detectionResponse = detectorModel.detect("src/main/resources/dog_bike_car.jpg"); log.info("目标检测结果：{}", JSONObject.toJSONString(detectionResponse)); } catch (Exception e) { @@ -100,7 +104,8 @@ public class ObjectDetection { */ @Test public void objectDetectionAndDraw(){ - try (DetectorModel detectorModel = ObjectDetectionModelFactory.getInstance().getModel()){ + try { + DetectorModel detectorModel = ObjectDetectionModelFactory.getInstance().getModel(); detectorModel.detectAndDraw("src/main/resources/object_detection.jpg","output/object_detection_detected.png"); } catch (Exception e) { e.printStackTrace(); @@ -112,7 +117,8 @@ public class ObjectDetection { */ @Test public void objectDetectionAndDraw2(){ - try (DetectorModel detectorModel = ObjectDetectionModelFactory.getInstance().getModel()){ + try { + DetectorModel detectorModel = ObjectDetectionModelFactory.getInstance().getModel(); String imagePath = "src/main/resources/object_detection.jpg"; BufferedImage image = ImageIO.read(new File(Paths.get(imagePath).toAbsolutePath().toString())); //可以根据后续业务场景使用detectedImage @@ -131,15 +137,16 @@ public class ObjectDetection { */ @Test public void objectDetectionWithOfficialModel(){ - DetectorModelConfig config = new DetectorModelConfig(); - config.setThreshold(0.3f); - //也支持YoloV8：YOLOV8_OFFICIAL 模型可以从文档中提供的地址下载 - config.setModelEnum(DetectorModelEnum.YOLOV12_OFFICIAL);//检测模型，目前支持19种模型 - // 指定模型路径，需要更改为自己的模型路径 - config.setModelPath("/Users/xxx/Documents/yolov12n.onnx"); - config.setDevice(device); - //一定要将yolo官方的类别文件：synset.txt（文档中下载）放在模型同目录下，否则报错 - try (DetectorModel detectorModel = ObjectDetectionModelFactory.getInstance().getModel(config)){ + try { + DetectorModelConfig config = new DetectorModelConfig(); + config.setThreshold(0.3f); + //也支持YoloV8：YOLOV8_OFFICIAL 模型可以从文档中提供的地址下载 + config.setModelEnum(DetectorModelEnum.YOLOV12_OFFICIAL);//检测模型，目前支持19种模型 + // 指定模型路径，需要更改为自己的模型路径 + config.setModelPath("/Users/xxx/Documents/yolov12n.onnx"); + config.setDevice(device); + //一定要将yolo官方的类别文件：synset.txt（文档中下载）放在模型同目录下，否则报错 + DetectorModel detectorModel = ObjectDetectionModelFactory.getInstance().getModel(config); DetectionResponse detect = detectorModel.detect("src/main/resources/dog_bike_car.jpg"); log.info("目标检测结果：{}", JSONObject.toJSONString(detect)); } catch (Exception e) { @@ -152,16 +159,21 @@ public class ObjectDetection { */ @Test public void objectDetectionWithCustomModel(){ - DetectorModelConfig config = new DetectorModelConfig(); - //也支持YoloV8：YOLOV8_CUSTOM 模型需要自己训练，训练教程可以查看文档 - config.setModelEnum(DetectorModelEnum.YOLOV12_CUSTOM);//自定义YOLOV12模型 - // 指定模型路径，需要更改为自己的模型路径 - config.setModelPath("/Users/xxx/Documents/develop/fire_model/best.onnx"); - config.putCustomParam("width", 640);//resize 宽 - config.putCustomParam("height", 640);// resize 高 - config.putCustomParam("nmsThreshold", 0.5f); - config.setDevice(device); - try (DetectorModel detectorModel = ObjectDetectionModelFactory.getInstance().getModel(config)){ + try { + DetectorModelConfig config = new DetectorModelConfig(); + //也支持YoloV8：YOLOV8_CUSTOM 模型需要自己训练，训练教程可以查看文档 + config.setModelEnum(DetectorModelEnum.YOLOV12_CUSTOM);//自定义YOLOV12模型 + // 指定模型路径，需要更改为自己的模型路径 + config.setModelPath("/Users/xxx/Documents/develop/fire_model/best.onnx"); + config.putCustomParam("width", 640);//resize 宽 + config.putCustomParam("height", 640);// resize 高 + config.putCustomParam("nmsThreshold", 0.5f); + // 指定允许的类别 +// config.setAllowedClasses(Arrays.asList("person")); + //指定返回检测数量 + config.setTopK(100); + config.setDevice(device); + DetectorModel detectorModel = ObjectDetectionModelFactory.getInstance().getModel(config); DetectionResponse detect = detectorModel.detect("src/main/resources/dog_bike_car.jpg"); log.info("目标检测结果：{}", JSONObject.toJSONString(detect)); } catch (Exception e) { @@ -176,7 +188,8 @@ public class ObjectDetection { */ @Test public void testDetectCamera(){ - try (DetectorModel detectorModel = ObjectDetectionModelFactory.getInstance().getModel()){ + try { + DetectorModel detectorModel = ObjectDetectionModelFactory.getInstance().getModel(); OpenCV.loadShared(); VideoCapture capture = new VideoCapture(0); if (!capture.isOpened()) { diff --git a/examples/ocr-examples/pom.xml b/examples/ocr-examples/pom.xml index 3506826..f89411e 100644 --- a/examples/ocr-examples/pom.xml +++ b/examples/ocr-examples/pom.xml @@ -12,7 +12,7 @@ 11 11 UTF-8 - 1.0.20 + 1.0.22 smartai.examples.ocr.common.OcrRecognizeDemo diff --git a/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrDetectionDemo.java b/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrDetectionDemo.java index e77a31e..e64cc30 100644 --- a/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrDetectionDemo.java +++ b/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrDetectionDemo.java @@ -11,7 +11,6 @@ import cn.smartjavaai.ocr.entity.OcrInfo; import cn.smartjavaai.ocr.enums.CommonDetModelEnum; import cn.smartjavaai.ocr.factory.OcrModelFactory; import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel; -import cn.smartjavaai.ocr.opencv.OcrOpenCVUtils; import cn.smartjavaai.ocr.utils.OcrUtils; import com.alibaba.fastjson.JSONObject; import lombok.extern.slf4j.Slf4j; @@ -60,7 +59,8 @@ public class OcrDetectionDemo { */ @Test public void detect(){ - try (OcrCommonDetModel model = getDetectionModel()){ + try { + OcrCommonDetModel model = getDetectionModel(); List boxes = model.detect("src/main/resources/ocr_1.jpg"); log.info("OCR检测结果：{}", JSONObject.toJSONString(boxes)); } catch (Exception e) { @@ -77,7 +77,8 @@ public class OcrDetectionDemo { */ @Test public void detectAndDraw(){ - try (OcrCommonDetModel model = getDetectionModel()){ + try { + OcrCommonDetModel model = getDetectionModel(); model.detectAndDraw("src/main/resources/ocr_1.jpg", "output/ocr_1_detected.jpg"); } catch (Exception e) { e.printStackTrace(); @@ -94,7 +95,8 @@ public class OcrDetectionDemo { */ @Test public void batchDetect(){ - try (OcrCommonDetModel model = getDetectionModel()){ + try { + OcrCommonDetModel model = getDetectionModel(); //批量检测要求图片宽高一致 String folderPath = "/Users/xxx/Downloads/testing33"; //读取文件夹中所有图片 diff --git a/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrDirectionDetDemo.java b/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrDirectionDetDemo.java index d9264fa..5e93842 100644 --- a/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrDirectionDetDemo.java +++ b/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrDirectionDetDemo.java @@ -67,7 +67,8 @@ public class OcrDirectionDetDemo { */ @Test public void detect(){ - try (OcrDirectionModel directionModel = getDirectionModel()){ + try { + OcrDirectionModel directionModel = getDirectionModel(); List itemList = directionModel.detect("src/main/resources/ocr_1.jpg"); log.info("OCR方向检测结果1：{}", JSONObject.toJSONString(itemList)); } catch (Exception e) { @@ -84,7 +85,8 @@ public class OcrDirectionDetDemo { */ @Test public void detectAndDraw(){ - try (OcrDirectionModel directionModel = getDirectionModel()){ + try { + OcrDirectionModel directionModel = getDirectionModel(); directionModel.detectAndDraw("src/main/resources/ocr_3.jpg", "output/ocr_3_detected.png"); } catch (Exception e) { e.printStackTrace(); diff --git a/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrRecognizeDemo.java b/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrRecognizeDemo.java index 037a278..263d5bd 100644 --- a/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrRecognizeDemo.java +++ b/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrRecognizeDemo.java @@ -106,7 +106,8 @@ public class OcrRecognizeDemo { */ @Test public void recognize(){ - try (OcrCommonRecModel recModel = getRecModel()){ + try { + OcrCommonRecModel recModel = getRecModel(); //不带方向矫正，分行返回文本 OcrRecOptions options = new OcrRecOptions(false, true); OcrInfo ocrInfo = recModel.recognize("src/main/resources/ocr_2.jpg",options); @@ -127,7 +128,8 @@ public class OcrRecognizeDemo { */ @Test public void recognizeHandWriting(){ - try (OcrCommonRecModel recModel = getRecModel()){ + try { + OcrCommonRecModel recModel = getRecModel(); OcrInfo ocrInfo = recModel.recognize("src/main/resources/handwriting_1.jpg",new OcrRecOptions()); log.info("OCR识别结果：{}", JSONObject.toJSONString(ocrInfo)); } catch (Exception e) { @@ -146,7 +148,8 @@ public class OcrRecognizeDemo { */ @Test public void recognize2(){ - try (OcrCommonRecModel recModel = getRecModelWithDirection()){ + try { + OcrCommonRecModel recModel = getRecModelWithDirection(); //带方向矫正，分行返回文本 OcrRecOptions options = new OcrRecOptions(true, true); OcrInfo ocrInfo = recModel.recognize("src/main/resources/ocr_3.jpg",options); @@ -168,7 +171,8 @@ public class OcrRecognizeDemo { */ @Test public void recognizeAndDraw(){ - try (OcrCommonRecModel recModel = getRecModelWithDirection()){ + try { + OcrCommonRecModel recModel = getRecModelWithDirection(); int fontSize = 18; recModel.recognizeAndDraw("src/main/resources/general_ocr_002.png", "output/ocr_4_recognized.jpg", fontSize, new OcrRecOptions()); } catch (Exception e) { @@ -184,7 +188,8 @@ public class OcrRecognizeDemo { */ @Test public void batchRecognize(){ - try (OcrCommonRecModel recModel = getRecModelWithDirection()){ + try { + OcrCommonRecModel recModel = getRecModelWithDirection(); //批量检测要求图片宽高一致 String folderPath = "/Users/xxx/Downloads/testing33"; //读取文件夹中所有图片 diff --git a/examples/ocr-examples/src/main/java/smartai/examples/ocr/plate/PlateRecDemo.java b/examples/ocr-examples/src/main/java/smartai/examples/ocr/plate/PlateRecDemo.java new file mode 100644 index 0000000..55137f4 --- /dev/null +++ b/examples/ocr-examples/src/main/java/smartai/examples/ocr/plate/PlateRecDemo.java @@ -0,0 +1,81 @@ +package smartai.examples.ocr.plate; + +import ai.djl.util.JsonUtils; +import cn.smartjavaai.common.entity.R; +import cn.smartjavaai.common.enums.DeviceEnum; +import cn.smartjavaai.common.utils.ImageUtils; +import cn.smartjavaai.ocr.config.PlateDetModelConfig; +import cn.smartjavaai.ocr.config.PlateRecModelConfig; +import cn.smartjavaai.ocr.entity.PlateInfo; +import cn.smartjavaai.ocr.enums.PlateDetModelEnum; +import cn.smartjavaai.ocr.enums.PlateRecModelEnum; +import cn.smartjavaai.ocr.factory.PlateModelFactory; +import cn.smartjavaai.ocr.model.plate.PlateDetModel; +import cn.smartjavaai.ocr.model.plate.PlateRecModel; +import lombok.extern.slf4j.Slf4j; +import org.junit.Test; + +import java.io.File; +import java.util.List; + +/** + * @author dwj + */ +@Slf4j +public class PlateRecDemo { + + //设备类型 + public static DeviceEnum device = DeviceEnum.CPU; + + /** + * 获取车牌检测模型 + * @return + */ + public PlateDetModel getPlateDetModel() { + PlateDetModelConfig config = new PlateDetModelConfig(); + config.setModelEnum(PlateDetModelEnum.YOLOV5); + config.setModelPath("/Users/wenjie/Downloads/modelscope-2b3747db53adb48370c6edaccac56eb1ca0da1b5/onnx_export/model.onnx"); +// config.setModelPath("/Users/wenjie/Documents/develop/model/plate/yolov8s.onnx"); + config.setModelPath("/Users/wenjie/Downloads/数据集/车牌/onnx/plate_detect.onnx"); + config.setPredictorPoolSize(3); + config.setDevice(device); + return PlateModelFactory.getInstance().getDetModel(config); + } + + /** + * 获取车牌识别模型 + * @return + */ + public PlateRecModel getPlateRecModel() { + PlateRecModelConfig recModelConfig = new PlateRecModelConfig(); + recModelConfig.setModelEnum(PlateRecModelEnum.PLATE_REC_CRNN); + recModelConfig.setModelPath("/Users/wenjie/Downloads/数据集/车牌/onnx/plate_rec_color.onnx"); + recModelConfig.setPlateDetModel(getPlateDetModel()); + return PlateModelFactory.getInstance().getRecModel(recModelConfig); + } + + @Test + public void testDetect() { + PlateRecModel plateRecModel = getPlateRecModel(); + R> result = plateRecModel.recognize("src/main/resources/plate/Quicker_20220930_180856.png"); + if(result.isSuccess()){ + log.info("车牌识别结果：{}", JsonUtils.toJson(result.getData())); + }else{ + log.error("车牌识别失败：{}", result.getMessage()); + } + } + + @Test + public void recognizeAndDraw() { + PlateRecModel plateRecModel = getPlateRecModel(); + R result = plateRecModel.recognizeAndDraw("src/main/resources/plate/single_green.jpg", "output/plate_recognized2.jpg"); + if(result.isSuccess()){ + log.info("车牌识别成功"); + }else{ + log.error("车牌识别失败：{}", result.getMessage()); + } + + } + + +} diff --git a/examples/ocr-examples/src/main/java/smartai/examples/ocr/table/TableRecDemo.java b/examples/ocr-examples/src/main/java/smartai/examples/ocr/table/TableRecDemo.java index 8b41f9b..0aba097 100644 --- a/examples/ocr-examples/src/main/java/smartai/examples/ocr/table/TableRecDemo.java +++ b/examples/ocr-examples/src/main/java/smartai/examples/ocr/table/TableRecDemo.java @@ -113,10 +113,11 @@ public class TableRecDemo { */ @Test public void recognize(){ - try (TableStructureModel tableStructureModel = getTableStructureModel(); - OcrCommonDetModel detModel = getDetectionModel(); - OcrCommonRecModel recModel = getRecModel(); - OcrDirectionModel directionModel = getDirectionModel()){ + try { + TableStructureModel tableStructureModel = getTableStructureModel(); + OcrCommonDetModel detModel = getDetectionModel(); + OcrCommonRecModel recModel = getRecModel(); + OcrDirectionModel directionModel = getDirectionModel(); //创建表格识别器 TableRecognizer tableRecognizer = TableRecognizer.builder() .withStructureModel(tableStructureModel) diff --git a/examples/ocr-examples/src/main/resources/plate/Quicker_20220930_180856.png b/examples/ocr-examples/src/main/resources/plate/Quicker_20220930_180856.png new file mode 100644 index 0000000..eed5017 Binary files /dev/null and b/examples/ocr-examples/src/main/resources/plate/Quicker_20220930_180856.png differ diff --git a/examples/ocr-examples/src/main/resources/plate/Quicker_20220930_180919.png b/examples/ocr-examples/src/main/resources/plate/Quicker_20220930_180919.png new file mode 100644 index 0000000..42abc0f Binary files /dev/null and b/examples/ocr-examples/src/main/resources/plate/Quicker_20220930_180919.png differ diff --git a/examples/ocr-examples/src/main/resources/plate/Quicker_20220930_180938.png b/examples/ocr-examples/src/main/resources/plate/Quicker_20220930_180938.png new file mode 100644 index 0000000..1e34d8c Binary files /dev/null and b/examples/ocr-examples/src/main/resources/plate/Quicker_20220930_180938.png differ diff --git a/examples/ocr-examples/src/main/resources/plate/Quicker_20220930_181044.png b/examples/ocr-examples/src/main/resources/plate/Quicker_20220930_181044.png new file mode 100644 index 0000000..cde51c1 Binary files /dev/null and b/examples/ocr-examples/src/main/resources/plate/Quicker_20220930_181044.png differ diff --git a/examples/ocr-examples/src/main/resources/plate/double_yellow.jpg b/examples/ocr-examples/src/main/resources/plate/double_yellow.jpg new file mode 100644 index 0000000..b42b61c Binary files /dev/null and b/examples/ocr-examples/src/main/resources/plate/double_yellow.jpg differ diff --git a/examples/ocr-examples/src/main/resources/plate/hongkang1.jpg b/examples/ocr-examples/src/main/resources/plate/hongkang1.jpg new file mode 100644 index 0000000..23099b1 Binary files /dev/null and b/examples/ocr-examples/src/main/resources/plate/hongkang1.jpg differ diff --git a/examples/ocr-examples/src/main/resources/plate/moto.png b/examples/ocr-examples/src/main/resources/plate/moto.png new file mode 100644 index 0000000..c0d49db Binary files /dev/null and b/examples/ocr-examples/src/main/resources/plate/moto.png differ diff --git a/examples/ocr-examples/src/main/resources/plate/police.jpg b/examples/ocr-examples/src/main/resources/plate/police.jpg new file mode 100644 index 0000000..9ac0bb1 Binary files /dev/null and b/examples/ocr-examples/src/main/resources/plate/police.jpg differ diff --git a/examples/ocr-examples/src/main/resources/plate/shi_lin_guan.jpg b/examples/ocr-examples/src/main/resources/plate/shi_lin_guan.jpg new file mode 100644 index 0000000..f7c216d Binary files /dev/null and b/examples/ocr-examples/src/main/resources/plate/shi_lin_guan.jpg differ diff --git a/examples/ocr-examples/src/main/resources/plate/single_blue.jpg b/examples/ocr-examples/src/main/resources/plate/single_blue.jpg new file mode 100644 index 0000000..28abf64 Binary files /dev/null and b/examples/ocr-examples/src/main/resources/plate/single_blue.jpg differ diff --git a/examples/ocr-examples/src/main/resources/plate/single_green.jpg b/examples/ocr-examples/src/main/resources/plate/single_green.jpg new file mode 100644 index 0000000..0a0e0a9 Binary files /dev/null and b/examples/ocr-examples/src/main/resources/plate/single_green.jpg differ diff --git a/examples/ocr-examples/src/main/resources/plate/single_yellow.jpg b/examples/ocr-examples/src/main/resources/plate/single_yellow.jpg new file mode 100644 index 0000000..ebe3a26 Binary files /dev/null and b/examples/ocr-examples/src/main/resources/plate/single_yellow.jpg differ diff --git a/examples/ocr-examples/src/main/resources/plate/tmp8F1F.png b/examples/ocr-examples/src/main/resources/plate/tmp8F1F.png new file mode 100644 index 0000000..aeaa7fa Binary files /dev/null and b/examples/ocr-examples/src/main/resources/plate/tmp8F1F.png differ diff --git a/examples/ocr-examples/src/main/resources/plate/tmpA5E3.png b/examples/ocr-examples/src/main/resources/plate/tmpA5E3.png new file mode 100644 index 0000000..b75872b Binary files /dev/null and b/examples/ocr-examples/src/main/resources/plate/tmpA5E3.png differ diff --git a/examples/ocr-examples/src/main/resources/plate/xue.jpg b/examples/ocr-examples/src/main/resources/plate/xue.jpg new file mode 100644 index 0000000..d583ac7 Binary files /dev/null and b/examples/ocr-examples/src/main/resources/plate/xue.jpg differ diff --git a/examples/ocr-examples/src/main/resources/table/medal_table.png b/examples/ocr-examples/src/main/resources/table/medal_table.png new file mode 100644 index 0000000..08e025d Binary files /dev/null and b/examples/ocr-examples/src/main/resources/table/medal_table.png differ diff --git a/examples/ocr-examples/src/main/resources/table/table_ch1.png b/examples/ocr-examples/src/main/resources/table/table_ch1.png new file mode 100644 index 0000000..16745bd Binary files /dev/null and b/examples/ocr-examples/src/main/resources/table/table_ch1.png differ diff --git a/examples/translation-example/pom.xml b/examples/translation-example/pom.xml index ec76a11..18f8eef 100644 --- a/examples/translation-example/pom.xml +++ b/examples/translation-example/pom.xml @@ -12,7 +12,7 @@ 11 11 UTF-8 - 1.0.20 + 1.0.22 smartai.examples.nlp.translation.TranslationDemo diff --git a/pom.xml b/pom.xml index a78b7ba..4142bfb 100644 --- a/pom.xml +++ b/pom.xml @@ -7,7 +7,7 @@ SmartJavaAI cn.smartjavaai smartjavaai-parent - 1.0.20 + 1.0.22 pom SmartJavaAI diff --git a/smartjavaai-all/pom.xml b/smartjavaai-all/pom.xml index e66d60c..e9c35f1 100644 --- a/smartjavaai-all/pom.xml +++ b/smartjavaai-all/pom.xml @@ -6,11 +6,11 @@ cn.smartjavaai smartjavaai-parent - 1.0.20 + 1.0.22 smartjavaai-all - 1.0.20 + 1.0.22 ${project.artifactId} SmartJavaAI https://github.com/geekwenjie/SmartJavaAI diff --git a/smartjavaai-bom/pom.xml b/smartjavaai-bom/pom.xml index d841897..1357965 100644 --- a/smartjavaai-bom/pom.xml +++ b/smartjavaai-bom/pom.xml @@ -6,10 +6,10 @@ cn.smartjavaai smartjavaai-parent - 1.0.20 + 1.0.22 - 1.0.20 + 1.0.22 smartjavaai-bom smartjavaai-bom 统一版本管理的 BOM 包，同时支持 import 和全量依赖 diff --git a/smartjavaai-common/pom.xml b/smartjavaai-common/pom.xml index 1314870..421744a 100644 --- a/smartjavaai-common/pom.xml +++ b/smartjavaai-common/pom.xml @@ -6,7 +6,7 @@ cn.smartjavaai smartjavaai-parent - 1.0.20 + 1.0.22 smartjavaai-common diff --git a/smartjavaai-common/src/main/java/cn/smartjavaai/common/config/ModelConfig.java b/smartjavaai-common/src/main/java/cn/smartjavaai/common/config/ModelConfig.java index fdbdfe7..14a1628 100644 --- a/smartjavaai-common/src/main/java/cn/smartjavaai/common/config/ModelConfig.java +++ b/smartjavaai-common/src/main/java/cn/smartjavaai/common/config/ModelConfig.java @@ -29,6 +29,11 @@ public class ModelConfig { */ private String batchifier; + /** + * 模型预测器池大小(默认为cpu核心数) + */ + private int predictorPoolSize; + /** * 个性化配置（按模型类型动态解析） */ diff --git a/smartjavaai-common/src/main/java/cn/smartjavaai/common/entity/R.java b/smartjavaai-common/src/main/java/cn/smartjavaai/common/entity/R.java index 30cc3d3..511a836 100644 --- a/smartjavaai-common/src/main/java/cn/smartjavaai/common/entity/R.java +++ b/smartjavaai-common/src/main/java/cn/smartjavaai/common/entity/R.java @@ -57,6 +57,7 @@ public class R { NO_FACE_DETECTED(3, "未检测到人脸"), PARAM_ERROR(4, "参数错误"), INVALID_VIDEO(5, "视频无效"), + NO_OBJECT_DETECTED(6, "未检测到目标"), Unknown(-1, "未知错误"); private final int code; diff --git a/smartjavaai-common/src/main/java/cn/smartjavaai/common/pool/PredictorFactory.java b/smartjavaai-common/src/main/java/cn/smartjavaai/common/pool/PredictorFactory.java index b71e513..6af4459 100644 --- a/smartjavaai-common/src/main/java/cn/smartjavaai/common/pool/PredictorFactory.java +++ b/smartjavaai-common/src/main/java/cn/smartjavaai/common/pool/PredictorFactory.java @@ -3,6 +3,7 @@ package cn.smartjavaai.common.pool; import ai.djl.inference.Predictor; import ai.djl.repository.zoo.ZooModel; import ai.djl.translate.Translator; +import lombok.extern.slf4j.Slf4j; import org.apache.commons.pool2.BasePooledObjectFactory; import org.apache.commons.pool2.PooledObject; import org.apache.commons.pool2.impl.DefaultPooledObject; @@ -12,6 +13,7 @@ import org.apache.commons.pool2.impl.DefaultPooledObject; * @author dwj * @date 2025/4/8 */ +@Slf4j public class PredictorFactory extends BasePooledObjectFactory> { private final ZooModel model; @@ -21,6 +23,7 @@ public class PredictorFactory extends BasePooledObjectFactory create() { + log.debug("create predictor"); return model.newPredictor(); } @@ -31,6 +34,7 @@ public class PredictorFactory extends BasePooledObjectFactory> p) { + log.debug("close predictor"); p.getObject().close(); } } diff --git a/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/ImageUtils.java b/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/ImageUtils.java index 1a3c75c..e71e281 100644 --- a/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/ImageUtils.java +++ b/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/ImageUtils.java @@ -8,6 +8,9 @@ import ai.djl.ndarray.NDArray; import cn.smartjavaai.common.entity.DetectionRectangle; import cn.smartjavaai.common.entity.DetectionResponse; import org.opencv.core.Mat; +import org.opencv.core.Scalar; +import org.opencv.core.Size; +import org.opencv.imgproc.Imgproc; import javax.imageio.ImageIO; import java.awt.*; @@ -364,6 +367,7 @@ public class ImageUtils { int width = metrics.stringWidth(text) + padding * 2 - stroke / 2; int height = metrics.getHeight() + metrics.getDescent(); int ascent = metrics.getAscent(); + y = Math.max(0, y - height); java.awt.Rectangle background = new java.awt.Rectangle(x, y, width, height); g.fill(background); g.setPaint(Color.WHITE); @@ -468,6 +472,33 @@ public class ImageUtils { return true; } + /** + * 在图像上绘制带白色背景、黑色文字的文本 + */ + public static void putTextWithBackground(Mat image, String text, org.opencv.core.Point origin, Scalar textColor, Scalar backgroundColor, int padding) { + // 默认字体 + int font = Imgproc.FONT_HERSHEY_SCRIPT_SIMPLEX; + // 默认字体缩放大小 + double fontScale = 1.0; + //线条粗细 + int thickness = 2; + //获取文字大小 + int[] baseLine = new int[1]; + Size textSize = Imgproc.getTextSize(text, font, fontScale, thickness, baseLine); + int textWidth = (int) textSize.width; + int textHeight = (int) textSize.height; + + //计算带padding的背景框 + org.opencv.core.Point bgTopLeft = new org.opencv.core.Point(origin.x - padding, origin.y - textHeight - padding); + org.opencv.core.Point bgBottomRight = new org.opencv.core.Point(origin.x + textWidth + padding, origin.y + baseLine[0] + padding); + + //绘制背景矩形 + Imgproc.rectangle(image, bgTopLeft, bgBottomRight, backgroundColor, Imgproc.FILLED); + + //绘制文字（黑色） + Imgproc.putText(image, text, origin, font, fontScale, textColor, thickness); + } + } diff --git a/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/LetterBoxUtils.java b/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/LetterBoxUtils.java new file mode 100644 index 0000000..678251b --- /dev/null +++ b/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/LetterBoxUtils.java @@ -0,0 +1,124 @@ +package cn.smartjavaai.common.utils; + +import ai.djl.modality.cv.util.NDImageUtils; +import ai.djl.ndarray.NDArray; +import ai.djl.ndarray.NDManager; +import ai.djl.ndarray.index.NDIndex; +import ai.djl.ndarray.types.DataType; +import ai.djl.ndarray.types.Shape; + +import java.util.Arrays; + +/** + * 按比例缩放，剩余空间用指定颜色填充 + * @author dwj + */ +public class LetterBoxUtils { + + public enum PaddingPosition { + CENTER, LEFT_TOP, RIGHT_BOTTOM + } + + public static class ResizeResult { + public NDArray image; + public float r; + public int left; + public int top; + } + + public static ResizeResult letterboxWithMeta(NDArray paddingImg, float r, int left, int top) { + // ... letterbox 逻辑不变 + ResizeResult result = new ResizeResult(); + result.image = paddingImg; + result.r = r; + result.left = left; + result.top = top; + return result; + } + + + /** + * 按比例缩放 + padding + * + * @param img 原图 NDArray HWC + * @param targetW 目标宽度 + * @param targetH 目标高度 + * @param padColor padding 填充颜色，RGB 归一化 0-1 + * @param position padding 位置：CENTER / LEFT_TOP / RIGHT_BOTTOM + * @return 处理后的 NDArray + */ + public static ResizeResult letterbox(NDManager manager, NDArray img, int targetW, int targetH, float padColor, PaddingPosition position) { + long origH = img.getShape().get(0); + long origW = img.getShape().get(1); + + float r = Math.min(targetW / (float) origW, targetH / (float) origH); + int newW = Math.round(origW * r); + int newH = Math.round(origH * r); + + img = NDImageUtils.resize(img, newW, newH); // HWC 0~1 + +// NDArray paddingImg = manager +// .full(new Shape(targetW, targetH, 3), padColor, DataType.UINT8); + + NDArray paddingImg = manager.zeros(new Shape(targetW, targetH, 3), DataType.FLOAT32); + paddingImg = paddingImg.add(114); + + int padW = targetW - newW; + int padH = targetH - newH; + int top = 0, left = 0; + + switch (position) { + case CENTER: + left = padW / 2; + top = padH / 2; + break; + case LEFT_TOP: + left = 0; + top = 0; + break; + case RIGHT_BOTTOM: + left = padW; + top = padH; + break; + } + paddingImg.set(new NDIndex(String.format("%d:%d,%d:%d", top, top + newH, left, left + newW)), img); + return letterboxWithMeta(paddingImg, r, left, top); + } + + + /** + * 恢复缩放后的 box + * @param boxes + * @param scaleRatio + * @param left + * @param top + * @param keypointStart + * @param keypointDim + * @return + */ + public static NDArray restoreBox(NDArray boxes, float scaleRatio, float left, float top, int keypointStart, int keypointDim) { + // 处理 bbox + NDArray x1 = boxes.get(":, 0").sub(left).div(scaleRatio); + NDArray y1 = boxes.get(":, 1").sub(top).div(scaleRatio); + NDArray x2 = boxes.get(":, 2").sub(left).div(scaleRatio); + NDArray y2 = boxes.get(":, 3").sub(top).div(scaleRatio); + + boxes.set(new NDIndex(":, 0"), x1); + boxes.set(new NDIndex(":, 1"), y1); + boxes.set(new NDIndex(":, 2"), x2); + boxes.set(new NDIndex(":, 3"), y2); + + if (keypointDim > 0) { + for (int i = 0; i < keypointDim; i += 2) { + int xIdx = keypointStart + i; + int yIdx = keypointStart + i + 1; + NDArray keyX = boxes.get(":, " + xIdx).sub(left).div(scaleRatio); + NDArray keyY = boxes.get(":, " + yIdx).sub(top).div(scaleRatio); + boxes.set(new NDIndex(":, " + xIdx), keyX); + boxes.set(new NDIndex(":, " + yIdx), keyY); + } + } + return boxes; + } + +} diff --git a/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/NMSUtils.java b/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/NMSUtils.java new file mode 100644 index 0000000..ef2799a --- /dev/null +++ b/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/NMSUtils.java @@ -0,0 +1,67 @@ +package cn.smartjavaai.common.utils; + +import ai.djl.ndarray.NDArray; + +import java.util.ArrayList; +import java.util.List; + +/** + * @author dwj + * @date 2025/7/23 + */ +public class NMSUtils { + + /** + * 通用 NMS 方法，输入 NDArray 形式的 boxes 和 scores，返回保留的索引列表 + * + * @param boxes NDArray 形状为 (N, 4)，格式为 [x1, y1, x2, y2] + * @param scores NDArray 形状为 (N,) 或 (N,1)，每个 box 的置信度 + * @param iouThreshold IOU 阈值，超过该阈值则认为有重叠 + * @return 保留框的索引列表 + */ + public static int[] nms(NDArray boxes, NDArray scores, float iouThreshold) { + if (boxes.isEmpty()) { + return new int[0]; + } + + NDArray x1 = boxes.get(":, 0"); + NDArray y1 = boxes.get(":, 1"); + NDArray x2 = boxes.get(":, 2"); + NDArray y2 = boxes.get(":, 3"); + + NDArray areas = x2.sub(x1).add(1).mul(y2.sub(y1).add(1)); + + // 按照置信度降序排序 + NDArray order = scores.argSort().flip(0); + + List keep = new ArrayList<>(); + + while (order.size() > 0) { + int idx = (int)order.getLong(0); + keep.add(idx); + + if (order.size() == 1) break; + + NDArray currentBox = boxes.get(idx); + NDArray others = boxes.get(order); + + NDArray xx1 = x1.get(order).maximum(x1.get(idx)); + NDArray yy1 = y1.get(order).maximum(y1.get(idx)); + NDArray xx2 = x2.get(order).minimum(x2.get(idx)); + NDArray yy2 = y2.get(order).minimum(y2.get(idx)); + + NDArray w = xx2.sub(xx1).add(1).maximum(0); + NDArray h = yy2.sub(yy1).add(1).maximum(0); + NDArray inter = w.mul(h); + + NDArray remAreas = areas.get(order); + NDArray union = remAreas.add(areas.get(idx)).sub(inter); + NDArray iou = inter.div(union); + + NDArray mask = iou.lte(iouThreshold); + order = order.get(mask); + } + return keep.stream().mapToInt(i -> i).toArray(); + } + +} diff --git a/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/OpenCVUtils.java b/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/OpenCVUtils.java index 35c8692..5ffd4ab 100644 --- a/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/OpenCVUtils.java +++ b/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/OpenCVUtils.java @@ -124,4 +124,20 @@ public class OpenCVUtils { mat.put(0, 0, data); return mat; } + + /** + * 透视变换 + * + * @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; + } } diff --git a/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/PointUtils.java b/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/PointUtils.java new file mode 100644 index 0000000..38514ef --- /dev/null +++ b/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/PointUtils.java @@ -0,0 +1,70 @@ +package cn.smartjavaai.common.utils; + +import ai.djl.modality.cv.output.Point; + +import java.util.Arrays; +import java.util.Comparator; +import java.util.List; + +/** + * @author dwj + */ +public class PointUtils { + + /** + * 对 4 个关键点进行排序，顺序为： + * 左上、右上、右下、左下 + */ + public static List orderPoints(List points) { + if (points == null || points.size() != 4) { + throw new IllegalArgumentException("必须提供 4 个点"); + } + + // 按 X 坐标升序排列 + points.sort(Comparator.comparingDouble(Point::getX)); + + List left = points.subList(0, 2); + List right = points.subList(2, 4); + + // 左侧两点按 Y 排序：上为 tl，下为 bl + Point tl = left.get(0).getY() < left.get(1).getY() ? left.get(0) : left.get(1); + Point bl = left.get(0).getY() >= left.get(1).getY() ? left.get(0) : left.get(1); + + // 右侧两点按 Y 排序：上为 tr，下为 br + Point tr = right.get(0).getY() < right.get(1).getY() ? right.get(0) : right.get(1); + Point br = right.get(0).getY() >= right.get(1).getY() ? right.get(0) : right.get(1); + + return Arrays.asList(tl, tr, br, bl); + } + + /** + * 欧式距离计算 + * + * @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; + } + + /** + * 欧式距离计算 + * + * @param point1 + * @param point2 + * @return + */ + public static float distance(Point point1, Point point2) { + double disX = point1.getX() - point2.getX(); + double disY = point1.getY() - point2.getY(); + float dis = (float) Math.sqrt(disX * disX + disY * disY); + return dis; + } + + + +} diff --git a/smartjavaai-face/pom.xml b/smartjavaai-face/pom.xml index beecee2..6475e07 100644 --- a/smartjavaai-face/pom.xml +++ b/smartjavaai-face/pom.xml @@ -6,11 +6,11 @@ cn.smartjavaai smartjavaai-parent - 1.0.20 + 1.0.22 smartjavaai-face - 1.0.20 + 1.0.22 smartjavaai-face SmartJavaAI https://github.com/geekwenjie/SmartJavaAI diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceAttributeConfig.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceAttributeConfig.java index a7cc598..badf3f3 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceAttributeConfig.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceAttributeConfig.java @@ -1,5 +1,6 @@ package cn.smartjavaai.face.config; +import cn.smartjavaai.common.config.ModelConfig; import cn.smartjavaai.common.enums.DeviceEnum; import cn.smartjavaai.face.enums.FaceAttributeModelEnum; import lombok.Data; @@ -9,7 +10,7 @@ import lombok.Data; * @author dwj */ @Data -public class FaceAttributeConfig { +public class FaceAttributeConfig extends ModelConfig { /** * 人脸属性识别模型枚举 @@ -21,16 +22,6 @@ public class FaceAttributeConfig { */ private String modelPath; - /** - * 设备类型 - */ - private DeviceEnum device; - - /** - * gpu设备ID 当device为GPU时生效 - */ - private int gpuId = 0; - /** * 是否启用年龄检测 */ diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceDetConfig.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceDetConfig.java index 3b44a04..4a23fc6 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceDetConfig.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceDetConfig.java @@ -1,5 +1,6 @@ package cn.smartjavaai.face.config; +import cn.smartjavaai.common.config.ModelConfig; import cn.smartjavaai.common.enums.DeviceEnum; import cn.smartjavaai.face.constant.FaceDetectConstant; import cn.smartjavaai.face.enums.FaceDetModelEnum; @@ -13,7 +14,7 @@ import java.util.Map; * @author dwj */ @Data -public class FaceDetConfig { +public class FaceDetConfig extends ModelConfig { /** * 人脸检测模型枚举 @@ -36,16 +37,6 @@ public class FaceDetConfig { */ private String modelPath; - /** - * 设备类型 - */ - private DeviceEnum device; - - /** - * 个性化配置（按模型类型动态解析） - */ - private Map customParams = new HashMap<>(); - public FaceDetConfig() { } @@ -59,19 +50,4 @@ public class FaceDetConfig { this.modelPath = modelPath; } - public T getCustomParam(String key, Class clazz) { - Object value = customParams.get(key); - if (value == null) return null; - return clazz.cast(value); - } - - /** - * 添加个性化配置项 - */ - public void putCustomParam(String key, Object value) { - if (customParams == null) { - customParams = new HashMap<>(); - } - customParams.put(key, value); - } } diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceExpressionConfig.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceExpressionConfig.java index 38cefa5..0ba53d4 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceExpressionConfig.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceExpressionConfig.java @@ -1,5 +1,6 @@ package cn.smartjavaai.face.config; +import cn.smartjavaai.common.config.ModelConfig; import cn.smartjavaai.common.enums.DeviceEnum; import cn.smartjavaai.face.enums.ExpressionModelEnum; import cn.smartjavaai.face.model.facedect.FaceDetModel; @@ -14,7 +15,7 @@ import java.util.Map; * @date 2025/7/1 */ @Data -public class FaceExpressionConfig { +public class FaceExpressionConfig extends ModelConfig { /** * 模型枚举 @@ -26,11 +27,6 @@ public class FaceExpressionConfig { */ private String modelPath; - /** - * 设备类型 - */ - private DeviceEnum device; - /** * 人脸检测模型 */ diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceRecConfig.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceRecConfig.java index 7803f38..90ffd0a 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceRecConfig.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/FaceRecConfig.java @@ -1,5 +1,6 @@ package cn.smartjavaai.face.config; +import cn.smartjavaai.common.config.ModelConfig; import cn.smartjavaai.common.enums.DeviceEnum; import cn.smartjavaai.face.constant.FaceDetectConstant; import cn.smartjavaai.face.enums.FaceRecModelEnum; @@ -16,7 +17,7 @@ import java.util.Map; * @author dwj */ @Data -public class FaceRecConfig { +public class FaceRecConfig extends ModelConfig { /** * 人脸模型枚举 @@ -28,11 +29,6 @@ public class FaceRecConfig { */ private String modelPath; - /** - * 设备类型 - */ - private DeviceEnum device; - /** * 向量数据库配置 @@ -60,10 +56,6 @@ public class FaceRecConfig { */ private FaceDetModel detectModel; - /** - * 个性化配置（按模型类型动态解析） - */ - private Map customParams = new HashMap<>(); public FaceRecConfig() { } @@ -77,19 +69,5 @@ public class FaceRecConfig { this.modelPath = modelPath; } - public T getCustomParam(String key, Class clazz) { - Object value = customParams.get(key); - if (value == null) return null; - return clazz.cast(value); - } - /** - * 添加个性化配置项 - */ - public void putCustomParam(String key, Object value) { - if (customParams == null) { - customParams = new HashMap<>(); - } - customParams.put(key, value); - } } diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/LivenessConfig.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/LivenessConfig.java index 8ca3a02..f39567d 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/LivenessConfig.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/LivenessConfig.java @@ -1,5 +1,6 @@ package cn.smartjavaai.face.config; +import cn.smartjavaai.common.config.ModelConfig; import cn.smartjavaai.common.enums.DeviceEnum; import cn.smartjavaai.face.constant.LivenessConstant; import cn.smartjavaai.face.enums.LivenessModelEnum; @@ -15,7 +16,7 @@ import java.util.Map; * @author dwj */ @Data -public class LivenessConfig { +public class LivenessConfig extends ModelConfig { /** * 活体检测模型枚举 @@ -27,20 +28,12 @@ public class LivenessConfig { */ private String modelPath; - /** - * 设备类型 - */ - private DeviceEnum device; /** * 人脸检测模型 */ private FaceDetModel detectModel; - /** - * 个性化配置（按模型类型动态解析） - */ - private Map customParams = new HashMap<>(); /** @@ -74,20 +67,4 @@ public class LivenessConfig { this.modelPath = modelPath; } - // 可选封装方法，便于类型转换和调用 - public T getCustomParam(String key, Class clazz) { - Object value = customParams.get(key); - if (value == null) return null; - return clazz.cast(value); - } - - /** - * 添加个性化配置项 - */ - public void putCustomParam(String key, Object value) { - if (customParams == null) { - customParams = new HashMap<>(); - } - customParams.put(key, value); - } } diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/QualityConfig.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/QualityConfig.java index 2ccd1dd..98a9993 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/QualityConfig.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/config/QualityConfig.java @@ -1,5 +1,6 @@ package cn.smartjavaai.face.config; +import cn.smartjavaai.common.config.ModelConfig; import cn.smartjavaai.common.enums.DeviceEnum; import cn.smartjavaai.face.constant.LivenessConstant; import cn.smartjavaai.face.enums.LivenessModelEnum; @@ -11,7 +12,7 @@ import lombok.Data; * @author dwj */ @Data -public class QualityConfig { +public class QualityConfig extends ModelConfig { /** * 活体检测模型枚举 @@ -23,16 +24,6 @@ public class QualityConfig { */ private String modelPath; - /** - * 设备类型 - */ - private DeviceEnum device; - - /** - * gpu设备ID 当device为GPU时生效 - */ - private int gpuId = 0; - public QualityConfig() { } diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/attribute/Seetaface6FaceAttributeModel.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/attribute/Seetaface6FaceAttributeModel.java index 095d8f2..c2a92c3 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/attribute/Seetaface6FaceAttributeModel.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/attribute/Seetaface6FaceAttributeModel.java @@ -1,5 +1,6 @@ package cn.smartjavaai.face.model.attribute; +import ai.djl.engine.Engine; import cn.smartjavaai.common.entity.*; import cn.smartjavaai.common.entity.face.FaceAttribute; import cn.smartjavaai.common.entity.face.FaceInfo; @@ -102,6 +103,19 @@ public class Seetaface6FaceAttributeModel implements FaceAttributeModel { this.eyeStateDetectorPool = new EyeStateDetectorPool(eyeStateDetectorPoolConfSetting); this.maskDetectorPool = new MaskDetectorPool(maskDetectorPoolConfSetting); this.poseEstimatorPool = new PoseEstimatorPool(poseEstimatorPoolConfSetting); + + int predictorPoolSize = config.getPredictorPoolSize(); + if(config.getPredictorPoolSize() <= 0){ + predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数 + } + faceDetectorPool.setMaxTotal(predictorPoolSize); + faceLandmarkerPool.setMaxTotal(predictorPoolSize); + genderPredictorPool.setMaxTotal(predictorPoolSize); + agePredictorPool.setMaxTotal(predictorPoolSize); + eyeStateDetectorPool.setMaxTotal(predictorPoolSize); + maskDetectorPool.setMaxTotal(predictorPoolSize); + poseEstimatorPool.setMaxTotal(predictorPoolSize); + log.debug("模型推理器线程池最大数量: " + predictorPoolSize); } catch (FileNotFoundException e) { throw new FaceException(e); } @@ -471,6 +485,35 @@ public class Seetaface6FaceAttributeModel implements FaceAttributeModel { } } + + public FaceDetectorPool getFaceDetectorPool() { + return faceDetectorPool; + } + + public GenderPredictorPool getGenderPredictorPool() { + return genderPredictorPool; + } + + public FaceLandmarkerPool getFaceLandmarkerPool() { + return faceLandmarkerPool; + } + + public AgePredictorPool getAgePredictorPool() { + return agePredictorPool; + } + + public EyeStateDetectorPool getEyeStateDetectorPool() { + return eyeStateDetectorPool; + } + + public MaskDetectorPool getMaskDetectorPool() { + return maskDetectorPool; + } + + public PoseEstimatorPool getPoseEstimatorPool() { + return poseEstimatorPool; + } + @Override public void close() throws Exception { if(Objects.nonNull(faceDetectorPool)){ diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/expression/CommonEmotionModel.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/expression/CommonEmotionModel.java index 0d36451..06ef818 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/expression/CommonEmotionModel.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/expression/CommonEmotionModel.java @@ -2,6 +2,7 @@ package cn.smartjavaai.face.model.expression; import ai.djl.Device; import ai.djl.MalformedModelException; +import ai.djl.engine.Engine; import ai.djl.inference.Predictor; import ai.djl.modality.Classifications; import ai.djl.modality.cv.Image; @@ -32,6 +33,7 @@ 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 org.opencv.face.Face; import javax.imageio.ImageIO; @@ -56,7 +58,7 @@ public class CommonEmotionModel implements ExpressionModel{ private ZooModel model; - private ObjectPool> predictorPool; + private GenericObjectPool> predictorPool; @Override public void loadModel(FaceExpressionConfig config) { @@ -73,6 +75,14 @@ public class CommonEmotionModel implements ExpressionModel{ try { model = criteria.loadModel(); 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 FaceException("DenseNetEmotionModel模型加载失败", e); } @@ -153,6 +163,7 @@ public class CommonEmotionModel implements ExpressionModel{ result.setClassifications(classifications); faceInfo.setExpressionResult(result); } + ((Mat)djlImage.getWrappedImage()).release(); return faceDetectionResponse; } @@ -226,6 +237,7 @@ public class CommonEmotionModel implements ExpressionModel{ result.setClassifications(classifications); expressionResults.add(result); } + ((Mat)djlImage.getWrappedImage()).release(); return R.ok(expressionResults); } @@ -276,6 +288,7 @@ public class CommonEmotionModel implements ExpressionModel{ FacialExpression expression = FacialExpression.fromLabel(bestClass.getClassName()); ExpressionResult result = new ExpressionResult(expression, (float)bestClass.getProbability()); result.setClassifications(classifications); + ((Mat)djlImage.getWrappedImage()).release(); return R.ok(result); } @@ -342,7 +355,10 @@ public class CommonEmotionModel implements ExpressionModel{ return detectTopFace(imageData); } - + @Override + public GenericObjectPool> getPool() { + return predictorPool; + } @Override public void close() { diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/expression/ExpressionModel.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/expression/ExpressionModel.java index 2821677..d5e4e16 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/expression/ExpressionModel.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/expression/ExpressionModel.java @@ -1,11 +1,16 @@ package cn.smartjavaai.face.model.expression; +import ai.djl.inference.Predictor; +import ai.djl.modality.Classifications; +import ai.djl.modality.cv.Image; +import ai.djl.modality.cv.output.DetectedObjects; import cn.smartjavaai.common.entity.DetectionRectangle; import cn.smartjavaai.common.entity.DetectionResponse; import cn.smartjavaai.common.entity.Point; import cn.smartjavaai.common.entity.R; import cn.smartjavaai.common.entity.face.ExpressionResult; import cn.smartjavaai.face.config.FaceExpressionConfig; +import org.apache.commons.pool2.impl.GenericObjectPool; import java.awt.image.BufferedImage; import java.util.List; @@ -190,6 +195,9 @@ public interface ExpressionModel extends AutoCloseable{ + default GenericObjectPool> getPool(){ + throw new UnsupportedOperationException("默认不支持该功能"); + } } diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/expression/criterial/EmotionCriteriaFactory.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/expression/criterial/EmotionCriteriaFactory.java index 26d8cd8..02d28ac 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/expression/criterial/EmotionCriteriaFactory.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/expression/criterial/EmotionCriteriaFactory.java @@ -26,7 +26,7 @@ public class EmotionCriteriaFactory { public static Criteria createCriteria(FaceExpressionConfig 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 criteria = null; if(config.getModelEnum() == ExpressionModelEnum.DensNet121){ @@ -50,6 +50,7 @@ public class EmotionCriteriaFactory { .optModelPath(Paths.get(config.getModelPath())) .optTranslator(new FrEmotionTranslator(224)) .optProgress(new ProgressBar()) + .optDevice(device) .build(); } return criteria; diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/CommonFaceDetModel.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/CommonFaceDetModel.java index 9e97ce2..b79710d 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/CommonFaceDetModel.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/CommonFaceDetModel.java @@ -24,6 +24,7 @@ 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; @@ -43,7 +44,7 @@ import java.util.Objects; @Slf4j public class CommonFaceDetModel implements FaceDetModel{ - private ObjectPool> predictorPool; + private GenericObjectPool> predictorPool; private ZooModel model; @@ -58,8 +59,14 @@ public class CommonFaceDetModel implements FaceDetModel{ try { model = criteria.loadModel(); 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 FaceException("人脸检测模型加载失败", e); } @@ -81,11 +88,16 @@ public class CommonFaceDetModel implements FaceDetModel{ Image img = null; try { img = ImageFactory.getInstance().fromFile(Paths.get(imagePath)); + DetectedObjects detection = detect(img); + return R.ok(FaceUtils.convertToDetectionResponse(detection,img)); } catch (IOException e) { throw new FaceException("无效的图片", e); + } finally { + if (img != null) { + ((Mat)img.getWrappedImage()).release(); + } } - DetectedObjects detection = detect(img); - return R.ok(FaceUtils.convertToDetectionResponse(detection,img)); + } /** @@ -99,12 +111,17 @@ public class CommonFaceDetModel implements FaceDetModel{ if(Objects.isNull(imageInputStream)){ return R.fail(R.Status.INVALID_IMAGE); } + Image img = null; try { - Image img = ImageFactory.getInstance().fromInputStream(imageInputStream); + img = ImageFactory.getInstance().fromInputStream(imageInputStream); DetectedObjects detection = detect(img); return R.ok(FaceUtils.convertToDetectionResponse(detection,img)); } catch (IOException e) { throw new FaceException("无效图片输入流", e); + } finally { + if (img != null) { + ((Mat)img.getWrappedImage()).release(); + } } } @@ -114,9 +131,19 @@ public class CommonFaceDetModel implements FaceDetModel{ if(!ImageUtils.isImageValid(image)){ return R.fail(R.Status.INVALID_IMAGE); } - Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image)); - DetectedObjects detection = detect(img); - return R.ok(FaceUtils.convertToDetectionResponse(detection,img)); + Image img = null; + try { + img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image)); + DetectedObjects detection = detect(img); + return R.ok(FaceUtils.convertToDetectionResponse(detection,img)); + } catch (Exception e) { + throw new FaceException(e); + } finally { + if (img != null) { + ((Mat)img.getWrappedImage()).release(); + } + } + } @Override @@ -124,11 +151,7 @@ public class CommonFaceDetModel implements FaceDetModel{ if(Objects.isNull(imageData)){ return R.fail(R.Status.INVALID_IMAGE); } - try { - return detect(ImageIO.read(new ByteArrayInputStream(imageData))); - } catch (IOException e) { - throw new FaceException("错误的图像", e); - } + return detect(new ByteArrayInputStream(imageData)); } @Override @@ -145,8 +168,9 @@ public class CommonFaceDetModel implements FaceDetModel{ if(!FileUtils.isFileExists(imagePath)){ return R.fail(R.Status.FILE_NOT_FOUND); } + Image img = null; try { - Image img = ImageFactory.getInstance().fromFile(Paths.get(imagePath)); + 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); @@ -158,6 +182,10 @@ public class CommonFaceDetModel implements FaceDetModel{ return R.ok(); } catch (IOException e) { throw new FaceException(e); + } finally { + if (img != null){ + ((Mat)img.getWrappedImage()).release(); + } } } @@ -181,6 +209,10 @@ public class CommonFaceDetModel implements FaceDetModel{ return R.ok(ImageIO.read(new ByteArrayInputStream(imageBytes))); } catch (IOException e) { throw new FaceException("导出图片失败", e); + } finally { + if (img != null){ + ((Mat)img.getWrappedImage()).release(); + } } } @@ -213,6 +245,11 @@ public class CommonFaceDetModel implements FaceDetModel{ } + @Override + public GenericObjectPool> getPool() { + return predictorPool; + } + @Override public void close() { try { diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/FaceDetModel.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/FaceDetModel.java index ea6e4ea..475cd7d 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/FaceDetModel.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/FaceDetModel.java @@ -1,8 +1,13 @@ package cn.smartjavaai.face.model.facedect; +import ai.djl.inference.Predictor; +import ai.djl.modality.Classifications; +import ai.djl.modality.cv.Image; +import ai.djl.modality.cv.output.DetectedObjects; import cn.smartjavaai.common.entity.DetectionResponse; import cn.smartjavaai.common.entity.R; import cn.smartjavaai.face.config.FaceDetConfig; +import org.apache.commons.pool2.impl.GenericObjectPool; import java.awt.image.BufferedImage; import java.io.InputStream; @@ -83,6 +88,10 @@ public interface FaceDetModel extends AutoCloseable{ throw new UnsupportedOperationException("默认不支持该功能"); } + default GenericObjectPool> getPool(){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + } diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/SeetaFace6FaceDetModel.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/SeetaFace6FaceDetModel.java index 362680d..ecba841 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/SeetaFace6FaceDetModel.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/SeetaFace6FaceDetModel.java @@ -1,5 +1,6 @@ package cn.smartjavaai.face.model.facedect; +import ai.djl.engine.Engine; import cn.smartjavaai.common.entity.DetectionResponse; import cn.smartjavaai.common.entity.R; import cn.smartjavaai.common.enums.DeviceEnum; @@ -48,13 +49,9 @@ public class SeetaFace6FaceDetModel implements FaceDetModel{ String[] faceDetectorModelPath = {config.getModelPath() + File.separator + "face_detector.csta"}; String[] faceLandmarkerModelPath = {config.getModelPath() + File.separator + "face_landmarker_pts5.csta"}; SeetaDevice device = SeetaDevice.SEETA_DEVICE_AUTO; - int gpuId = 0; + int gpuId = config.getGpuId(); if(Objects.nonNull(config.getDevice())){ device = config.getDevice() == DeviceEnum.CPU ? SeetaDevice.SEETA_DEVICE_CPU : SeetaDevice.SEETA_DEVICE_GPU; - Integer gpuIdValue = config.getCustomParam("gpuId", Integer.class); - if(Objects.nonNull(gpuIdValue) && device == SeetaDevice.SEETA_DEVICE_GPU){ - gpuId = gpuIdValue; - } } try { SeetaModelSetting faceDetectorPoolSetting = new SeetaModelSetting(gpuId, faceDetectorModelPath, device); @@ -65,6 +62,14 @@ public class SeetaFace6FaceDetModel implements FaceDetModel{ this.faceDetectorPool = new FaceDetectorPool(faceDetectorPoolConfSetting); this.faceLandmarkerPool = new FaceLandmarkerPool(faceLandmarkerPoolConfSetting); + + int predictorPoolSize = config.getPredictorPoolSize(); + if(config.getPredictorPoolSize() <= 0){ + predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数 + } + faceDetectorPool.setMaxTotal(predictorPoolSize); + faceLandmarkerPool.setMaxTotal(predictorPoolSize); + log.debug("模型推理器线程池最大数量: " + predictorPoolSize); } catch (FileNotFoundException e) { throw new FaceException(e); } @@ -212,8 +217,13 @@ public class SeetaFace6FaceDetModel implements FaceDetModel{ } + public FaceDetectorPool getFaceDetectorPool() { + return faceDetectorPool; + } - + public FaceLandmarkerPool getFaceLandmarkerPool() { + return faceLandmarkerPool; + } @Override public void close() throws Exception { diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/criterial/FaceDetCriteriaFactory.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/criterial/FaceDetCriteriaFactory.java index a0f236a..38edb36 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/criterial/FaceDetCriteriaFactory.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facedect/criterial/FaceDetCriteriaFactory.java @@ -31,7 +31,7 @@ public class FaceDetCriteriaFactory { public static Criteria createCriteria(FaceDetConfig 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 criteria = null; if(config.getModelEnum() == FaceDetModelEnum.RETINA_FACE){ diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/CommonFaceRecModel.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/CommonFaceRecModel.java index 63cd24f..38e8941 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/CommonFaceRecModel.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/CommonFaceRecModel.java @@ -9,6 +9,7 @@ import ai.djl.ndarray.NDManager; import ai.djl.repository.zoo.Criteria; import ai.djl.repository.zoo.ModelNotFoundException; import ai.djl.repository.zoo.ZooModel; +import cn.hutool.core.lang.UUID; import cn.smartjavaai.common.entity.*; import cn.smartjavaai.common.entity.face.FaceInfo; import cn.smartjavaai.common.entity.face.FaceSearchResult; @@ -41,6 +42,7 @@ 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; @@ -68,7 +70,7 @@ public class CommonFaceRecModel implements FaceRecModel{ private static volatile boolean isLoadCompleted = false; - private ObjectPool> predictorPool; + private GenericObjectPool> predictorPool; private ZooModel model; @@ -101,8 +103,14 @@ public class CommonFaceRecModel implements FaceRecModel{ model = faceFeatureCriteria.loadModel(); // 创建池子：每个线程独享 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 FaceException("模型加载失败", e); } @@ -286,6 +294,7 @@ public class CommonFaceRecModel implements FaceRecModel{ faceInfo.setFeature(features); } } + ((Mat)djlImage.getWrappedImage()).release(); return detectedResult; } @@ -350,6 +359,7 @@ public class CommonFaceRecModel implements FaceRecModel{ } features = featureExtraction(subImage); } + ((Mat)djlImage.getWrappedImage()).release(); return Objects.isNull(features) ? R.fail(R.Status.Unknown) : R.ok(features); } @@ -703,5 +713,8 @@ public class CommonFaceRecModel implements FaceRecModel{ } - + @Override + public GenericObjectPool> getPool() { + return predictorPool; + } } diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/FaceRecModel.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/FaceRecModel.java index 1a695bf..2d3b9ef 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/FaceRecModel.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/FaceRecModel.java @@ -1,5 +1,7 @@ package cn.smartjavaai.face.model.facerec; +import ai.djl.inference.Predictor; +import ai.djl.modality.cv.Image; import cn.smartjavaai.common.entity.DetectionResponse; import cn.smartjavaai.common.entity.R; import cn.smartjavaai.face.config.FaceRecConfig; @@ -7,6 +9,7 @@ import cn.smartjavaai.face.entity.FaceRegisterInfo; import cn.smartjavaai.face.entity.FaceSearchParams; import cn.smartjavaai.common.entity.face.FaceSearchResult; import cn.smartjavaai.face.vector.entity.FaceVector; +import org.apache.commons.pool2.impl.GenericObjectPool; import java.awt.image.BufferedImage; import java.io.InputStream; @@ -367,4 +370,9 @@ public interface FaceRecModel extends AutoCloseable{ throw new UnsupportedOperationException("默认不支持该功能"); } + + default GenericObjectPool> getPool() { + throw new UnsupportedOperationException("默认不支持该功能"); + } + } diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/SeetaFace6FaceRecModel.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/SeetaFace6FaceRecModel.java index 64b58d0..bcb884d 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/SeetaFace6FaceRecModel.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/SeetaFace6FaceRecModel.java @@ -1,5 +1,6 @@ package cn.smartjavaai.face.model.facerec; +import ai.djl.engine.Engine; import cn.smartjavaai.common.entity.DetectionInfo; import cn.smartjavaai.common.entity.DetectionResponse; import cn.smartjavaai.common.entity.R; @@ -89,13 +90,9 @@ public class SeetaFace6FaceRecModel implements FaceRecModel{ String[] faceRecognizerModelPath = {config.getModelPath() + File.separator + "face_recognizer.csta"}; String[] faceLandmarkerModelPath = {config.getModelPath() + File.separator + "face_landmarker_pts5.csta"}; SeetaDevice device = SeetaDevice.SEETA_DEVICE_AUTO; - int gpuId = 0; + int gpuId = config.getGpuId(); if(Objects.nonNull(config.getDevice())){ device = config.getDevice() == DeviceEnum.CPU ? SeetaDevice.SEETA_DEVICE_CPU : SeetaDevice.SEETA_DEVICE_GPU; - Integer gpuIdValue = config.getCustomParam("gpuId", Integer.class); - if(Objects.nonNull(gpuIdValue) && device == SeetaDevice.SEETA_DEVICE_GPU){ - gpuId = gpuIdValue; - } } try { SeetaModelSetting faceDetectorPoolSetting = new SeetaModelSetting(gpuId, faceDetectorModelPath, device); @@ -115,6 +112,16 @@ public class SeetaFace6FaceRecModel implements FaceRecModel{ this.faceLandmarkerPool = new FaceLandmarkerPool(faceLandmarkerPoolConfSetting); this.faceDatabasePool = new FaceDatabasePool(faceDatabasePoolConfSetting); + int predictorPoolSize = config.getPredictorPoolSize(); + if(config.getPredictorPoolSize() <= 0){ + predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数 + } + faceDetectorPool.setMaxTotal(predictorPoolSize); + faceRecognizerPool.setMaxTotal(predictorPoolSize); + faceLandmarkerPool.setMaxTotal(predictorPoolSize); + faceDatabasePool.setMaxTotal(predictorPoolSize); + log.debug("模型推理器线程池最大数量: " + predictorPoolSize); + //初始化人脸库 if(config.getVectorDBConfig() != null && config.getVectorDBConfig().getType() != null){ @@ -739,6 +746,36 @@ public class SeetaFace6FaceRecModel implements FaceRecModel{ } } + /** + * 特征提取 + * @param image + * @return + */ + public float[] featureExtraction(BufferedImage image){ + FaceRecognizer faceRecognizer = null; + try { + faceRecognizer = faceRecognizerPool.borrowObject(); + SeetaImageData imageData = new SeetaImageData(image.getWidth(), image.getHeight(), 3); + imageData.data = ImageUtils.getMatrixBGR(image); + //提取特征 + float[] features = new float[faceRecognizer.GetExtractFeatureSize()]; + faceRecognizer.ExtractCroppedFace(imageData, features); + return features; + } catch (FaceException e) { + throw e; + } catch (Exception e) { + throw new FaceException("目标检测错误", e); + }finally { + if (faceRecognizer != null) { + try { + faceRecognizerPool.returnObject(faceRecognizer); //归还 + } catch (Exception e) { + log.warn("归还Predictor失败", e); + } + } + } + } + @Override public R extractTopFaceFeature(String imagePath) { if(!FileUtils.isFileExists(imagePath)){ @@ -878,4 +915,19 @@ public class SeetaFace6FaceRecModel implements FaceRecModel{ } + public FaceDetectorPool getFaceDetectorPool() { + return faceDetectorPool; + } + + public FaceRecognizerPool getFaceRecognizerPool() { + return faceRecognizerPool; + } + + public FaceLandmarkerPool getFaceLandmarkerPool() { + return faceLandmarkerPool; + } + + public FaceDatabasePool getFaceDatabasePool() { + return faceDatabasePool; + } } diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/criteria/FaceRecCriteriaFactory.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/criteria/FaceRecCriteriaFactory.java index 47b27a9..de6b708 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/criteria/FaceRecCriteriaFactory.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/criteria/FaceRecCriteriaFactory.java @@ -27,7 +27,7 @@ public class FaceRecCriteriaFactory { public static Criteria createCriteria(FaceRecConfig 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 criteria = null; if(config.getModelEnum() == FaceRecModelEnum.FACENET_MODEL){ @@ -58,6 +58,7 @@ public class FaceRecCriteriaFactory { .optTranslator(new FaceFeatureTranslator()) .optEngine("PyTorch") // Use PyTorch engine .optProgress(new ProgressBar()) + .optDevice(device) .build(); }else if (config.getModelEnum() == FaceRecModelEnum.INSIGHT_FACE_IRSE50_MODEL){ if(StringUtils.isBlank(config.getModelPath())){ @@ -73,6 +74,7 @@ public class FaceRecCriteriaFactory { // .optArgument("resize", "112,112") .optTranslator(new FaceFeatureTranslator()) .optEngine("PyTorch") // Use PyTorch engine + .optDevice(device) .optProgress(new ProgressBar()) .build(); }else if (config.getModelEnum() == FaceRecModelEnum.ELASTIC_FACE_MODEL){ @@ -89,6 +91,7 @@ public class FaceRecCriteriaFactory { // .optArgument("resize", "112,112") .optTranslator(new FaceFeatureTranslator()) .optEngine("PyTorch") // Use PyTorch engine + .optDevice(device) .optProgress(new ProgressBar()) .build(); } diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/translator/FaceFeatureTranslator.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/translator/FaceFeatureTranslator.java index 52c1e9e..6b623af 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/translator/FaceFeatureTranslator.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/translator/FaceFeatureTranslator.java @@ -31,7 +31,7 @@ public final class FaceFeatureTranslator implements Translator { NDArray array = input.toNDArray(ctx.getNDManager(), Image.Flag.COLOR); Pipeline pipeline = new Pipeline(); if(input.getWidth() != 112 || input.getHeight() != 112){ - pipeline.add(new Resize(112)); + pipeline.add(new Resize(112,112)); } pipeline .add(new ToTensor()) diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/CommonLivenessModel.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/CommonLivenessModel.java index 491fbac..f1fbd0a 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/CommonLivenessModel.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/CommonLivenessModel.java @@ -2,6 +2,7 @@ package cn.smartjavaai.face.model.liveness; import ai.djl.Device; 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; @@ -35,6 +36,7 @@ import org.apache.commons.pool2.impl.GenericObjectPool; import org.bytedeco.javacv.FFmpegFrameGrabber; import org.bytedeco.javacv.Frame; import org.bytedeco.javacv.Java2DFrameUtils; +import org.opencv.core.Mat; import javax.imageio.ImageIO; import java.awt.image.BufferedImage; @@ -54,7 +56,7 @@ import java.util.*; @Slf4j public class CommonLivenessModel implements LivenessDetModel{ - protected ObjectPool> predictorPool; + protected GenericObjectPool> predictorPool; protected LivenessConfig config; @@ -78,6 +80,15 @@ public class CommonLivenessModel implements LivenessDetModel{ } catch (IOException | ModelNotFoundException | MalformedModelException e) { throw new FaceException("阿里通义实验室活体检测模型加载失败", e); } + + 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); } @@ -87,6 +98,7 @@ public class CommonLivenessModel implements LivenessDetModel{ throw new FaceException("图像无效"); } Predictor predictor = null; + Image djlImage = null; try { predictor = predictorPool.borrowObject(); //预处理图片 @@ -101,7 +113,7 @@ public class CommonLivenessModel implements LivenessDetModel{ .setCenterCropSize(112) .process(); } - Image djlImage = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(processedImage)); + djlImage = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(processedImage)); Float result = predictor.predict(djlImage); if(result >= config.getRealityThreshold()){ return R.ok(new LivenessResult(LivenessStatus.LIVE, result)); @@ -124,6 +136,10 @@ public class CommonLivenessModel implements LivenessDetModel{ } } } + if (djlImage != null){ + ((Mat)djlImage.getWrappedImage()).release(); + } + } } @@ -403,6 +419,12 @@ public class CommonLivenessModel implements LivenessDetModel{ return R.fail(R.Status.Unknown); } + + @Override + public GenericObjectPool> getPool() { + return predictorPool; + } + @Override public void close() throws Exception { try { diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/LivenessDetModel.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/LivenessDetModel.java index c0c4f49..5256c4e 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/LivenessDetModel.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/LivenessDetModel.java @@ -1,11 +1,14 @@ package cn.smartjavaai.face.model.liveness; +import ai.djl.inference.Predictor; +import ai.djl.modality.cv.Image; import cn.smartjavaai.common.entity.DetectionRectangle; import cn.smartjavaai.common.entity.DetectionResponse; import cn.smartjavaai.common.entity.Point; import cn.smartjavaai.common.entity.R; import cn.smartjavaai.common.entity.face.LivenessResult; import cn.smartjavaai.face.config.LivenessConfig; +import org.apache.commons.pool2.impl.GenericObjectPool; import java.awt.image.BufferedImage; import java.io.InputStream; @@ -292,6 +295,9 @@ public interface LivenessDetModel extends AutoCloseable{ + default GenericObjectPool> getPool() { + throw new UnsupportedOperationException("默认不支持该功能"); + } diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/MiniVisionLivenessModel.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/MiniVisionLivenessModel.java index fadb52b..0c56f36 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/MiniVisionLivenessModel.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/MiniVisionLivenessModel.java @@ -2,6 +2,7 @@ package cn.smartjavaai.face.model.liveness; import ai.djl.Device; 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; @@ -33,6 +34,7 @@ import org.apache.commons.pool2.impl.GenericObjectPool; import org.bytedeco.javacv.FFmpegFrameGrabber; import org.bytedeco.javacv.Frame; import org.bytedeco.javacv.Java2DFrameUtils; +import org.opencv.core.Mat; import javax.imageio.ImageIO; import java.awt.image.BufferedImage; @@ -58,9 +60,9 @@ public class MiniVisionLivenessModel extends CommonLivenessModel{ */ private static final String SE_MODEL_PATH_KEY = "seModelPath"; - private ObjectPool> predictorPool; + private GenericObjectPool> predictorPool; - private ObjectPool> sePredictorPool; + private GenericObjectPool> sePredictorPool; /** @@ -87,7 +89,7 @@ public class MiniVisionLivenessModel extends CommonLivenessModel{ this.config.setRealityThreshold(realityThreshold); 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()); } if(StringUtils.isNotBlank(config.getModelPath()) && StringUtils.isBlank(seModelPath)){ //2.7_80x80_MiniFASNetV2 @@ -109,6 +111,7 @@ public class MiniVisionLivenessModel extends CommonLivenessModel{ .optModelPath(Paths.get(config.getModelPath())) .optTranslator(new MiniVisionTranslator()) .optProgress(new ProgressBar()) + .optDevice(device) .build(); try { model = criteria.loadModel(); @@ -128,6 +131,7 @@ public class MiniVisionLivenessModel extends CommonLivenessModel{ .optModelPath(Paths.get(seModelPath)) .optTranslator(new MiniVisionTranslator()) .optProgress(new ProgressBar()) + .optDevice(device) .build(); try { seModel = seCriteria.loadModel(); @@ -136,6 +140,16 @@ public class MiniVisionLivenessModel extends CommonLivenessModel{ throw new FaceException("MiniFASNetV1SE模型加载失败", e); } } + + int predictorPoolSize = config.getPredictorPoolSize(); + if(config.getPredictorPoolSize() <= 0){ + predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数 + } + predictorPool.setMaxTotal(predictorPoolSize); + sePredictorPool.setMaxTotal(predictorPoolSize); + log.debug("当前设备: " + model.getNDManager().getDevice()); + log.debug("当前引擎: " + Engine.getInstance().getEngineName()); + log.debug("模型推理器线程池最大数量: " + predictorPoolSize); } @@ -160,6 +174,8 @@ public class MiniVisionLivenessModel extends CommonLivenessModel{ predictor = predictorPool.borrowObject(); Image djlImage = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(processedImage)); result = predictor.predict(djlImage); + ((Mat)djlImage.getWrappedImage()).release(); + } if(Objects.nonNull(sePredictorPool)){ //预处理图片 @@ -172,6 +188,7 @@ public class MiniVisionLivenessModel extends CommonLivenessModel{ Image djlImage = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(processedImage)); sePredictor = sePredictorPool.borrowObject(); seResult = sePredictor.predict(djlImage); + ((Mat)djlImage.getWrappedImage()).release(); } if(Objects.isNull(result) && Objects.isNull(seResult)){ throw new FaceException("活体检测错误"); @@ -221,6 +238,14 @@ public class MiniVisionLivenessModel extends CommonLivenessModel{ } } + public GenericObjectPool> getPredictorPool() { + return predictorPool; + } + + public GenericObjectPool> getSePredictorPool() { + return sePredictorPool; + } + @Override public void close() throws Exception { try { diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/Seetaface6LivenessModel.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/Seetaface6LivenessModel.java index d286af2..119f22c 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/Seetaface6LivenessModel.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/Seetaface6LivenessModel.java @@ -1,5 +1,6 @@ package cn.smartjavaai.face.model.liveness; +import ai.djl.engine.Engine; import cn.smartjavaai.common.entity.*; import cn.smartjavaai.common.entity.face.FaceInfo; import cn.smartjavaai.common.entity.face.LivenessResult; @@ -56,13 +57,9 @@ public class Seetaface6LivenessModel implements LivenessDetModel{ String[] faceAntiSpoofingModelPath = {config.getModelPath() + File.separator + "fas_first.csta",config.getModelPath() + File.separator + "fas_second.csta"}; String[] faceLandmarkerModelPath = {config.getModelPath() + File.separator + "face_landmarker_pts5.csta"}; SeetaDevice device = SeetaDevice.SEETA_DEVICE_AUTO; - int gpuId = 0; + int gpuId = config.getGpuId(); if(Objects.nonNull(config.getDevice())){ device = config.getDevice() == DeviceEnum.CPU ? SeetaDevice.SEETA_DEVICE_CPU : SeetaDevice.SEETA_DEVICE_GPU; - Integer gpuIdValue = config.getCustomParam("gpuId", Integer.class); - if(Objects.nonNull(gpuIdValue) && device == SeetaDevice.SEETA_DEVICE_GPU){ - gpuId = gpuIdValue; - } } try { @@ -79,6 +76,16 @@ public class Seetaface6LivenessModel implements LivenessDetModel{ this.faceAntiSpoofingPool = new FaceAntiSpoofingPool(faceAntiSpoofingPoolConfSetting); this.faceLandmarkerPool = new FaceLandmarkerPool(faceLandmarkerPoolConfSetting); + int predictorPoolSize = config.getPredictorPoolSize(); + if(config.getPredictorPoolSize() <= 0){ + predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数 + } + + faceDetectorPool.setMaxTotal(predictorPoolSize); + faceAntiSpoofingPool.setMaxTotal(predictorPoolSize); + faceLandmarkerPool.setMaxTotal(predictorPoolSize); + log.debug("模型推理器线程池最大数量: " + predictorPoolSize); + //初始化模型参数 initConfig(); } catch (FileNotFoundException e) { @@ -550,6 +557,18 @@ public class Seetaface6LivenessModel implements LivenessDetModel{ return R.fail(1000, "有效帧数量不足，无法完成活体检测"); } + public FaceDetectorPool getFaceDetectorPool() { + return faceDetectorPool; + } + + public FaceAntiSpoofingPool getFaceAntiSpoofingPool() { + return faceAntiSpoofingPool; + } + + public FaceLandmarkerPool getFaceLandmarkerPool() { + return faceLandmarkerPool; + } + @Override public void close() throws Exception { try { diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/criterial/LivenessCriteriaFactory.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/criterial/LivenessCriteriaFactory.java index 086797b..c3003c7 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/criterial/LivenessCriteriaFactory.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/liveness/criterial/LivenessCriteriaFactory.java @@ -31,7 +31,7 @@ public class LivenessCriteriaFactory { public static Criteria createCriteria(LivenessConfig 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 criteria = null; if(config.getModelEnum() == LivenessModelEnum.IIC_FL_MODEL){ diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/quality/Seetaface6QualityModel.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/quality/Seetaface6QualityModel.java index f8d9dfe..7dd8469 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/quality/Seetaface6QualityModel.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/quality/Seetaface6QualityModel.java @@ -1,5 +1,6 @@ package cn.smartjavaai.face.model.quality; +import ai.djl.engine.Engine; import cn.smartjavaai.common.entity.*; import cn.smartjavaai.common.enums.DeviceEnum; import cn.smartjavaai.common.utils.FileUtils; @@ -74,6 +75,8 @@ public class Seetaface6QualityModel implements FaceQualityModel { */ private QualityOfResolutionPool qualityOfResolutionPool; + int predictorPoolSize = 0; + @Override public void loadModel(QualityConfig config) { @@ -84,6 +87,11 @@ public class Seetaface6QualityModel implements FaceQualityModel { //加载依赖库 NativeLoader.loadNativeLibraries(device); this.config = config; + int predictorPoolSize = config.getPredictorPoolSize(); + if(config.getPredictorPoolSize() <= 0){ + predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数 + } + log.debug("模型推理器线程池最大数量: " + predictorPoolSize); log.debug("Loading seetaFace6 library successfully."); } @@ -104,6 +112,7 @@ public class Seetaface6QualityModel implements FaceQualityModel { try { if(Objects.isNull(this.qualityOfBrightnessPool)){ this.qualityOfBrightnessPool = new QualityOfBrightnessPool(new SeetaConfSetting()); + qualityOfBrightnessPool.setMaxTotal(predictorPoolSize); } qualityOfBrightness = qualityOfBrightnessPool.borrowObject(); SeetaImageData imageData = new SeetaImageData(image.getWidth(), image.getHeight(), 3); @@ -170,6 +179,7 @@ public class Seetaface6QualityModel implements FaceQualityModel { try { if(Objects.isNull(this.qualityOfClarityPool)){ this.qualityOfClarityPool = new QualityOfClarityPool(new SeetaConfSetting()); + qualityOfClarityPool.setMaxTotal(predictorPoolSize); } qualityOfClarity = qualityOfClarityPool.borrowObject(); SeetaImageData imageData = new SeetaImageData(image.getWidth(), image.getHeight(), 3); @@ -236,6 +246,7 @@ public class Seetaface6QualityModel implements FaceQualityModel { try { if(Objects.isNull(this.qualityOfIntegrityPool)){ this.qualityOfIntegrityPool = new QualityOfIntegrityPool(new SeetaConfSetting()); + qualityOfIntegrityPool.setMaxTotal(predictorPoolSize); } qualityOfIntegrity = qualityOfIntegrityPool.borrowObject(); SeetaImageData imageData = new SeetaImageData(image.getWidth(), image.getHeight(), 3); @@ -302,6 +313,7 @@ public class Seetaface6QualityModel implements FaceQualityModel { try { if(Objects.isNull(this.qualityOfPosePool)){ this.qualityOfPosePool = new QualityOfPosePool(new SeetaConfSetting()); + qualityOfPosePool.setMaxTotal(predictorPoolSize); } qualityOfPose = qualityOfPosePool.borrowObject(); SeetaImageData imageData = new SeetaImageData(image.getWidth(), image.getHeight(), 3); @@ -368,6 +380,7 @@ public class Seetaface6QualityModel implements FaceQualityModel { try { if(Objects.isNull(this.qualityOfResolutionPool)){ this.qualityOfResolutionPool = new QualityOfResolutionPool(new SeetaConfSetting()); + qualityOfResolutionPool.setMaxTotal(predictorPoolSize); } qualityOfResolution = qualityOfResolutionPool.borrowObject(); SeetaImageData imageData = new SeetaImageData(image.getWidth(), image.getHeight(), 3); @@ -438,6 +451,7 @@ public class Seetaface6QualityModel implements FaceQualityModel { } SeetaConfSetting setting = getClarityMLSetting(); this.qualityOfLBNPool = new QualityOfLBNPool(setting); + qualityOfLBNPool.setMaxTotal(predictorPoolSize); } qualityOfLBN = qualityOfLBNPool.borrowObject(); SeetaImageData imageData = new SeetaImageData(image.getWidth(), image.getHeight(), 3); @@ -509,6 +523,7 @@ public class Seetaface6QualityModel implements FaceQualityModel { } SeetaConfSetting setting = getPoseMLSetting(); this.qualityOfPoseExPool = new QualityOfPoseExPool(setting); + qualityOfPoseExPool.setMaxTotal(predictorPoolSize); } qualityOfPoseEx = qualityOfPoseExPool.borrowObject(); SeetaImageData imageData = new SeetaImageData(image.getWidth(), image.getHeight(), 3); @@ -628,18 +643,23 @@ public class Seetaface6QualityModel implements FaceQualityModel { try { if(Objects.isNull(this.qualityOfBrightnessPool)){ this.qualityOfBrightnessPool = new QualityOfBrightnessPool(new SeetaConfSetting()); + qualityOfBrightnessPool.setMaxTotal(predictorPoolSize); } if(Objects.isNull(this.qualityOfClarityPool)){ this.qualityOfClarityPool = new QualityOfClarityPool(new SeetaConfSetting()); + qualityOfClarityPool.setMaxTotal(predictorPoolSize); } if(Objects.isNull(this.qualityOfIntegrityPool)){ this.qualityOfIntegrityPool = new QualityOfIntegrityPool(new SeetaConfSetting()); + qualityOfIntegrityPool.setMaxTotal(predictorPoolSize); } if(Objects.isNull(this.qualityOfPosePool)){ this.qualityOfPosePool = new QualityOfPosePool(new SeetaConfSetting()); + qualityOfPosePool.setMaxTotal(predictorPoolSize); } if(Objects.isNull(this.qualityOfResolutionPool)){ this.qualityOfResolutionPool = new QualityOfResolutionPool(new SeetaConfSetting()); + qualityOfResolutionPool.setMaxTotal(predictorPoolSize); } FaceQualitySummary summary = new FaceQualitySummary(); qualityOfBrightness = qualityOfBrightnessPool.borrowObject(); @@ -690,6 +710,34 @@ public class Seetaface6QualityModel implements FaceQualityModel { } } + public QualityOfBrightnessPool getQualityOfBrightnessPool() { + return qualityOfBrightnessPool; + } + + public QualityOfClarityPool getQualityOfClarityPool() { + return qualityOfClarityPool; + } + + public QualityOfLBNPool getQualityOfLBNPool() { + return qualityOfLBNPool; + } + + public QualityOfIntegrityPool getQualityOfIntegrityPool() { + return qualityOfIntegrityPool; + } + + public QualityOfPosePool getQualityOfPosePool() { + return qualityOfPosePool; + } + + public QualityOfPoseExPool getQualityOfPoseExPool() { + return qualityOfPoseExPool; + } + + public QualityOfResolutionPool getQualityOfResolutionPool() { + return qualityOfResolutionPool; + } + @Override public void close() throws Exception { if(Objects.nonNull(qualityOfBrightnessPool)){ diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/preprocess/DJLImagePreprocessor.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/preprocess/DJLImagePreprocessor.java index de62d4b..9060cc5 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/preprocess/DJLImagePreprocessor.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/preprocess/DJLImagePreprocessor.java @@ -60,15 +60,13 @@ public class DJLImagePreprocessor { // 处理流程 public Image process() { Image result = image; - - if (enableCrop) { - result = result.getSubImage(cropRect.x, cropRect.y, cropRect.width, cropRect.height); - } - - if (enableAffine) { + if(enableAffine){ result = warpAffine(keyPoints, affineTargetWidth, affineTargetHeight); + }else { + if(enableCrop){ + result = result.getSubImage(cropRect.x, cropRect.y, cropRect.width, cropRect.height); + } } - return result; } diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/config/MilvusConfig.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/config/MilvusConfig.java index 51b4948..27ef509 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/config/MilvusConfig.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/config/MilvusConfig.java @@ -62,6 +62,16 @@ public class MilvusConfig extends VectorDBConfig { */ private boolean useMemoryCache = true; + /** + * 用户名 + */ + private String username; + + /** + * 密码 + */ + private String password; + /** * 构造函数 */ diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/core/MilvusClient.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/core/MilvusClient.java index 2f7a466..39595ff 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/core/MilvusClient.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/core/MilvusClient.java @@ -49,10 +49,13 @@ public class MilvusClient implements VectorDBClient { @Override public void initialize() { try { - ConnectParam connectParam = ConnectParam.newBuilder() + ConnectParam.Builder builder = ConnectParam.newBuilder() .withHost(config.getHost()) - .withPort(config.getPort()) - .build(); + .withPort(config.getPort()); + if (StringUtils.isNotBlank(config.getUsername()) && StringUtils.isNotBlank(config.getPassword())) { + builder.withAuthorization(config.getUsername(), config.getPassword()); + } + ConnectParam connectParam = builder.build(); serviceClient = new MilvusServiceClient(connectParam); collectionName = StringUtils.isNotBlank(config.getCollectionName()) ? config.getCollectionName() : VectorDBConstants.Defaults.DEFAULT_COLLECTION_NAME; createCollection(collectionName, config.getDimension()); diff --git a/smartjavaai-objectdetection/pom.xml b/smartjavaai-objectdetection/pom.xml index 8b10643..51a2dc8 100644 --- a/smartjavaai-objectdetection/pom.xml +++ b/smartjavaai-objectdetection/pom.xml @@ -6,11 +6,11 @@ cn.smartjavaai smartjavaai-parent - 1.0.20 + 1.0.22 smartjavaai-objectdetection - 1.0.20 + 1.0.22 smartjavaai-objectdetection SmartJavaAI https://github.com/geekwenjie/SmartJavaAI diff --git a/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/config/DetectorModelConfig.java b/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/config/DetectorModelConfig.java index 36e2624..b0b1a24 100644 --- a/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/config/DetectorModelConfig.java +++ b/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/config/DetectorModelConfig.java @@ -1,11 +1,13 @@ package cn.smartjavaai.objectdetection.config; +import cn.smartjavaai.common.config.ModelConfig; import cn.smartjavaai.common.enums.DeviceEnum; import cn.smartjavaai.objectdetection.constant.DetectorConstant; import cn.smartjavaai.objectdetection.enums.DetectorModelEnum; import lombok.Data; import java.util.HashMap; +import java.util.List; import java.util.Map; /** @@ -15,7 +17,7 @@ import java.util.Map; * @date 2025/4/4 */ @Data -public class DetectorModelConfig { +public class DetectorModelConfig extends ModelConfig { /** * 模型 @@ -27,10 +29,6 @@ public class DetectorModelConfig { */ private float threshold = DetectorConstant.DEFAULT_THRESHOLD; - /** - * 设备类型 - */ - private DeviceEnum device; /** * 模型路径 @@ -45,35 +43,27 @@ public class DetectorModelConfig { private int maxBox; /** - * 个性化配置 + * 允许的类列表 */ - private Map customParams = new HashMap<>(); + private List allowedClasses; + + /** + * 检测结果数量 + */ + private int topK; + public DetectorModelConfig() { } public DetectorModelConfig(DetectorModelEnum modelEnum, DeviceEnum device) { this.modelEnum = modelEnum; - this.device = device; + setDevice(device); } public DetectorModelConfig(DetectorModelEnum modelEnum) { this.modelEnum = modelEnum; } - public T getCustomParam(String key, Class clazz) { - Object value = customParams.get(key); - if (value == null) return null; - return clazz.cast(value); - } - /** - * 添加个性化配置项 - */ - public void putCustomParam(String key, Object value) { - if (customParams == null) { - customParams = new HashMap<>(); - } - customParams.put(key, value); - } } diff --git a/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/criteria/DJLModelCriteriaBuilder.java b/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/criteria/DJLModelCriteriaBuilder.java index bb64684..8cdd635 100644 --- a/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/criteria/DJLModelCriteriaBuilder.java +++ b/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/criteria/DJLModelCriteriaBuilder.java @@ -25,7 +25,7 @@ public class DJLModelCriteriaBuilder implements CriteriaBuilderStrategy { public Criteria buildCriteria(DetectorModelConfig 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 criteria = Criteria.builder() .optApplication(Application.CV.OBJECT_DETECTION) diff --git a/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/criteria/YoloCriteriaBuilder.java b/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/criteria/YoloCriteriaBuilder.java index 7e4f3c0..501d7e5 100644 --- a/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/criteria/YoloCriteriaBuilder.java +++ b/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/criteria/YoloCriteriaBuilder.java @@ -25,7 +25,7 @@ public class YoloCriteriaBuilder implements CriteriaBuilderStrategy { public Criteria buildCriteria(DetectorModelConfig 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()); } Map customParams = getDefaultConfig(); diff --git a/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/model/DetectorModel.java b/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/model/DetectorModel.java index bd517b0..b4e3b40 100644 --- a/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/model/DetectorModel.java +++ b/smartjavaai-objectdetection/src/main/java/cn/smartjavaai/objectdetection/model/DetectorModel.java @@ -1,16 +1,22 @@ package cn.smartjavaai.objectdetection.model; 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.modality.cv.translator.YoloV8TranslatorFactory; import ai.djl.repository.zoo.Criteria; import ai.djl.repository.zoo.ModelNotFoundException; import ai.djl.repository.zoo.ZooModel; import ai.djl.training.util.ProgressBar; +import cn.smartjavaai.common.entity.DetectionInfo; +import cn.smartjavaai.common.entity.DetectionRectangle; import cn.smartjavaai.common.entity.DetectionResponse; +import cn.smartjavaai.common.entity.Point; +import cn.smartjavaai.common.entity.face.FaceInfo; import cn.smartjavaai.common.pool.PredictorFactory; import cn.smartjavaai.common.utils.FileUtils; import cn.smartjavaai.common.utils.ImageUtils; @@ -21,39 +27,53 @@ import cn.smartjavaai.objectdetection.criteria.CriteriaBuilderFactory; import cn.smartjavaai.objectdetection.exception.DetectionException; import cn.smartjavaai.objectdetection.utils.DetectorUtils; 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.Mat; import javax.imageio.ImageIO; import java.awt.image.BufferedImage; import java.io.*; import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; import java.util.Objects; +import java.util.stream.Collectors; /** * 目标检测模型 * @author dwj - * @date 2025/4/4 */ @Slf4j public class DetectorModel implements AutoCloseable{ private ZooModel model; - private ObjectPool> predictorPool; + private GenericObjectPool> predictorPool; + + private DetectorModelConfig config; public void loadModel(DetectorModelConfig config){ if(Objects.isNull(config.getModelEnum())){ throw new DetectionException("未配置模型枚举"); } Criteria criteria = CriteriaBuilderFactory.createCriteria(config); - + this.config = config; try { model = criteria.loadModel(); // 创建池子：每个线程独享 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 DetectionException("模型加载失败", e); } @@ -72,11 +92,16 @@ public class DetectorModel implements AutoCloseable{ Image image = null; try { image = ImageFactory.getInstance().fromFile(Paths.get(imagePath)); - } catch (IOException e) { - throw new DetectionException("图片转换错误", e); + DetectedObjects detectedObjects = detect(image); + return DetectorUtils.convertToDetectionResponse(detectedObjects, image); + } catch (Exception e) { + throw new DetectionException(e); + } finally { + if (image != null){ + ((Mat)image.getWrappedImage()).release(); + } } - DetectedObjects detectedObjects = detect(image); - return DetectorUtils.convertToDetectionResponse(detectedObjects, image); + } @@ -89,8 +114,9 @@ public class DetectorModel implements AutoCloseable{ if(!FileUtils.isFileExists(imagePath)){ throw new DetectionException("图像文件不存在"); } + Image img = null; try { - Image img = ImageFactory.getInstance().fromFile(Paths.get(imagePath)); + img = ImageFactory.getInstance().fromFile(Paths.get(imagePath)); DetectedObjects detectedObjects = detect(img); img.drawBoundingBoxes(detectedObjects); ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); @@ -98,6 +124,10 @@ public class DetectorModel implements AutoCloseable{ img.save(new FileOutputStream(Paths.get(outputPath).toAbsolutePath().toString()), "png"); } catch (IOException e) { throw new DetectionException(e); + } finally { + if (img != null){ + ((Mat)img.getWrappedImage()).release(); + } } } @@ -129,9 +159,19 @@ public class DetectorModel implements AutoCloseable{ if(!ImageUtils.isImageValid(image)){ throw new DetectionException("图像无效"); } - Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image)); - DetectedObjects detectedObjects = detect(img); - return DetectorUtils.convertToDetectionResponse(detectedObjects, img); + Image img = null; + try { + img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image)); + DetectedObjects detectedObjects = detect(img); + return DetectorUtils.convertToDetectionResponse(detectedObjects, img); + } catch (Exception e) { + throw new DetectionException(e); + } finally { + if (img != null) { + ((Mat)img.getWrappedImage()).release(); + } + } + } /** @@ -155,6 +195,10 @@ public class DetectorModel implements AutoCloseable{ return ImageIO.read(new ByteArrayInputStream(imageBytes)); } catch (IOException e) { throw new DetectionException("导出图片失败", e); + } finally { + if (img != null) { + ((Mat)img.getWrappedImage()).release(); + } } } @@ -163,11 +207,13 @@ public class DetectorModel implements AutoCloseable{ * @param image * @return */ - private DetectedObjects detect(Image image){ + public DetectedObjects detect(Image image){ Predictor predictor = null; try { predictor = predictorPool.borrowObject(); - return predictor.predict(image); + DetectedObjects detectedObjects = predictor.predict(image); + detectedObjects = filterDetections(detectedObjects); + return detectedObjects; } catch (Exception e) { throw new DetectionException("目标检测错误", e); }finally { @@ -187,6 +233,50 @@ public class DetectorModel implements AutoCloseable{ } } + /** + * 筛选检测结果 + * @param detectedObjects + * @return + */ + private DetectedObjects filterDetections(DetectedObjects detectedObjects) { + if (Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0) { + return detectedObjects; + } + List items = detectedObjects.items(); + // 按照允许的类别进行过滤 + List filtered = new ArrayList<>(); + //过滤类别 + if(!CollectionUtils.isEmpty(config.getAllowedClasses())){ + for (DetectedObjects.DetectedObject obj : items) { + if(config.getAllowedClasses().contains(obj.getClassName())){ + filtered.add(obj); + } + } + }else{ + filtered = items; + } + // 按照概率进行排序 + filtered.sort((o1, o2) -> Double.compare(o2.getProbability(), o1.getProbability())); + if(config.getTopK() > 0 && filtered.size() > config.getTopK()){ + filtered = filtered.subList(0, config.getTopK()); + } + // 构建新的 DetectedObjects 返回 + List names = new ArrayList<>(); + List probs = new ArrayList<>(); + List boxes = new ArrayList<>(); + + for (DetectedObjects.DetectedObject obj : filtered) { + names.add(obj.getClassName()); + probs.add(obj.getProbability()); + boxes.add(obj.getBoundingBox()); + } + return new DetectedObjects(names, probs, boxes); + } + + public GenericObjectPool> getPool() { + return predictorPool; + } + /** * 显式释放资源 diff --git a/smartjavaai-ocr/pom.xml b/smartjavaai-ocr/pom.xml index 70315cd..12effac 100644 --- a/smartjavaai-ocr/pom.xml +++ b/smartjavaai-ocr/pom.xml @@ -6,7 +6,7 @@ cn.smartjavaai smartjavaai-parent - 1.0.20 + 1.0.22 smartjavaai-ocr @@ -42,7 +42,7 @@ - 1.0.20 + 1.0.22 smartjavaai-ocr SmartJavaAI https://github.com/geekwenjie/SmartJavaAI diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/PlateDetModelConfig.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/PlateDetModelConfig.java new file mode 100644 index 0000000..2c586be --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/PlateDetModelConfig.java @@ -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; + + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/PlateRecModelConfig.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/PlateRecModelConfig.java new file mode 100644 index 0000000..a8af403 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/PlateRecModelConfig.java @@ -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; + + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/PlateInfo.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/PlateInfo.java new file mode 100644 index 0000000..aa8d7d9 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/PlateInfo.java @@ -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; + + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/PlateResult.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/PlateResult.java new file mode 100644 index 0000000..45cfb22 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/PlateResult.java @@ -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 + '\'' + + '}'; + } +} + diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/PlateDetModelEnum.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/PlateDetModelEnum.java new file mode 100644 index 0000000..9e68368 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/PlateDetModelEnum.java @@ -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); + } + + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/PlateRecModelEnum.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/PlateRecModelEnum.java new file mode 100644 index 0000000..26a27e1 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/PlateRecModelEnum.java @@ -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); + } + + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/PlateType.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/PlateType.java new file mode 100644 index 0000000..9703d0c --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/PlateType.java @@ -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; + } + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/factory/PlateModelFactory.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/factory/PlateModelFactory.java new file mode 100644 index 0000000..8ffe729 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/factory/PlateModelFactory.java @@ -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 detModelMap = new ConcurrentHashMap<>(); + + /** + * 模型缓存 + */ + private static final ConcurrentHashMap recModelMap = new ConcurrentHashMap<>(); + + + /** + * 模型注册表 + */ + private static final Map> detModelRegistry = + new ConcurrentHashMap<>(); + + /** + * 模型注册表 + */ + private static final Map> 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 clazz) { + detModelRegistry.put(plateDetModelEnum, clazz); + } + + /** + * 注册模型 + * @param plateRecModelEnum + * @param clazz + */ + private static void registerRecModel(PlateRecModelEnum plateRecModelEnum, Class 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()); + } + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/OcrCommonDetModel.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/OcrCommonDetModel.java index aa29177..214775e 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/OcrCommonDetModel.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/OcrCommonDetModel.java @@ -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> getPool(){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + } diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/OcrCommonDetModelImpl.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/OcrCommonDetModelImpl.java index 648c68c..369883f 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/OcrCommonDetModelImpl.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/OcrCommonDetModelImpl.java @@ -43,7 +43,7 @@ import java.util.*; @Slf4j public class OcrCommonDetModelImpl implements OcrCommonDetModel{ - private ObjectPool> detPredictorPool; + private GenericObjectPool> detPredictorPool; private ZooModel 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> getPool() { + return detPredictorPool; + } + @Override public void close() throws Exception { try { diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/criteria/OcrCommonDetCriterialFactory.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/criteria/OcrCommonDetCriterialFactory.java index 74a6059..bbf9708 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/criteria/OcrCommonDetCriterialFactory.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/criteria/OcrCommonDetCriterialFactory.java @@ -25,7 +25,7 @@ public class OcrCommonDetCriterialFactory { public static Criteria 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 criteria = null; ConcurrentHashMap params = new ConcurrentHashMap(); diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/OcrDirectionModel.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/OcrDirectionModel.java index ef2e8a4..7e96f1c 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/OcrDirectionModel.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/OcrDirectionModel.java @@ -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> getPool() { + throw new UnsupportedOperationException("默认不支持该功能"); + } + } diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/PPOCRMobileV2ClsModel.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/PPOCRMobileV2ClsModel.java index 483ada0..3567460 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/PPOCRMobileV2ClsModel.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/PPOCRMobileV2ClsModel.java @@ -55,7 +55,7 @@ import java.util.concurrent.ConcurrentHashMap; public class PPOCRMobileV2ClsModel implements OcrDirectionModel { - private ObjectPool> predictorPool; + private GenericObjectPool> 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(); 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> getPool() { + return predictorPool; + } + @Override public void close() throws Exception { try { diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/criteria/DirectionCriteriaFactory.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/criteria/DirectionCriteriaFactory.java index 540b698..566831c 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/criteria/DirectionCriteriaFactory.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/criteria/DirectionCriteriaFactory.java @@ -28,7 +28,7 @@ public class DirectionCriteriaFactory { public static Criteria 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 criteria = null; ConcurrentHashMap params = new ConcurrentHashMap(); diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/OcrCommonRecModel.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/OcrCommonRecModel.java index e9ac8b2..8b4a4cd 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/OcrCommonRecModel.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/OcrCommonRecModel.java @@ -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> getPool() { + throw new UnsupportedOperationException("默认不支持该功能"); + } + } diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/OcrCommonRecModelImpl.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/OcrCommonRecModelImpl.java index 775efd0..7fd4806 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/OcrCommonRecModelImpl.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/OcrCommonRecModelImpl.java @@ -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> recPredictorPool; + private GenericObjectPool> 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> getRecPredictorPool() { + return recPredictorPool; + } + @Override public void close() throws Exception { try { diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/criteria/OcrCommonRecCriterialFactory.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/criteria/OcrCommonRecCriterialFactory.java index 6eb61e5..93e07ce 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/criteria/OcrCommonRecCriterialFactory.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/criteria/OcrCommonRecCriterialFactory.java @@ -24,7 +24,7 @@ public class OcrCommonRecCriterialFactory { public static Criteria 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 criteria = null; ConcurrentHashMap params = new ConcurrentHashMap(); diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/CRNNPlateRecModel.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/CRNNPlateRecModel.java new file mode 100644 index 0000000..5fb00c9 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/CRNNPlateRecModel.java @@ -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> recPredictorPool; + + private ZooModel 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 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> 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> plateResult = recognize(img); + return plateResult; + } catch (IOException e) { + throw new OcrException("无效的图片", e); + } finally { + ((Mat)img.getWrappedImage()).release(); + } + } + + @Override + public R> 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> recognize(BufferedImage image) { + if(!ImageUtils.isImageValid(image)){ + return R.fail(R.Status.INVALID_IMAGE); + } + Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image)); + R> plateResult = recognize(img); + ((Mat)img.getWrappedImage()).release(); + return plateResult; + } + + @Override + public R> recognize(byte[] imageData) { + if(Objects.isNull(imageData)){ + return R.fail(R.Status.INVALID_IMAGE); + } + return recognize(new ByteArrayInputStream(imageData)); + } + + @Override + public R> 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 plateInfoList = OcrUtils.convertToPlateInfo(detectedObjects, image); + Predictor 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 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 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> 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 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> 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 recognizeAndDraw(BufferedImage sourceImage) { + if(!ImageUtils.isImageValid(sourceImage)){ + return R.fail(R.Status.INVALID_IMAGE); + } + Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(sourceImage)); + try { + R> 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> 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); + } + } +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/PlateDetModel.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/PlateDetModel.java new file mode 100644 index 0000000..2f31059 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/PlateDetModel.java @@ -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> detect(String imagePath) { + throw new UnsupportedOperationException("默认不支持该功能"); + } + + + /** + * 车牌检测 + * @param inputStream + * @return + */ + default R> detect(InputStream inputStream) { + throw new UnsupportedOperationException("默认不支持该功能"); + } + + + /** + * 车牌检测 + * @param base64Image + * @return + */ + default R> detectBase64(String base64Image){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + + + /** + * 车牌检测 + * @param image BufferedImage + * @return + */ + default R> detect(BufferedImage image) { + throw new UnsupportedOperationException("默认不支持该功能"); + } + + + /** + * 车牌检测 + * @param imageData 图片字节数组 + * @return + */ + default R> 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 detectAndDraw(String imagePath, String outputPath) { + throw new UnsupportedOperationException("默认不支持该功能"); + } + + /** + * 检测并绘制结果 + * @param sourceImage + * @return + */ + default R detectAndDraw(BufferedImage sourceImage){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + + default GenericObjectPool> getPool(){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + + + + + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/PlateRecModel.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/PlateRecModel.java new file mode 100644 index 0000000..68b27d0 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/PlateRecModel.java @@ -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> recognize(String imagePath) { + throw new UnsupportedOperationException("默认不支持该功能"); + } + + + /** + * 车牌识别 + * @param inputStream + * @return + */ + default R> recognize(InputStream inputStream) { + throw new UnsupportedOperationException("默认不支持该功能"); + } + + + /** + * 车牌识别 + * @param base64Image + * @return + */ + default R> recognizeBase64(String base64Image){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + + + /** + * 车牌识别 + * @param image BufferedImage + * @return + */ + default R> recognize(BufferedImage image) { + throw new UnsupportedOperationException("默认不支持该功能"); + } + + + /** + * 车牌识别 + * @param imageData 图片字节数组 + * @return + */ + default R> recognize(byte[] imageData) { + throw new UnsupportedOperationException("默认不支持该功能"); + } + + /** + * 车牌识别 + * @param image DJL Image + * @return + */ + default R> recognize(Image image){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + + /** + * 识别裁剪后的图片 + * @return + */ + default PlateResult recognizeCropped(Image image){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + + + /** + * 检测并绘制结果 + * @param imagePath 图片输入路径（包含文件名称） + * @param outputPath 图片输出路径（包含文件名称） + */ + default R recognizeAndDraw(String imagePath, String outputPath) { + throw new UnsupportedOperationException("默认不支持该功能"); + } + + /** + * 检测并绘制结果 + * @param sourceImage + * @return + */ + default R recognizeAndDraw(BufferedImage sourceImage){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + + default GenericObjectPool> getPool() { + throw new UnsupportedOperationException("默认不支持该功能"); + } + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/Yolov5PlateDetModel.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/Yolov5PlateDetModel.java new file mode 100644 index 0000000..be221e6 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/Yolov5PlateDetModel.java @@ -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> detPredictorPool; + + private ZooModel 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 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> 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 plateInfoList = OcrUtils.convertToPlateInfo(detectedObjects, img); + ((Mat)img.getWrappedImage()).release(); + return R.ok(plateInfoList); + } + + @Override + public R> 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> 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 plateInfoList = OcrUtils.convertToPlateInfo(detectedObjects, img); + ((Mat)img.getWrappedImage()).release(); + return R.ok(plateInfoList); + } + + @Override + public R> 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 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> 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 plateInfoList = OcrUtils.convertToPlateInfo(detection, img); + ((Mat)img.getWrappedImage()).release(); + return R.ok(plateInfoList); + } catch (IOException e) { + throw new OcrException("无效图片输入流", e); + } + } + + @Override + public R 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 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> 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); + } + } +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/criteria/PlateDetCriterialFactory.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/criteria/PlateDetCriterialFactory.java new file mode 100644 index 0000000..2b45326 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/criteria/PlateDetCriterialFactory.java @@ -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 createCriteria(PlateDetModelConfig config) { + Device device = null; + if(!Objects.isNull(config.getDevice())){ + device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId()); + } + Criteria criteria = null; + ConcurrentHashMap params = new ConcurrentHashMap(); + 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; + } +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/criteria/PlateRecCriterialFactory.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/criteria/PlateRecCriterialFactory.java new file mode 100644 index 0000000..ef6c553 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/criteria/PlateRecCriterialFactory.java @@ -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 createCriteria(PlateRecModelConfig config) { + Device device = null; + if(!Objects.isNull(config.getDevice())){ + device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId()); + } + Criteria 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; + } +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/CRNNPlateRecTranslator.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/CRNNPlateRecTranslator.java new file mode 100644 index 0000000..1614bc7 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/CRNNPlateRecTranslator.java @@ -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 { + + 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 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; // 非批量任务 + } +} + diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/Yolo5PlateDetectTranslator.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/Yolo5PlateDetectTranslator.java new file mode 100644 index 0000000..899b964 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/Yolo5PlateDetectTranslator.java @@ -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 { + + 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 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 classNames = new ArrayList<>(); + List probabilities = new ArrayList<>(); + List 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 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; + } + + + + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/Yolov7PlateDetectTranslator.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/Yolov7PlateDetectTranslator.java new file mode 100644 index 0000000..6c47fb3 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/Yolov7PlateDetectTranslator.java @@ -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 { + + 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 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 classNames = new ArrayList<>(); + List probabilities = new ArrayList<>(); + List 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 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; + } + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/Yolov8PlateDetectTranslator.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/Yolov8PlateDetectTranslator.java new file mode 100644 index 0000000..55ed1f3 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/plate/translator/Yolov8PlateDetectTranslator.java @@ -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 { + + 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 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 classNames = new ArrayList<>(); + List probabilities = new ArrayList<>(); + List 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 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); + } + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/CommonTableStructureModel.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/CommonTableStructureModel.java index a8a506b..55c2d4e 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/CommonTableStructureModel.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/CommonTableStructureModel.java @@ -47,7 +47,7 @@ public class CommonTableStructureModel implements TableStructureModel{ private ZooModel model; - private ObjectPool> predictorPool; + private GenericObjectPool> 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> getPool() { + return predictorPool; + } + @Override public void close() throws Exception { try { diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/TableStructureModel.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/TableStructureModel.java index a1ea4e3..098d593 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/TableStructureModel.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/TableStructureModel.java @@ -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 detect(Image image){ throw new UnsupportedOperationException("默认不支持该功能"); } + + default GenericObjectPool> getPool() { + throw new UnsupportedOperationException("默认不支持该功能"); + } } diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/opencv/OcrOpenCVUtils.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/opencv/OcrOpenCVUtils.java deleted file mode 100644 index 38af44e..0000000 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/opencv/OcrOpenCVUtils.java +++ /dev/null @@ -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; - } -} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/utils/OcrUtils.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/utils/OcrUtils.java index 797181b..d392f34 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/utils/OcrUtils.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/utils/OcrUtils.java @@ -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 landMarks){ + if (landMarks == null || landMarks.size() != 4) { + throw new IllegalArgumentException("必须提供4个关键点"); + } + + // 步骤 1：排序为 左上、右上、右下、左下 + List 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 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 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 convertToPlateInfo(DetectedObjects detectedObjects, Image image) { + List 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 keyPoints = new ArrayList(); + 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 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 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(); + } + } + } diff --git a/smartjavaai-translate/pom.xml b/smartjavaai-translate/pom.xml index d856263..e861921 100644 --- a/smartjavaai-translate/pom.xml +++ b/smartjavaai-translate/pom.xml @@ -20,7 +20,7 @@ - 1.0.20 + 1.0.22 smartjavaai-translate SmartJavaAI https://github.com/geekwenjie/SmartJavaAI diff --git a/smartjavaai-translate/src/main/java/cn/smartjavaai/translation/config/TranslationModelConfig.java b/smartjavaai-translate/src/main/java/cn/smartjavaai/translation/config/TranslationModelConfig.java index 55e6b7f..14b809f 100644 --- a/smartjavaai-translate/src/main/java/cn/smartjavaai/translation/config/TranslationModelConfig.java +++ b/smartjavaai-translate/src/main/java/cn/smartjavaai/translation/config/TranslationModelConfig.java @@ -1,5 +1,6 @@ package cn.smartjavaai.translation.config; +import cn.smartjavaai.common.config.ModelConfig; import cn.smartjavaai.common.enums.DeviceEnum; import cn.smartjavaai.translation.enums.TranslationModeEnum; @@ -11,17 +12,12 @@ import lombok.Data; * @date 2025/6/05 */ @Data -public class TranslationModelConfig { +public class TranslationModelConfig extends ModelConfig { /** * 翻译模型 */ private TranslationModeEnum modelEnum; - /** - * 设备类型 - */ - private DeviceEnum device; - /** * 翻译模型路径 */ diff --git a/smartjavaai-translate/src/main/java/cn/smartjavaai/translation/model/NllbModel.java b/smartjavaai-translate/src/main/java/cn/smartjavaai/translation/model/NllbModel.java index daff744..71653f7 100644 --- a/smartjavaai-translate/src/main/java/cn/smartjavaai/translation/model/NllbModel.java +++ b/smartjavaai-translate/src/main/java/cn/smartjavaai/translation/model/NllbModel.java @@ -48,11 +48,11 @@ import java.util.Objects; @Slf4j public class NllbModel implements TranslationModel{ - private ObjectPool> encodePredictorPool; + private GenericObjectPool> encodePredictorPool; - private ObjectPool> decodePredictorPool; + private GenericObjectPool> decodePredictorPool; - private ObjectPool> decode2PredictorPool; + private GenericObjectPool> decode2PredictorPool; private ZooModel nllbModel; private HuggingFaceTokenizer tokenizer; @@ -69,7 +69,7 @@ public class NllbModel implements TranslationModel{ } 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()); } this.config = config; Path modelPath = Paths.get(config.getModelPath()); @@ -91,8 +91,16 @@ public class NllbModel implements TranslationModel{ tokenizer = HuggingFaceTokenizer.newInstance(tokenizerPath); //初始化searchConfig this.searchConfig = new NllbSearchConfig(); + int predictorPoolSize = config.getPredictorPoolSize(); + if(config.getPredictorPoolSize() <= 0){ + predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数 + } + encodePredictorPool.setMaxTotal(predictorPoolSize); + decodePredictorPool.setMaxTotal(predictorPoolSize); + decode2PredictorPool.setMaxTotal(predictorPoolSize); log.debug("当前设备: " + nllbModel.getNDManager().getDevice()); log.debug("当前引擎: " + Engine.getInstance().getEngineName()); + log.debug("模型推理器线程池最大数量: " + predictorPoolSize); } catch (IOException | ModelNotFoundException | MalformedModelException e) { throw new TranslationException("模型加载失败", e); } @@ -252,4 +260,54 @@ public class NllbModel implements TranslationModel{ return next_tokens.expandDims(0); } + public GenericObjectPool> getEncodePredictorPool() { + return encodePredictorPool; + } + + public GenericObjectPool> getDecodePredictorPool() { + return decodePredictorPool; + } + + public GenericObjectPool> getDecode2PredictorPool() { + return decode2PredictorPool; + } + + @Override + public void close() throws Exception { + try { + if (nllbModel != null) { + nllbModel.close(); + } + } catch (Exception e) { + log.warn("关闭 model 失败", e); + } + try { + if (tokenizer != null) { + tokenizer.close(); + } + } catch (Exception e) { + log.warn("关闭 tokenizer 失败", e); + } + try { + if (encodePredictorPool != null) { + encodePredictorPool.close(); + } + } catch (Exception e) { + log.warn("关闭 encodePredictorPool 失败", e); + } + try { + if (decodePredictorPool != null) { + decodePredictorPool.close(); + } + } catch (Exception e) { + log.warn("关闭 decodePredictorPool 失败", e); + } + try { + if (decode2PredictorPool != null) { + decode2PredictorPool.close(); + } + } catch (Exception e) { + log.warn("关闭 decode2PredictorPool 失败", e); + } + } } diff --git a/smartjavaai-translate/src/main/java/cn/smartjavaai/translation/model/TranslationModel.java b/smartjavaai-translate/src/main/java/cn/smartjavaai/translation/model/TranslationModel.java index a663f02..f0b35d2 100644 --- a/smartjavaai-translate/src/main/java/cn/smartjavaai/translation/model/TranslationModel.java +++ b/smartjavaai-translate/src/main/java/cn/smartjavaai/translation/model/TranslationModel.java @@ -9,7 +9,7 @@ import cn.smartjavaai.translation.entity.TranslateParam; * @author lwx * @date 2025/6/05 */ -public interface TranslationModel { +public interface TranslationModel extends AutoCloseable{ /** * 加载模型