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

OCR文字识别

+ - 表格识别
+
+ + +
+ +
+ +
@@ -222,23 +235,25 @@ SmartJavaAI是专为JAVA 开发者打造的一个功能丰富、开箱即用的 - 支持多种主流模型：兼容 YOLOv3、YOLOv5、YOLOv8、YOLOv11、YOLOv12、SSD 等目标检测算法 - 支持自定义模型加载：可无缝加载并部署用户自行训练的目标检测模型 - **OCR文字识别** - - 支持PaddleOCR 3.0模型：集成最新PP-OCRv5模型 + - 支持PaddleOCR 3.0模型：集成最新PP-OCRv5、PP-OCRv4、表格结构识别模型(SLANet_plus)、文本行方向分类模型 - 支持任意角度识别，方向校准 - 支持通用文字识别，通用手写字识别 + - 支持表格识别 - **机器翻译** - 集成NLLB-200模型：支持200+语言互相翻译 ## 🌟 AI集成方式对比 -| 方案 | 技术特点 | 优点 | 缺点 | -|---------------------|--------------------------------------------------------------------------|----------------------------------------------------------------------|----------------------------------------------------------------------| -| **OpenCV** | 传统图像处理方案 | ✅ 提供java接口
✅ 轻量级部署
✅ 社区资源丰富 | ❌ 基于传统算法精度低(60%-75%)
❌ 需本地安装环境 | -| ‌**虹软SDK**‌ | 商业级闭源解决方案 | ✅ 开箱即用
✅ 提供完整文档和SDK
✅ 支持离线活体检测 | ❌ 免费版需年度授权更新
❌ 商业授权费用高
❌ 代码不可控 | +| 方案 | 技术特点 | 优点 | 缺点 | +|---------------------|--------------------------------------------------------------------------|---------------------------------------------------------------------|---------------------------------------------------------------------| +| **OpenCV** | 传统图像处理方案 | ✅ 提供java接口
✅ 轻量级部署
✅ 社区资源丰富 | ❌ 基于传统算法精度低(60%-75%)
❌ 需本地安装环境 | +| ‌**商业闭源SDK（如虹软等）**‌ | 商业级闭源解决方案 | ✅ 开箱即用
✅ 提供完整文档和SDK
✅ 支持离线活体检测 | ❌ 免费版需年度授权更新
❌ 商业授权费用高
❌ 代码不可控 | | ‌**云API(阿里云)**‌ | SaaS化云端服务 | ✅ 零部署成本
✅ 支持高并发
✅ 自带模型迭代 | ❌ 网络延迟风险(200-800ms)
❌ 按调用量计费
❌ 有数据安全风险 | | ‌**Python混合调用**‌ | 跨语言调用方案 | ✅ 可集成PyTorch/TF等框架
✅ 支持自定义算法
✅ 识别精度高 | ❌ 需维护双语言环境
❌ 进程通信性能损耗(30%+)
❌ 异常处理复杂度翻倍 | -| ‌**DJL框架**‌ | 深度学习框架 | ✅ 纯Java实现
✅ 支持主流深度学习框架
✅ 可加载预训练模型(99%+) | ❌ 需掌握DL知识
❌ 需处理模型加载、预处理、后处理等复杂技术细节 | -| ‌**SmartJavaAI**‌ | java深度学习工具包 | ✅ 支持主流深度学习框架
✅ 提供丰富、开箱即用API
✅ 上手简单，单一Jar包集成 | ❌要求JDK版本11及以上 | +| ‌**JNI/JNA**‌ | 跨语言底层调用方案 | ✅ 直接调用 C/C++ 高性能算法库✅ 支持调用各种原生成熟库✅ 可封装成通用工具Jar |❌ 开发成本高，JNI更复杂❌ 跨平台兼容性差 | +| ‌**DJL框架**‌ | 深度学习框架 | ✅ 纯Java实现
✅ 支持主流深度学习框架
✅ 可加载预训练模型(99%+) | ❌ 需掌握DL知识
❌ 需处理模型加载、预处理、后处理等复杂技术细节 | +| ‌**SmartJavaAI**‌ | java深度学习工具包 | ✅ 支持主流深度学习框架
✅ 提供丰富、开箱即用API
✅ 上手简单，单一Jar包集成 | ❌要求JDK版本11及以上 | @@ -273,7 +288,7 @@ SmartJavaAI是专为JAVA 开发者打造的一个功能丰富、开箱即用的 cn.smartjavaai smartjavaai-all - 1.0.19 + 1.0.20 ``` ### 3、完整示例代码 @@ -427,19 +442,37 @@ SmartJavaAI是专为JAVA 开发者打造的一个功能丰富、开箱即用的 | 模型名称 | 模型简介 | 模型开源网站 | | ------------| ------------------- |--------------------------------------------------------------------------------------------------------------| -| PP-OCRv5_server_det | 飞桨PaddleOCR 3.0 | [Github](https://github.com/PaddlePaddle/PaddleOCR/blob/main/docs/version3.x/algorithm/PP-OCRv5/PP-OCRv5.md) | +| PP-OCRv5_server_det | 服务端文本检测模型，精度更高，适合在性能较好的服务器上部署 | [Github](https://github.com/PaddlePaddle/PaddleOCR/blob/main/docs/version3.x/algorithm/PP-OCRv5/PP-OCRv5.md) | +| PP-OCRv5_mobile_det | 轻量文本检测模型，效率更高，适合在端侧设备部署 | [Github](https://github.com/PaddlePaddle/PaddleOCR/blob/main/docs/version3.x/algorithm/PP-OCRv5/PP-OCRv5.md) | +| PP-OCRv4_server_det | 服务端文本检测模型，精度更高，适合在性能较好的服务器上部署 | [Github](https://github.com/PaddlePaddle/PaddleOCR/blob/main/docs/version3.x/algorithm/PP-OCRv5/PP-OCRv5.md) | +| PP-OCRv4_mobile_det | 轻量文本检测模型，效率更高，适合在端侧设备部署 | [Github](https://github.com/PaddlePaddle/PaddleOCR/blob/main/docs/version3.x/algorithm/PP-OCRv5/PP-OCRv5.md) | **文本识别模型** -| 模型名称 | 模型简介 | 模型开源网站 | -| ------------| ------------------- |--------------------------------------------------------------------------------------------------------------| -| PP-OCRv5_server_rec | 飞桨PaddleOCR 3.0 | [Github](https://github.com/PaddlePaddle/PaddleOCR/blob/main/docs/version3.x/algorithm/PP-OCRv5/PP-OCRv5.md) | +| 模型名称 | 模型简介 | 模型开源网站 | +| ------------|----------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------| +| PP-OCRv5_server_rec | （服务端）致力于以单一模型高效、精准地支持简体中文、繁体中文、英文、日文四种主要语言，以及手写、竖版、拼音、生僻字等复杂文本场景的识别。在保持识别效果的同时，兼顾推理速度和模型鲁棒性，为各种场景下的文档理解提供高效、精准的技术支撑。 | [Github](https://github.com/PaddlePaddle/PaddleOCR/blob/main/docs/version3.x/algorithm/PP-OCRv5/PP-OCRv5.md) | +| PP-OCRv5_mobile_rec | （轻量）致力于以单一模型高效、精准地支持简体中文、繁体中文、英文、日文四种主要语言，以及手写、竖版、拼音、生僻字等复杂文本场景的识别。在保持识别效果的同时，兼顾推理速度和模型鲁棒性，为各种场景下的文档理解提供高效、精准的技术支撑。 | [Github](https://github.com/PaddlePaddle/PaddleOCR/blob/main/docs/version3.x/algorithm/PP-OCRv5/PP-OCRv5.md) | +| PP-OCRv4_server_rec | （服务端）推理精度高，可以部署在多种不同的服务器上 | [Github](https://github.com/PaddlePaddle/PaddleOCR/blob/main/docs/version3.x/algorithm/PP-OCRv5/PP-OCRv5.md) | +| PP-OCRv4_mobile_rec | （轻量） 效率更高，适合在端侧设备部署 | [Github](https://github.com/PaddlePaddle/PaddleOCR/blob/main/docs/version3.x/algorithm/PP-OCRv5/PP-OCRv5.md) | + **文本方向分类模型(cls)** -| 模型名称 | 模型简介 | 模型开源网站 | -| ------------| ------------------- |------------| -| ch_ppocr_mobile_v2.0_cls | 原始分类器模型，对检测到的文本行文字角度分类 | [Github](https://paddlepaddle.github.io/PaddleOCR/latest/en/version2.x/ppocr/model_list.html?h=models+list#13-multilingual-detection-model) | +| 模型名称 | 模型简介 | 模型开源网站 | +| ------------|------------------------------|------------| +| ch_ppocr_mobile_v2.0_cls | 原始分类器模型，对检测到的文本行文字角度分类 | [Github](https://paddlepaddle.github.io/PaddleOCR/latest/en/version2.x/ppocr/model_list.html?h=models+list#13-multilingual-detection-model) | +| PP_LCNET_X0_25 | （轻量）基于PP-LCNet_x0_25的文本行分类模型 | [Github](https://paddlepaddle.github.io/PaddleOCR/v3.1.0/version3.x/module_usage/textline_orientation_classification.html) | +| PP_LCNET_X1_0 | 基于PP-LCNet_x1_0的文本行分类模型 | [Github](https://paddlepaddle.github.io/PaddleOCR/v3.1.0/version3.x/module_usage/textline_orientation_classification.html) | + + +**表格结构识别(Table Structure Recognition)** + +| 模型名称 | 模型简介 | 模型开源网站 | +| ------------|------------------------------|------------| +| SLANet | 该模型通过轻量级骨干 PP-LCNet、CSP-PAN 融合与 SLA Head 解码，有效提升表格结构识别的精度与速度。 | [Github](https://paddlepaddle.github.io/PaddleOCR/v3.1.0/version3.x/module_usage/table_structure_recognition.html#_3) | +| 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) | + --- @@ -474,7 +507,7 @@ SmartJavaAI是专为JAVA 开发者打造的一个功能丰富、开箱即用的 ## 近期更新日志 -## [v1.0.19] - 2025-07-06 +## [v1.0.20] - 2025-07-06 - 人脸模块：新增小视科技（MiniVision）活体检测模型 - 人脸模块：新增阿里通义工作室活体检测模型 - 人脸模块：新增 2 个表情识别模型 diff --git a/examples/face-example/pom.xml b/examples/face-example/pom.xml index 3c21eef..ff2d93e 100644 --- a/examples/face-example/pom.xml +++ b/examples/face-example/pom.xml @@ -12,7 +12,7 @@ 11 11 UTF-8 - 1.0.19 + 1.0.20 smartai.examples.face.facedet.FaceDetDemo 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 b6670d4..031c1e5 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 @@ -20,6 +20,7 @@ import cn.smartjavaai.face.model.facerec.FaceRecModel; import cn.smartjavaai.face.utils.SimilarityUtil; import cn.smartjavaai.face.vector.config.MilvusConfig; import cn.smartjavaai.face.vector.config.SQLiteConfig; +import cn.smartjavaai.face.vector.entity.FaceVector; import com.alibaba.fastjson.JSONArray; import com.alibaba.fastjson.JSONObject; import lombok.extern.slf4j.Slf4j; @@ -360,6 +361,53 @@ public class FaceRecDemo { } + /** + * 获取人脸信息 + */ + @Test + public void getFaceInfo(){ + //使用ID获取人脸信息 + try (FaceRecModel faceRecModel = getFaceRecModelWithSQLiteConfig()){ + //等待加载人脸库结束 + while (!faceRecModel.isLoadFaceCompleted()){ + Thread.sleep(100); + } + //ID需改为你需要查询的ID + R faceInfoResult = faceRecModel.getFaceInfoById("9c4c316d53a74b1184195c1714c250c4"); + if(faceInfoResult.isSuccess()){ + log.info("人脸信息：{}", JSONObject.toJSONString(faceInfoResult.getData())); + }else{ + log.info("获取人脸信息失败：{}", faceInfoResult.getMessage()); + } + } catch (Exception e) { + e.printStackTrace(); + } + } + + /** + * 获取人脸信息 + */ + @Test + public void listFaces(){ + //使用ID获取人脸信息 + try (FaceRecModel faceRecModel = getFaceRecModelWithDbConfig()){ + //等待加载人脸库结束 + while (!faceRecModel.isLoadFaceCompleted()){ + Thread.sleep(100); + } + //ID需改为你需要查询的ID + R> faceInfoResult = faceRecModel.listFaces(1, 10); + if(faceInfoResult.isSuccess()){ + log.info("人脸信息：{}", JSONObject.toJSONString(faceInfoResult.getData())); + }else{ + log.info("获取人脸信息失败：{}", faceInfoResult.getMessage()); + } + } catch (Exception e) { + e.printStackTrace(); + } + } + + } 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 1c807b4..bf1f966 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 @@ -79,6 +79,8 @@ public class LivenessDetDemo { 这个数量相当于多帧识别结果融合的融合的帧数。当输入的帧数超过设定帧数的时候，会采用滑动窗口的方式，返回融合的最近输入的帧融合的识别结果。 一般来说，在10以内，帧数越多，结果越稳定，相对性能越好，但是得到结果的延时越高。*/ config.setFrameCount(LivenessConstant.DEFAULT_FRAME_COUNT); + //视频最大检测帧数 + config.setMaxVideoDetectFrames(LivenessConstant.DEFAULT_MAX_VIDEO_DETECT_FRAMES); //指定人脸检测模型 config.setDetectModel(getFaceDetModel()); return LivenessModelFactory.getInstance().getModel(config); @@ -103,6 +105,8 @@ public class LivenessDetDemo { 这个数量相当于多帧识别结果融合的融合的帧数。当输入的帧数超过设定帧数的时候，会采用滑动窗口的方式，返回融合的最近输入的帧融合的识别结果。 一般来说，在10以内，帧数越多，结果越稳定，相对性能越好，但是得到结果的延时越高。*/ config.setFrameCount(LivenessConstant.DEFAULT_FRAME_COUNT); + //视频最大检测帧数 + config.setMaxVideoDetectFrames(LivenessConstant.DEFAULT_MAX_VIDEO_DETECT_FRAMES); //指定人脸检测模型 config.setDetectModel(getFaceDetModel()); return LivenessModelFactory.getInstance().getModel(config); diff --git a/examples/objectdetection-example/pom.xml b/examples/objectdetection-example/pom.xml index 7b5bdf9..7d3cdce 100644 --- a/examples/objectdetection-example/pom.xml +++ b/examples/objectdetection-example/pom.xml @@ -12,7 +12,7 @@ 11 11 UTF-8 - 1.0.19 + 1.0.20 smartai.examples.objectdetection.ObjectDetection diff --git a/examples/ocr-examples/pom.xml b/examples/ocr-examples/pom.xml index 5e8700b..3506826 100644 --- a/examples/ocr-examples/pom.xml +++ b/examples/ocr-examples/pom.xml @@ -12,9 +12,9 @@ 11 11 UTF-8 - 1.0.19 + 1.0.20 - smartai.examples.ocr.OcrRecognizeDemo + smartai.examples.ocr.common.OcrRecognizeDemo 1.5.10 diff --git a/examples/ocr-examples/src/main/java/smartai/examples/ocr/OcrDetectionDemo.java b/examples/ocr-examples/src/main/java/smartai/examples/ocr/OcrDetectionDemo.java deleted file mode 100644 index ead2d71..0000000 --- a/examples/ocr-examples/src/main/java/smartai/examples/ocr/OcrDetectionDemo.java +++ /dev/null @@ -1,76 +0,0 @@ -package smartai.examples.ocr; - -import cn.smartjavaai.common.entity.DetectionResponse; -import cn.smartjavaai.common.enums.DeviceEnum; -import cn.smartjavaai.ocr.config.OcrDetModelConfig; -import cn.smartjavaai.ocr.entity.OcrBox; -import cn.smartjavaai.ocr.enums.CommonDetModelEnum; -import cn.smartjavaai.ocr.factory.OcrModelFactory; -import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel; -import com.alibaba.fastjson.JSONObject; -import lombok.extern.slf4j.Slf4j; -import org.junit.Test; - -import javax.imageio.ImageIO; -import java.awt.image.BufferedImage; -import java.io.File; -import java.nio.file.Paths; -import java.util.List; - -/** - * OCR 文本检测 示例 - * 模型下载地址：https://pan.baidu.com/s/15Noz2xHQzqMQSl1B19BobQ?pwd=1234 提取码: 1234 - * @author dwj - */ -@Slf4j -public class OcrDetectionDemo { - - - //设备类型 - public static DeviceEnum device = DeviceEnum.CPU; - - - /** - * 文本检测 - * 检测图像中的文本区域，仅返回文本框位置，不识别文字内容 - * 模型需要放在单独文件夹 - */ - @Test - public void detect(){ - OcrDetModelConfig config = new OcrDetModelConfig(); - //指定检测模型 - config.setModelEnum(CommonDetModelEnum.PADDLEOCR_V5_DET_MODEL); - //指定模型位置，需要更改为自己的模型路径（下载地址请查看文档） - config.setDetModelPath("/Users/xxx/Documents/develop/ocr模型/PP-OCRv5_server_det_infer/PP-OCRv5_server_det.onnx"); - config.setDevice(device); - try (OcrCommonDetModel model = OcrModelFactory.getInstance().getDetModel(config)){ - List boxes = model.detect("src/main/resources/ocr_1.jpg"); - log.info("OCR检测结果：{}", JSONObject.toJSONString(boxes)); - } catch (Exception e) { - e.printStackTrace(); - } - } - - /** - * 文本检测并绘制结果 - * 检测图像中的文本区域，仅检测文本框位置，不识别文字内容 - * 模型需要放在单独文件夹 - */ - @Test - public void detectAndDraw(){ - OcrDetModelConfig config = new OcrDetModelConfig(); - //指定检测模型 - config.setModelEnum(CommonDetModelEnum.PADDLEOCR_V5_DET_MODEL); - //指定模型位置，需要更改为自己的模型路径（下载地址请查看文档） - config.setDetModelPath("/PP-OCRv5_server_det_infer/PP-OCRv5_server_det.onnx"); - try (OcrCommonDetModel model = OcrModelFactory.getInstance().getDetModel(config)){ - model.detectAndDraw("src/main/resources/ocr_1.jpg", "output/ocr_1_detected.jpg"); - } catch (Exception e) { - e.printStackTrace(); - } - } - - - - -} diff --git a/examples/ocr-examples/src/main/java/smartai/examples/ocr/OcrRecognizeDemo.java b/examples/ocr-examples/src/main/java/smartai/examples/ocr/OcrRecognizeDemo.java deleted file mode 100644 index 57b3310..0000000 --- a/examples/ocr-examples/src/main/java/smartai/examples/ocr/OcrRecognizeDemo.java +++ /dev/null @@ -1,145 +0,0 @@ -package smartai.examples.ocr; - -import cn.smartjavaai.common.enums.DeviceEnum; -import cn.smartjavaai.ocr.config.OcrDetModelConfig; -import cn.smartjavaai.ocr.config.OcrRecModelConfig; -import cn.smartjavaai.ocr.entity.OcrBox; -import cn.smartjavaai.ocr.entity.OcrInfo; -import cn.smartjavaai.ocr.enums.CommonDetModelEnum; -import cn.smartjavaai.ocr.enums.CommonRecModelEnum; -import cn.smartjavaai.ocr.enums.DirectionModelEnum; -import cn.smartjavaai.ocr.factory.OcrModelFactory; -import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel; -import cn.smartjavaai.ocr.model.common.recognize.OcrCommonRecModel; -import com.alibaba.fastjson.JSONObject; -import lombok.extern.slf4j.Slf4j; -import org.junit.Test; - -import java.io.File; -import java.util.List; - -/** - * OCR 文本识别 示例 - * 模型下载地址：https://pan.baidu.com/s/1MLfd73Vjdpnuls9-oqc9uw?pwd=1234 提取码: 1234 - * @author dwj - * @date 2025/5/25 - */ -@Slf4j -public class OcrRecognizeDemo { - - //设备类型 - public static DeviceEnum device = DeviceEnum.CPU; - - /** - * 获取通用识别模型（不带方向矫正） - * @return - */ - public OcrCommonRecModel getRecModel(){ - OcrRecModelConfig recModelConfig = new OcrRecModelConfig(); - //指定检测模型 - recModelConfig.setDetModelEnum(CommonDetModelEnum.PADDLEOCR_V5_DET_MODEL); - //指定检测模型位置，需要更改为自己的模型路径（下载地址请查看文档） - recModelConfig.setDetModelPath("/Users/xxx/Documents/develop/ocr模型/PP-OCRv5_server_det_infer/PP-OCRv5_server_det.onnx"); - //指定识别模型 - recModelConfig.setRecModelEnum(CommonRecModelEnum.PADDLEOCR_V5_REC_MODEL); - //指定识别模型位置，需要更改为自己的模型路径（下载地址请查看文档） - recModelConfig.setRecModelPath("/Users/xxx/Documents/develop/ocr模型/PP-OCRv5_server_rec_infer/PP-OCRv5_server_rec.onnx"); - recModelConfig.setDevice(device); - return OcrModelFactory.getInstance().getRecModel(recModelConfig); - } - - /** - * 获取通用识别模型(带方向矫正) - * @return - */ - public OcrCommonRecModel getRecModelWithDirection() { - OcrRecModelConfig recModelConfig = new OcrRecModelConfig(); - //指定检测模型 - recModelConfig.setDetModelEnum(CommonDetModelEnum.PADDLEOCR_V5_DET_MODEL); - //指定检测模型位置，需要更改为自己的模型路径（下载地址请查看文档） - recModelConfig.setDetModelPath("/Users/xxx/Documents/develop/ocr模型/PP-OCRv5_server_det_infer/PP-OCRv5_server_det.onnx"); - //指定识别模型 - recModelConfig.setRecModelEnum(CommonRecModelEnum.PADDLEOCR_V5_REC_MODEL); - //指定识别模型位置，需要更改为自己的模型路径（下载地址请查看文档） - recModelConfig.setRecModelPath("/Users/xxx/Documents/develop/ocr模型/PP-OCRv5_server_rec_infer/PP-OCRv5_server_rec.onnx"); - //指定方向检测模型 - recModelConfig.setDirectionModelEnum(DirectionModelEnum.CH_PPOCR_MOBILE_V2_CLS); - //指定方向模型位置，需要更改为自己的模型路径（下载地址请查看文档） - recModelConfig.setDirectionModelPath("/Users/xxx/Documents/develop/ocr模型/ch_ppocr_mobile_v2.0_cls.onnx"); - recModelConfig.setDevice(device); - return OcrModelFactory.getInstance().getRecModel(recModelConfig); - } - - - /** - * 文本识别 - * 支持简体中文、繁体中文、英文、日文四种主要语言，以及手写、竖版、拼音、生僻字 - * 流程：文本检测 -> 文本识别 - * 模型需要放在单独文件夹 - */ - @Test - public void recognize(){ - try (OcrCommonRecModel recModel = getRecModel()){ - OcrInfo ocrInfo = recModel.recognize("src/main/resources/ocr_2.jpg"); - log.info("OCR识别结果：{}", JSONObject.toJSONString(ocrInfo)); - } catch (Exception e) { - e.printStackTrace(); - } - } - - - /** - * 文本识别（手写字） - * 支持简体中文、繁体中文、英文、日文四种主要语言，以及手写、竖版、拼音、生僻字 - * 流程：文本检测 -> 文本识别 - * 模型需要放在单独文件夹 - */ - @Test - public void recognizeHandWriting(){ - try (OcrCommonRecModel recModel = getRecModel()){ - OcrInfo ocrInfo = recModel.recognize("src/main/resources/handwriting_1.jpg"); - log.info("OCR识别结果：{}", JSONObject.toJSONString(ocrInfo)); - } catch (Exception e) { - e.printStackTrace(); - } - } - - /** - * 文本识别（带方向矫正） - * 支持简体中文、繁体中文、英文、日文四种主要语言，以及手写、竖版、拼音、生僻字 - * 本方法支持多角度文字识别 - * 流程：文本检测 -> 方向检测 -> 方向矫正 -> 文本识别 - * 模型需要放在单独文件夹 - */ - @Test - public void recognize2(){ - try (OcrCommonRecModel recModel = getRecModelWithDirection()){ - OcrInfo ocrInfo = recModel.recognize("src/main/resources/ocr_4.jpg"); - log.info("OCR识别结果：{}", JSONObject.toJSONString(ocrInfo)); - } catch (Exception e) { - e.printStackTrace(); - } - } - - - - /** - * 文本识别并绘制结果 - * 支持简体中文、繁体中文、英文、日文四种主要语言，以及手写、竖版、拼音、生僻字 - * 流程：文本检测 -> 文本识别 - * 模型需要放在单独文件夹 - */ - @Test - public void recognizeAndDraw(){ - try (OcrCommonRecModel recModel = getRecModelWithDirection()){ - int fontSize = 25; - recModel.recognizeAndDraw("src/main/resources/ocr_4.jpg", "output/ocr_4_recognized.jpg", fontSize); - } catch (Exception e) { - e.printStackTrace(); - } - } - - - - -} 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 new file mode 100644 index 0000000..e77a31e --- /dev/null +++ b/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrDetectionDemo.java @@ -0,0 +1,113 @@ +package smartai.examples.ocr.common; + +import ai.djl.modality.cv.Image; +import cn.smartjavaai.common.entity.DetectionResponse; +import cn.smartjavaai.common.enums.DeviceEnum; +import cn.smartjavaai.common.utils.ImageUtils; +import cn.smartjavaai.ocr.config.OcrDetModelConfig; +import cn.smartjavaai.ocr.config.OcrRecOptions; +import cn.smartjavaai.ocr.entity.OcrBox; +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; +import org.junit.Test; +import org.opencv.core.Mat; + +import javax.imageio.ImageIO; +import java.awt.image.BufferedImage; +import java.io.File; +import java.nio.file.Paths; +import java.util.List; + +/** + * OCR 文本检测 示例 + * 模型下载地址：https://pan.baidu.com/s/15Noz2xHQzqMQSl1B19BobQ?pwd=1234 提取码: 1234 + * @author dwj + */ +@Slf4j +public class OcrDetectionDemo { + + + //设备类型 + public static DeviceEnum device = DeviceEnum.CPU; + + /** + * 获取文本检测模型 + * @return + */ + public OcrCommonDetModel getDetectionModel() { + OcrDetModelConfig config = new OcrDetModelConfig(); + //指定检测模型 + config.setModelEnum(CommonDetModelEnum.PP_OCR_V5_MOBILE_DET_MODEL); + //指定模型位置，需要更改为自己的模型路径（下载地址请查看文档） + config.setDetModelPath("/Users/xxx/Documents/develop/model/ocr/PP-OCRv5_mobile_det_infer/PP-OCRv5_mobile_det_infer.onnx"); + config.setDevice(device); + return OcrModelFactory.getInstance().getDetModel(config); + } + + + /** + * 文本检测 + * 检测图像中的文本区域，仅返回文本框位置，不识别文字内容 + * 注意事项： + * 1、批量检测时，模型应统一放在外层 try 中使用，避免重复加载，自动释放资源更安全。 + * 2、模型文件需要放在单独文件夹 + */ + @Test + public void detect(){ + try (OcrCommonDetModel model = getDetectionModel()){ + List boxes = model.detect("src/main/resources/ocr_1.jpg"); + log.info("OCR检测结果：{}", JSONObject.toJSONString(boxes)); + } catch (Exception e) { + e.printStackTrace(); + } + } + + /** + * 文本检测并绘制结果 + * 检测图像中的文本区域，仅检测文本框位置，不识别文字内容 + * 注意事项： + * 1、批量检测时，模型应统一放在外层 try 中使用，避免重复加载，自动释放资源更安全。 + * 2、模型文件需要放在单独文件夹 + */ + @Test + public void detectAndDraw(){ + try (OcrCommonDetModel model = getDetectionModel()){ + model.detectAndDraw("src/main/resources/ocr_1.jpg", "output/ocr_1_detected.jpg"); + } catch (Exception e) { + e.printStackTrace(); + } + } + + + /** + * 批量文本检测：批量检测要求图片宽高一致 + * 检测图像中的文本区域，仅返回文本框位置，不识别文字内容 + * 注意事项： + * 1、批量检测时，模型应统一放在外层 try 中使用，避免重复加载，自动释放资源更安全。 + * 2、模型文件需要放在单独文件夹 + */ + @Test + public void batchDetect(){ + try (OcrCommonDetModel model = getDetectionModel()){ + //批量检测要求图片宽高一致 + String folderPath = "/Users/xxx/Downloads/testing33"; + //读取文件夹中所有图片 + List images = ImageUtils.readImagesFromFolder(folderPath); + List> ocrResult = model.batchDetectDJLImage(images); + for(int i = 0; i < ocrResult.size(); i++){ + log.info("图片" + i + "文本检测结果：{}", JSONObject.toJSONString(ocrResult.get(i))); + } + } catch (Exception e) { + e.printStackTrace(); + } + } + + + +} diff --git a/examples/ocr-examples/src/main/java/smartai/examples/ocr/OcrDirectionDetDemo.java b/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrDirectionDetDemo.java similarity index 66% rename from examples/ocr-examples/src/main/java/smartai/examples/ocr/OcrDirectionDetDemo.java rename to examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrDirectionDetDemo.java index c2165fd..d9264fa 100644 --- a/examples/ocr-examples/src/main/java/smartai/examples/ocr/OcrDirectionDetDemo.java +++ b/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrDirectionDetDemo.java @@ -1,4 +1,4 @@ -package smartai.examples.ocr; +package smartai.examples.ocr.common; import cn.smartjavaai.common.enums.DeviceEnum; import cn.smartjavaai.ocr.config.DirectionModelConfig; @@ -17,7 +17,7 @@ import org.junit.Test; import java.util.List; /** - * OCR 文本方向检测 示例 + * OCR 行文本方向检测 示例 * 模型下载地址：https://pan.baidu.com/s/1MLfd73Vjdpnuls9-oqc9uw?pwd=1234 提取码: 1234 * @author dwj * @date 2025/5/25 @@ -34,18 +34,29 @@ public class OcrDirectionDetDemo { */ public OcrDirectionModel getDirectionModel(){ DirectionModelConfig directionModelConfig = new DirectionModelConfig(); - //指定检测模型 - directionModelConfig.setDetModelEnum(CommonDetModelEnum.PADDLEOCR_V5_DET_MODEL); - //指定检测模型位置，需要更改为自己的模型路径（下载地址请查看文档） - directionModelConfig.setDetModelPath("/Users/xxx/Documents/develop/ocr模型/PP-OCRv5_server_det_infer/PP-OCRv5_server_det.onnx"); - //指定文本方向检测模型 - directionModelConfig.setModelEnum(DirectionModelEnum.CH_PPOCR_MOBILE_V2_CLS); - //指定文本方向检测模型路径，需要更改为自己的模型路径（下载地址请查看文档） - directionModelConfig.setModelPath("/Users/xxx/Documents/develop/ocr模型/ch_ppocr_mobile_v2.0_cls.onnx"); + //指定行文本方向检测模型 + directionModelConfig.setModelEnum(DirectionModelEnum.PP_LCNET_X0_25); + //指定行文本方向检测模型路径，需要更改为自己的模型路径（下载地址请查看文档） + directionModelConfig.setModelPath("/Users/xxx/Documents/develop/model/ocr/PP-LCNet_x0_25_textline_ori_infer/PP-LCNet_x0_25_textline_ori_infer.onnx"); directionModelConfig.setDevice(device); + directionModelConfig.setTextDetModel(getDetectionModel()); return OcrModelFactory.getInstance().getDirectionModel(directionModelConfig); } + /** + * 获取文本检测模型 + * @return + */ + public OcrCommonDetModel getDetectionModel() { + OcrDetModelConfig config = new OcrDetModelConfig(); + //指定检测模型 + config.setModelEnum(CommonDetModelEnum.PP_OCR_V5_MOBILE_DET_MODEL); + //指定模型位置，需要更改为自己的模型路径（下载地址请查看文档） + config.setDetModelPath("/Users/xxx/Documents/develop/model/ocr/PP-OCRv5_mobile_det_infer/PP-OCRv5_mobile_det_infer.onnx"); + config.setDevice(device); + return OcrModelFactory.getInstance().getDetModel(config); + } + /** * 文本方向检测 @@ -57,8 +68,8 @@ public class OcrDirectionDetDemo { @Test public void detect(){ try (OcrDirectionModel directionModel = getDirectionModel()){ - List itemList = directionModel.detect("src/main/resources/ocr_3.jpg"); - log.info("OCR方向检测结果：{}", JSONObject.toJSONString(itemList)); + List itemList = directionModel.detect("src/main/resources/ocr_1.jpg"); + log.info("OCR方向检测结果1：{}", JSONObject.toJSONString(itemList)); } 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 new file mode 100644 index 0000000..037a278 --- /dev/null +++ b/examples/ocr-examples/src/main/java/smartai/examples/ocr/common/OcrRecognizeDemo.java @@ -0,0 +1,206 @@ +package smartai.examples.ocr.common; + +import ai.djl.modality.cv.Image; +import cn.smartjavaai.common.enums.DeviceEnum; +import cn.smartjavaai.common.utils.ImageUtils; +import cn.smartjavaai.ocr.config.DirectionModelConfig; +import cn.smartjavaai.ocr.config.OcrDetModelConfig; +import cn.smartjavaai.ocr.config.OcrRecModelConfig; +import cn.smartjavaai.ocr.config.OcrRecOptions; +import cn.smartjavaai.ocr.entity.OcrBox; +import cn.smartjavaai.ocr.entity.OcrInfo; +import cn.smartjavaai.ocr.enums.CommonDetModelEnum; +import cn.smartjavaai.ocr.enums.CommonRecModelEnum; +import cn.smartjavaai.ocr.enums.DirectionModelEnum; +import cn.smartjavaai.ocr.factory.OcrModelFactory; +import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel; +import cn.smartjavaai.ocr.model.common.direction.OcrDirectionModel; +import cn.smartjavaai.ocr.model.common.recognize.OcrCommonRecModel; +import com.alibaba.fastjson.JSONObject; +import lombok.extern.slf4j.Slf4j; +import org.junit.Test; + +import java.io.File; +import java.util.List; + +/** + * OCR 文本识别 示例 + * 模型下载地址：https://pan.baidu.com/s/1MLfd73Vjdpnuls9-oqc9uw?pwd=1234 提取码: 1234 + * @author dwj + * @date 2025/5/25 + */ +@Slf4j +public class OcrRecognizeDemo { + + //设备类型 + public static DeviceEnum device = DeviceEnum.CPU; + + /** + * 获取通用识别模型（不带方向矫正） + * @return + */ + public OcrCommonRecModel getRecModel(){ + OcrRecModelConfig recModelConfig = new OcrRecModelConfig(); + //指定文本识别模型 + recModelConfig.setRecModelEnum(CommonRecModelEnum.PP_OCR_V5_MOBILE_REC_MODEL); + //指定识别模型位置，需要更改为自己的模型路径（下载地址请查看文档） + recModelConfig.setRecModelPath("/Users/xxx/Documents/develop/model/ocr/PP-OCRv5_mobile_rec_infer/PP-OCRv5_mobile_rec_infer.onnx"); + recModelConfig.setDevice(device); + recModelConfig.setTextDetModel(getDetectionModel()); + return OcrModelFactory.getInstance().getRecModel(recModelConfig); + } + + /** + * 获取文本检测模型 + * @return + */ + public OcrCommonDetModel getDetectionModel() { + OcrDetModelConfig config = new OcrDetModelConfig(); + //指定检测模型 + config.setModelEnum(CommonDetModelEnum.PP_OCR_V5_MOBILE_DET_MODEL); + //指定模型位置，需要更改为自己的模型路径（下载地址请查看文档） + config.setDetModelPath("/Users/xxx/Documents/develop/model/ocr/PP-OCRv5_mobile_det_infer/PP-OCRv5_mobile_det_infer.onnx"); + config.setDevice(device); + return OcrModelFactory.getInstance().getDetModel(config); + } + + /** + * 获取方向检测模型 + * @return + */ + public OcrDirectionModel getDirectionModel(){ + DirectionModelConfig directionModelConfig = new DirectionModelConfig(); + //指定行文本方向检测模型 + directionModelConfig.setModelEnum(DirectionModelEnum.PP_LCNET_X0_25); + //指定行文本方向检测模型路径，需要更改为自己的模型路径（下载地址请查看文档） + directionModelConfig.setModelPath("/Users/xxx/Documents/develop/model/ocr/PP-LCNet_x0_25_textline_ori_infer/PP-LCNet_x0_25_textline_ori_infer.onnx"); + directionModelConfig.setDevice(device); + return OcrModelFactory.getInstance().getDirectionModel(directionModelConfig); + } + + + /** + * 获取通用识别模型(带方向矫正) + * @return + */ + public OcrCommonRecModel getRecModelWithDirection() { + OcrRecModelConfig recModelConfig = new OcrRecModelConfig(); + //指定文本识别模型 + recModelConfig.setRecModelEnum(CommonRecModelEnum.PP_OCR_V5_MOBILE_REC_MODEL); + //指定识别模型位置，需要更改为自己的模型路径（下载地址请查看文档） + recModelConfig.setRecModelPath("/Users/xxx/Documents/develop/model/ocr/PP-OCRv5_mobile_rec_infer/PP-OCRv5_mobile_rec_infer.onnx"); + recModelConfig.setDevice(device); + recModelConfig.setTextDetModel(getDetectionModel()); + recModelConfig.setDirectionModel(getDirectionModel()); + return OcrModelFactory.getInstance().getRecModel(recModelConfig); + } + + + /** + * 文本识别 + * 支持简体中文、繁体中文、英文、日文四种主要语言，以及手写、竖版、拼音、生僻字 + * 流程：文本检测 -> 文本识别 + * 注意事项： + * 1、批量检测时，模型应统一放在外层 try 中使用，避免重复加载，自动释放资源更安全。 + * 2、模型文件需要放在单独文件夹 + */ + @Test + public void recognize(){ + try (OcrCommonRecModel recModel = getRecModel()){ + //不带方向矫正，分行返回文本 + OcrRecOptions options = new OcrRecOptions(false, true); + OcrInfo ocrInfo = recModel.recognize("src/main/resources/ocr_2.jpg",options); + log.info("OCR识别结果：{}", JSONObject.toJSONString(ocrInfo)); + } catch (Exception e) { + e.printStackTrace(); + } + } + + + /** + * 文本识别（手写字） + * 支持简体中文、繁体中文、英文、日文四种主要语言，以及手写、竖版、拼音、生僻字 + * 流程：文本检测 -> 文本识别 + * 注意事项： + * 1、批量检测时，模型应统一放在外层 try 中使用，避免重复加载，自动释放资源更安全。 + * 2、模型文件需要放在单独文件夹 + */ + @Test + public void recognizeHandWriting(){ + 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) { + e.printStackTrace(); + } + } + + /** + * 文本识别（带方向矫正） + * 支持简体中文、繁体中文、英文、日文四种主要语言，以及手写、竖版、拼音、生僻字 + * 本方法支持多角度文字识别 + * 流程：文本检测 -> 方向检测 -> 方向矫正 -> 文本识别 + * 注意事项： + * 1、批量检测时，模型应统一放在外层 try 中使用，避免重复加载，自动释放资源更安全。 + * 2、模型文件需要放在单独文件夹 + */ + @Test + public void recognize2(){ + try (OcrCommonRecModel recModel = getRecModelWithDirection()){ + //带方向矫正，分行返回文本 + OcrRecOptions options = new OcrRecOptions(true, true); + OcrInfo ocrInfo = recModel.recognize("src/main/resources/ocr_3.jpg",options); + log.info("OCR识别结果：{}", JSONObject.toJSONString(ocrInfo)); + } catch (Exception e) { + e.printStackTrace(); + } + } + + + + /** + * 文本识别并绘制结果 + * 支持简体中文、繁体中文、英文、日文四种主要语言，以及手写、竖版、拼音、生僻字 + * 流程：文本检测 -> 文本识别 + * 注意事项： + * 1、批量检测时，模型应统一放在外层 try 中使用，避免重复加载，自动释放资源更安全。 + * 2、模型文件需要放在单独文件夹 + */ + @Test + public void recognizeAndDraw(){ + 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) { + e.printStackTrace(); + } + } + + /** + * 批量识别 + * 注意事项： + * 1、批量检测时，模型应统一放在外层 try 中使用，避免重复加载，自动释放资源更安全。 + * 2、模型文件需要放在单独文件夹 + */ + @Test + public void batchRecognize(){ + try (OcrCommonRecModel recModel = getRecModelWithDirection()){ + //批量检测要求图片宽高一致 + String folderPath = "/Users/xxx/Downloads/testing33"; + //读取文件夹中所有图片 + List images = ImageUtils.readImagesFromFolder(folderPath); + //带方向矫正，分行返回文本 + OcrRecOptions options = new OcrRecOptions(true, true); + List ocrResult = recModel.batchRecognizeDJLImage(images, options); + for(int i = 0; i < ocrResult.size(); i++){ + log.info("图片" + i + "文本识别结果：{}", JSONObject.toJSONString(ocrResult.get(i))); + } + } catch (Exception e) { + e.printStackTrace(); + } + } + + + + +} 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 new file mode 100644 index 0000000..8b41f9b --- /dev/null +++ b/examples/ocr-examples/src/main/java/smartai/examples/ocr/table/TableRecDemo.java @@ -0,0 +1,144 @@ +package smartai.examples.ocr.table; + +import ai.djl.modality.cv.Image; +import cn.hutool.core.io.FileUtil; +import cn.smartjavaai.common.entity.R; +import cn.smartjavaai.common.enums.DeviceEnum; +import cn.smartjavaai.common.utils.ImageUtils; +import cn.smartjavaai.ocr.config.*; +import cn.smartjavaai.ocr.entity.OcrInfo; +import cn.smartjavaai.ocr.entity.TableStructureResult; +import cn.smartjavaai.ocr.enums.CommonDetModelEnum; +import cn.smartjavaai.ocr.enums.CommonRecModelEnum; +import cn.smartjavaai.ocr.enums.DirectionModelEnum; +import cn.smartjavaai.ocr.enums.TableStructureModelEnum; +import cn.smartjavaai.ocr.factory.OcrModelFactory; +import cn.smartjavaai.ocr.factory.TableRecModelFactory; +import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel; +import cn.smartjavaai.ocr.model.common.direction.OcrDirectionModel; +import cn.smartjavaai.ocr.model.common.recognize.OcrCommonRecModel; +import cn.smartjavaai.ocr.model.table.TableRecognizer; +import cn.smartjavaai.ocr.model.table.TableStructureModel; +import com.alibaba.fastjson.JSONObject; +import lombok.extern.slf4j.Slf4j; +import org.junit.Test; + +import javax.imageio.ImageIO; +import java.awt.image.BufferedImage; +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.List; + +/** + * OCR 文本识别 示例 + * 模型下载地址：https://pan.baidu.com/s/1MLfd73Vjdpnuls9-oqc9uw?pwd=1234 提取码: 1234 + * @author dwj + * @date 2025/5/25 + */ +@Slf4j +public class TableRecDemo { + + //设备类型 + public static DeviceEnum device = DeviceEnum.CPU; + + /** + * 获取通用识别模型（不带方向矫正） + * @return + */ + public OcrCommonRecModel getRecModel(){ + OcrRecModelConfig recModelConfig = new OcrRecModelConfig(); + //指定文本识别模型 + recModelConfig.setRecModelEnum(CommonRecModelEnum.PP_OCR_V5_MOBILE_REC_MODEL); + //指定识别模型位置，需要更改为自己的模型路径（下载地址请查看文档） + recModelConfig.setRecModelPath("/Users/xxx/Documents/develop/model/ocr/PP-OCRv5_mobile_rec_infer/PP-OCRv5_mobile_rec_infer.onnx"); + recModelConfig.setDevice(device); + recModelConfig.setTextDetModel(getDetectionModel()); + return OcrModelFactory.getInstance().getRecModel(recModelConfig); + } + + /** + * 获取文本检测模型 + * @return + */ + public OcrCommonDetModel getDetectionModel() { + OcrDetModelConfig config = new OcrDetModelConfig(); + //指定检测模型 + config.setModelEnum(CommonDetModelEnum.PP_OCR_V5_MOBILE_DET_MODEL); + //指定模型位置，需要更改为自己的模型路径（下载地址请查看文档） + config.setDetModelPath("/Users/xxx/Documents/develop/model/ocr/PP-OCRv5_mobile_det_infer/PP-OCRv5_mobile_det_infer.onnx"); +// config.setDetModelPath("/Users/xxx/Documents/develop/model/ocr/PP-OCRv5_server_det_infer/PP-OCRv5_server_det.onnx"); + config.setDevice(device); + return OcrModelFactory.getInstance().getDetModel(config); + } + + /** + * 获取方向检测模型 + * @return + */ + public OcrDirectionModel getDirectionModel(){ + DirectionModelConfig directionModelConfig = new DirectionModelConfig(); + //指定行文本方向检测模型 + directionModelConfig.setModelEnum(DirectionModelEnum.PP_LCNET_X0_25); + //指定行文本方向检测模型路径，需要更改为自己的模型路径（下载地址请查看文档） + directionModelConfig.setModelPath("/Users/xxx/Documents/develop/model/ocr/PP-LCNet_x0_25_textline_ori_infer/PP-LCNet_x0_25_textline_ori_infer.onnx"); + directionModelConfig.setDevice(device); + return OcrModelFactory.getInstance().getDirectionModel(directionModelConfig); + } + + /** + * 创建表格结构识别模型 + * @return + */ + public TableStructureModel getTableStructureModel(){ + TableStructureConfig config = new TableStructureConfig(); + //指定行文本方向检测模型 + config.setModelEnum(TableStructureModelEnum.SLANET_PLUS); + //指定行文本方向检测模型路径，需要更改为自己的模型路径（下载地址请查看文档） + config.setModelPath("/Users/xxx/Documents/develop/model/ocr/slanet-plus/slanet-plus.onnx"); +// config.setModelPath("/Users/xxx/Documents/develop/model/ocr/SLANet_infer/SLANet.onnx"); + config.setDevice(device); + return TableRecModelFactory.getInstance().getTableStructureModel(config); + } + + + + /** + * 表格识别 + * 仅支持简单表格 + * 流程：表格结构识别 -> 文本检测 -> 文本识别 -> 合成html table + * 注意事项： + * 1、批量检测时，模型应统一放在外层 try 中使用，避免重复加载，自动释放资源更安全。 + * 2、模型文件需要放在单独文件夹 + */ + @Test + public void recognize(){ + try (TableStructureModel tableStructureModel = getTableStructureModel(); + OcrCommonDetModel detModel = getDetectionModel(); + OcrCommonRecModel recModel = getRecModel(); + OcrDirectionModel directionModel = getDirectionModel()){ + //创建表格识别器 + TableRecognizer tableRecognizer = TableRecognizer.builder() + .withStructureModel(tableStructureModel) + .withTextDetModel(detModel) +// .withDirectionModel(getDirectionModel()) //如果表格中存在旋转的文字，可以使用方向分类模型 + .withTextRecModel(recModel).build(); + String imagePath = "src/main/resources/table/table_ch1.png"; + BufferedImage image = ImageIO.read(new File(Paths.get(imagePath).toAbsolutePath().toString())); + R result = tableRecognizer.recognize(image); + if(result.isSuccess()){ + log.info("result: {}", result.getData().getHtml()); + //导出html内容到文件 + Path outputPath = Paths.get("output/table_ch2_result.html"); + FileUtil.writeUtf8String(result.getData().getHtml(), outputPath.toAbsolutePath().toString()); + //绘制表格结构 + tableRecognizer.drawTable(result.getData(), image, "output/table_ch2_result.jpg"); + //导出excel，如果导出失败，可能是因为表格结果识别的结果是错乱的 + tableRecognizer.exportExcel(result.getData().getHtml(), "output/table_ch2_result.xls"); + } + } catch (Exception e) { + e.printStackTrace(); + } + } + +} diff --git a/examples/ocr-examples/src/main/resources/ocr_2.jpg b/examples/ocr-examples/src/main/resources/ocr_2.jpg index ed91b8c..ecf5212 100644 Binary files a/examples/ocr-examples/src/main/resources/ocr_2.jpg and b/examples/ocr-examples/src/main/resources/ocr_2.jpg differ diff --git a/examples/translation-example/pom.xml b/examples/translation-example/pom.xml index 6823205..ec76a11 100644 --- a/examples/translation-example/pom.xml +++ b/examples/translation-example/pom.xml @@ -12,7 +12,7 @@ 11 11 UTF-8 - 1.0.19 + 1.0.20 smartai.examples.nlp.translation.TranslationDemo diff --git a/pom.xml b/pom.xml index 6df95c2..a78b7ba 100644 --- a/pom.xml +++ b/pom.xml @@ -7,7 +7,7 @@ SmartJavaAI cn.smartjavaai smartjavaai-parent - 1.0.19 + 1.0.20 pom SmartJavaAI diff --git a/smartjavaai-all/pom.xml b/smartjavaai-all/pom.xml index 3bd7b2d..e66d60c 100644 --- a/smartjavaai-all/pom.xml +++ b/smartjavaai-all/pom.xml @@ -6,11 +6,11 @@ cn.smartjavaai smartjavaai-parent - 1.0.19 + 1.0.20 smartjavaai-all - 1.0.19 + 1.0.20 ${project.artifactId} SmartJavaAI https://github.com/geekwenjie/SmartJavaAI diff --git a/smartjavaai-bom/pom.xml b/smartjavaai-bom/pom.xml index 22c1a01..d841897 100644 --- a/smartjavaai-bom/pom.xml +++ b/smartjavaai-bom/pom.xml @@ -6,10 +6,10 @@ cn.smartjavaai smartjavaai-parent - 1.0.19 + 1.0.20 - 1.0.19 + 1.0.20 smartjavaai-bom smartjavaai-bom 统一版本管理的 BOM 包，同时支持 import 和全量依赖 diff --git a/smartjavaai-common/pom.xml b/smartjavaai-common/pom.xml index 8767b3d..1314870 100644 --- a/smartjavaai-common/pom.xml +++ b/smartjavaai-common/pom.xml @@ -6,7 +6,7 @@ cn.smartjavaai smartjavaai-parent - 1.0.19 + 1.0.20 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 new file mode 100644 index 0000000..fdbdfe7 --- /dev/null +++ b/smartjavaai-common/src/main/java/cn/smartjavaai/common/config/ModelConfig.java @@ -0,0 +1,56 @@ +package cn.smartjavaai.common.config; + +import cn.smartjavaai.common.enums.DeviceEnum; +import lombok.Data; + +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +/** + * 模型配置 + * @author dwj + */ +@Data +public class ModelConfig { + + /** + * 设备类型 + */ + private DeviceEnum device; + + /** + * gpu设备ID 当device为GPU时生效 + */ + private int gpuId = 0; + + /** + * 批量数据打包方式：stack，padding + */ + private String batchifier; + + /** + * 个性化配置（按模型类型动态解析） + */ + private ConcurrentHashMap customParams = new ConcurrentHashMap<>(); + + 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 ConcurrentHashMap<>(); + } + customParams.put(key, value); + } + + + + +} diff --git a/smartjavaai-common/src/main/java/cn/smartjavaai/common/entity/DetectionRectangle.java b/smartjavaai-common/src/main/java/cn/smartjavaai/common/entity/DetectionRectangle.java index a4bc898..06fb5c5 100644 --- a/smartjavaai-common/src/main/java/cn/smartjavaai/common/entity/DetectionRectangle.java +++ b/smartjavaai-common/src/main/java/cn/smartjavaai/common/entity/DetectionRectangle.java @@ -23,4 +23,5 @@ public class DetectionRectangle { this.width = width; this.height = height; } + } diff --git a/smartjavaai-common/src/main/java/cn/smartjavaai/common/entity/ocr/TableStructure.java b/smartjavaai-common/src/main/java/cn/smartjavaai/common/entity/ocr/TableStructure.java new file mode 100644 index 0000000..4b7c214 --- /dev/null +++ b/smartjavaai-common/src/main/java/cn/smartjavaai/common/entity/ocr/TableStructure.java @@ -0,0 +1,8 @@ +package cn.smartjavaai.common.entity.ocr; + +/** + * 表格结构 + * @author dwj + */ +public class TableStructure { +} 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 e30fafb..1a3c75c 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 @@ -370,6 +370,104 @@ public class ImageUtils { g.drawString(text, x + padding, y + ascent); } + /** + * 计算左上角，右下角坐标 （x0,y0,x1,y1） + * Get absolute coordinations + * + * @param rect + * @param width + * @param height + * @return + */ + public static int[] rectXYXY(ai.djl.modality.cv.output.Rectangle rect, int width, int height) { + int left = Math.max((int) (width * rect.getX()), 0); + int top = Math.max((int) (height * rect.getY()), 0); + int right = Math.min((int) (width * (rect.getX() + rect.getWidth())), width - 1); + int bottom = Math.min((int) (height * (rect.getY() + rect.getHeight())), height - 1); + return new int[] {left, top, right, bottom}; + } + + /** + * 列出文件夹下的所有图片文件 + * List all image files under the folder + * + * @param folderPath + * @return + */ + public static List listImageFiles(String folderPath) { + File folder = new File(folderPath); + List imageFiles = new ArrayList<>(); + if (folder.exists() && folder.isDirectory()) { + File[] files = folder.listFiles(); + if (files == null) { + return imageFiles; + } + for (File file : files) { + if (file.isFile()) { + String name = file.getName().toLowerCase(); + if (name.endsWith(".jpg") || name.endsWith(".jpeg") || + name.endsWith(".png") || name.endsWith(".bmp") || + name.endsWith(".gif") || name.endsWith(".tiff") || + name.endsWith(".webp")) { + imageFiles.add(file); + } + } + } + } + return imageFiles; + } + + /** + * 读取指定目录下所有图片，返回 List（DJL 格式） + * + * @param folderPath 图片文件夹路径 + * @return List + * @throws IOException + */ + public static List readImagesFromFolder(String folderPath) throws IOException { + File folder = new File(folderPath); + List imageList = new ArrayList<>(); + if (folder.exists() && folder.isDirectory()) { + File[] files = folder.listFiles(); + if (files == null) { + return imageList; + } + for (File file : files) { + if (file.isFile()) { + String name = file.getName().toLowerCase(); + if (name.endsWith(".jpg") || name.endsWith(".jpeg") || + name.endsWith(".png") || name.endsWith(".bmp") || + name.endsWith(".gif") || name.endsWith(".tiff") || + name.endsWith(".webp")) { + + Image img = ImageFactory.getInstance().fromInputStream(Files.newInputStream(file.toPath())); + imageList.add(img); + } + } + } + } + return imageList; + } + + /** + * 判断所有图片尺寸是否一致 + * + * @param images 图片列表 + */ + public static boolean isAllImageSizeEqual(List images) { + if (images == null || images.isEmpty()) { + return true; // 空集合视为一致 + } + int width = images.get(0).getWidth(); + int height = images.get(0).getHeight(); + for (Image img : images) { + if (img.getWidth() != width || img.getHeight() != height) { + return false; + } + } + return true; + } + } diff --git a/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/VideoUtils.java b/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/VideoUtils.java new file mode 100644 index 0000000..d43ddda --- /dev/null +++ b/smartjavaai-common/src/main/java/cn/smartjavaai/common/utils/VideoUtils.java @@ -0,0 +1,77 @@ +package cn.smartjavaai.common.utils; + +import org.bytedeco.ffmpeg.global.avcodec; +import org.bytedeco.javacv.FFmpegFrameGrabber; +import org.bytedeco.javacv.FFmpegFrameRecorder; +import org.bytedeco.javacv.Frame; +import org.bytedeco.javacv.OpenCVFrameConverter; +import org.bytedeco.opencv.global.opencv_core; +import org.bytedeco.opencv.opencv_core.Mat; + +/** + * 视频工具类 + * @author dwj + * @date 2025/7/17 + */ +public class VideoUtils { + + /** + * 视频旋转 + * @param inputPath 输入视频路径 + * @param outputPath 输出视频路径 + * @param angle 旋转角度 + * @param format 视频格式 + * @param videoCodec 视频编码器 + * @throws FFmpegFrameRecorder.Exception + * @throws FFmpegFrameGrabber.Exception + */ + public static void rotateVideo(String inputPath, String outputPath, int angle, String format, int videoCodec) throws FFmpegFrameRecorder.Exception, FFmpegFrameGrabber.Exception { + FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(inputPath); + grabber.start(); + int inputWidth = grabber.getImageWidth(); + int inputHeight = grabber.getImageHeight(); + int outputWidth = inputWidth; + int outputHeight = inputHeight; + + if (angle == 90 || angle == 270) { + outputWidth = inputHeight; + outputHeight = inputWidth; + } + FFmpegFrameRecorder recorder = new FFmpegFrameRecorder(outputPath, + outputWidth, outputHeight, grabber.getAudioChannels()); + recorder.setVideoCodec(videoCodec); + recorder.setFormat(format); + recorder.start(); + Frame frame; + OpenCVFrameConverter.ToMat converter = new OpenCVFrameConverter.ToMat(); + while ((frame = grabber.grab()) != null) { + if (frame.image != null) { + Mat mat = converter.convert(frame); + Mat rotated = new Mat(); + switch (angle) { + case 90: + opencv_core.transpose(mat, rotated); + opencv_core.flip(rotated, rotated, 1); + break; + case 180: + opencv_core.flip(mat, rotated, -1); + break; + case 270: + opencv_core.transpose(mat, rotated); + opencv_core.flip(rotated, rotated, 0); + break; + default: + rotated = mat.clone(); + break; + } + frame = converter.convert(rotated); + recorder.record(frame); + } + } + recorder.stop(); + recorder.release(); + grabber.stop(); + grabber.release(); + } + +} diff --git a/smartjavaai-face/pom.xml b/smartjavaai-face/pom.xml index e5980da..beecee2 100644 --- a/smartjavaai-face/pom.xml +++ b/smartjavaai-face/pom.xml @@ -6,11 +6,11 @@ cn.smartjavaai smartjavaai-parent - 1.0.19 + 1.0.20 smartjavaai-face - 1.0.19 + 1.0.20 smartjavaai-face SmartJavaAI https://github.com/geekwenjie/SmartJavaAI 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 48fbbb8..7803f38 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 @@ -23,31 +23,11 @@ public class FaceRecConfig { */ private FaceRecModelEnum modelEnum; - /** - * 置信度阈值 - */ - private double confidenceThreshold = FaceDetectConstant.DEFAULT_CONFIDENCE_THRESHOLD; - - /** - * 相似度阈值 作用：判断是否为同一人脸 - */ - //private double similarityThreshold = 0D; - - /** - * 非极大抑制阈值 作用：消除重叠检测框，保留最优结果 - */ - private double nmsThresh = FaceDetectConstant.NMS_THRESHOLD; - /** * 模型路径 */ private String modelPath; - /** - * 人脸库路径 - */ - private String faceDbPath; - /** * 设备类型 */ 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 6085a90..8ca3a02 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 @@ -48,6 +48,11 @@ public class LivenessConfig { */ private int frameCount = LivenessConstant.DEFAULT_FRAME_COUNT; + /** + * 视频检测最大帧数 + */ + private int maxVideoDetectFrames = LivenessConstant.DEFAULT_MAX_VIDEO_DETECT_FRAMES; + /** * 真人阈值 */ diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/constant/LivenessConstant.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/constant/LivenessConstant.java index 4f8969f..54737ed 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/constant/LivenessConstant.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/constant/LivenessConstant.java @@ -24,4 +24,9 @@ public class LivenessConstant { * 视频默认检测帧数 */ public static final int DEFAULT_FRAME_COUNT = 10; + + /** + * 视频默认最大检测帧数 + */ + public static final int DEFAULT_MAX_VIDEO_DETECT_FRAMES = Integer.MAX_VALUE; } diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/dao/FaceDao.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/dao/FaceDao.java index 5a0c035..972dd60 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/dao/FaceDao.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/dao/FaceDao.java @@ -1,6 +1,5 @@ package cn.smartjavaai.face.dao; -import cn.smartjavaai.face.entity.FaceData; import cn.smartjavaai.face.sqllite.RowMapper; import cn.smartjavaai.face.sqllite.SqliteHelper; import cn.smartjavaai.face.utils.VectorUtils; @@ -80,7 +79,7 @@ public class FaceDao { */ public FaceVector findById(String id) throws SQLException, ClassNotFoundException { SqliteHelper sqliteHelper = SqliteHelper.getInstance(dbFilePath); - String sql = "select \"id\",\"vector\",\"metadata\" from " + FACE_TABLE_NAME + " where \"id\"=" + id; + String sql = "select \"id\",\"vector\",\"metadata\" from " + FACE_TABLE_NAME + " where \"id\" = '" + id + "'"; List faceVectors = sqliteHelper.executeQuery(sql, new RowMapper() { @Override public FaceVector mapRow(ResultSet rs, int id) throws SQLException { @@ -164,8 +163,9 @@ public class FaceDao { * @throws ClassNotFoundException 类未找到异常 */ public List findFace(int pageNo, int pageSize) throws SQLException, ClassNotFoundException { + long offset = (pageNo - 1) * pageSize; String sql = "select \"id\",\"vector\",\"metadata\" from " + FACE_TABLE_NAME + - " limit " + pageNo * pageSize + "," + pageSize; + " limit " + offset + "," + pageSize; SqliteHelper sqliteHelper = SqliteHelper.getInstance(dbFilePath); return sqliteHelper.executeQuery(sql, new RowMapper() { @Override diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/entity/FaceData.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/entity/FaceData.java deleted file mode 100644 index 2f33f1f..0000000 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/entity/FaceData.java +++ /dev/null @@ -1,19 +0,0 @@ -package cn.smartjavaai.face.entity; - -import lombok.Data; - -/** - * 人脸数据 - * @author dwj - */ -@Data -public class FaceData { - - private String key; - private long index; - private byte[] imgData; - private int width = 256; - private int height = 256; - private int channel = 3; - -} diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/factory/ExpressionModelFactory.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/factory/ExpressionModelFactory.java index 39e89b8..d5327dd 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/factory/ExpressionModelFactory.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/factory/ExpressionModelFactory.java @@ -24,6 +24,9 @@ public class ExpressionModelFactory { // 使用 volatile 和双重检查锁定来确保线程安全的单例模式 private static volatile ExpressionModelFactory instance; + /** + * 模型缓存 + */ private static final ConcurrentHashMap modelMap = new ConcurrentHashMap<>(); /** 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 749825e..63cd24f 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,8 +9,6 @@ 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.hutool.core.lang.generator.UUIDGenerator; import cn.smartjavaai.common.entity.*; import cn.smartjavaai.common.entity.face.FaceInfo; import cn.smartjavaai.common.entity.face.FaceSearchResult; @@ -28,7 +26,7 @@ import cn.smartjavaai.face.enums.SimilarityType; import cn.smartjavaai.face.exception.FaceException; import cn.smartjavaai.face.factory.FaceDetModelFactory; import cn.smartjavaai.face.model.facedect.FaceDetModel; -import cn.smartjavaai.face.model.facerec.criterial.FaceRecCriteriaFactory; +import cn.smartjavaai.face.model.facerec.criteria.FaceRecCriteriaFactory; import cn.smartjavaai.face.preprocess.DJLImagePreprocessor; import cn.smartjavaai.face.utils.*; import cn.smartjavaai.face.vector.config.MilvusConfig; @@ -646,6 +644,22 @@ public class CommonFaceRecModel implements FaceRecModel{ } } + @Override + public R getFaceInfoById(String id) { + if(vectorDBClient == null){ + return R.fail(1000, "向量数据库未初始化成功"); + } + return R.ok(vectorDBClient.getFaceInfoById(id)); + } + + @Override + public R> listFaces(long pageNum, long pageSize) { + if(vectorDBClient == null){ + return R.fail(1000, "向量数据库未初始化成功"); + } + return R.ok(vectorDBClient.listFaces(pageNum, pageSize)); + } + @Override public void loadFaceFeatures() { if(Objects.isNull(vectorDBClient)){ 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 b353323..1a695bf 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 @@ -6,6 +6,7 @@ import cn.smartjavaai.face.config.FaceRecConfig; 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 java.awt.image.BufferedImage; import java.io.InputStream; @@ -245,8 +246,26 @@ public interface FaceRecModel extends AutoCloseable{ } + /** + * 使用人脸ID获取人脸信息 + * @param id + * @return + */ + default R getFaceInfoById(String id){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + /** + * 获取人脸列表 + * @param pageNum + * @param pageSize + * @return + */ + default R> listFaces(long pageNum, long pageSize){ + 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 1624b64..64b58d0 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 @@ -486,13 +486,13 @@ public class SeetaFace6FaceRecModel implements FaceRecModel{ * 检查是否存在人脸库 * @return */ - private boolean checkFaceDb(){ - if(Objects.nonNull(config) && StringUtils.isNotBlank(config.getFaceDbPath())){ - File file = new File(config.getFaceDbPath()); - return file.exists() && file.isFile(); - } - return false; - } +// private boolean checkFaceDb(){ +// if(Objects.nonNull(config) && StringUtils.isNotBlank(config.getFaceDbPath())){ +// File file = new File(config.getFaceDbPath()); +// return file.exists() && file.isFile(); +// } +// return false; +// } private FaceResult searchFaceDb(long index,float similar) { if(index >= 0){ @@ -837,6 +837,22 @@ public class SeetaFace6FaceRecModel implements FaceRecModel{ FaceRecModel.super.upsertFace(faceRegisterInfo, imageData); } + @Override + public R getFaceInfoById(String id) { + if(vectorDBClient == null){ + return R.fail(1000, "向量数据库未初始化成功"); + } + return R.ok(vectorDBClient.getFaceInfoById(id)); + } + + @Override + public R> listFaces(long pageNum, long pageSize) { + if(vectorDBClient == null){ + return R.fail(1000, "向量数据库未初始化成功"); + } + return R.ok(vectorDBClient.listFaces(pageNum, pageSize)); + } + @Override public void close() throws Exception { if(Objects.nonNull(faceDetectorPool)){ diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/criterial/FaceRecCriteriaFactory.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/criteria/FaceRecCriteriaFactory.java similarity index 86% rename from smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/criterial/FaceRecCriteriaFactory.java rename to smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/criteria/FaceRecCriteriaFactory.java index be2e88a..47b27a9 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/criterial/FaceRecCriteriaFactory.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/model/facerec/criteria/FaceRecCriteriaFactory.java @@ -1,28 +1,15 @@ -package cn.smartjavaai.face.model.facerec.criterial; +package cn.smartjavaai.face.model.facerec.criteria; import ai.djl.Device; import ai.djl.modality.cv.Image; -import ai.djl.modality.cv.output.DetectedObjects; -import ai.djl.modality.cv.transform.Normalize; -import ai.djl.modality.cv.transform.Resize; -import ai.djl.modality.cv.transform.ToTensor; -import ai.djl.modality.cv.translator.ImageFeatureExtractor; -import ai.djl.modality.cv.translator.ImageFeatureExtractorFactory; import ai.djl.repository.zoo.Criteria; import ai.djl.training.util.ProgressBar; -import ai.djl.translate.Translator; import cn.smartjavaai.common.enums.DeviceEnum; -import cn.smartjavaai.face.config.FaceDetConfig; import cn.smartjavaai.face.config.FaceRecConfig; -import cn.smartjavaai.face.constant.FaceDetectConstant; import cn.smartjavaai.face.constant.FaceNetConstant; -import cn.smartjavaai.face.constant.RetinaFaceConstant; -import cn.smartjavaai.face.constant.UltraLightFastGenericFaceConstant; -import cn.smartjavaai.face.enums.FaceDetModelEnum; import cn.smartjavaai.face.enums.FaceRecModelEnum; import cn.smartjavaai.face.model.facerec.translator.FaceFeatureTranslator; import cn.smartjavaai.face.model.facerec.translator.FaceNetRecTranslator; -import cn.smartjavaai.face.translator.FaceDetectionTranslator; import org.apache.commons.lang3.StringUtils; import java.nio.file.Paths; 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 d72ed8b..491fbac 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 @@ -354,6 +354,9 @@ public class CommonLivenessModel implements LivenessDetModel{ } // 逐帧处理视频 for (int frameIndex = 0; frameIndex < totalFrames; frameIndex++) { + if(frameIndex >= config.getMaxVideoDetectFrames()){ + return R.fail(10002, "超出最大检测帧数：" + config.getMaxVideoDetectFrames()); + } // 获取当前帧 Frame frame = grabber.grabImage(); if (frame != null) { 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 130426d..d286af2 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 @@ -515,6 +515,9 @@ public class Seetaface6LivenessModel implements LivenessDetModel{ } // 逐帧处理视频 for (int frameIndex = 0; frameIndex < totalFrames; frameIndex++) { + if(frameIndex >= config.getMaxVideoDetectFrames()){ + return R.fail(10002, "超出最大检测帧数：" + config.getMaxVideoDetectFrames()); + } // 获取当前帧 Frame frame = grabber.grabImage(); if (frame != null) { 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 f71eb9e..2f7a466 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 @@ -14,8 +14,11 @@ import io.milvus.param.*; import io.milvus.param.collection.*; import io.milvus.param.dml.*; import io.milvus.param.index.CreateIndexParam; +import io.milvus.response.DescCollResponseWrapper; import io.milvus.response.QueryResultsWrapper; import io.milvus.response.SearchResultsWrapper; +import io.milvus.v2.service.collection.request.DescribeCollectionReq; +import io.milvus.v2.service.collection.response.DescribeCollectionResp; import lombok.extern.slf4j.Slf4j; import org.apache.commons.lang3.StringUtils; @@ -53,6 +56,14 @@ public class MilvusClient implements VectorDBClient { serviceClient = new MilvusServiceClient(connectParam); collectionName = StringUtils.isNotBlank(config.getCollectionName()) ? config.getCollectionName() : VectorDBConstants.Defaults.DEFAULT_COLLECTION_NAME; createCollection(collectionName, config.getDimension()); + + boolean isAutoID = isAutoID(collectionName); + if(isAutoID && config.getIdStrategy() != IdStrategy.AUTO){ + throw new VectorDBException("ID策略与当前Collection不匹配"); + } + if(!isAutoID && config.getIdStrategy() == IdStrategy.AUTO){ + throw new VectorDBException("ID策略与当前Collection不匹配"); + } if(config.isUseMemoryCache()){ // 加载集合到内存 loadFaceFeatures(); @@ -152,6 +163,21 @@ public class MilvusClient implements VectorDBClient { } } + /** + * 判断是否为自增长ID + * @param collectionName + * @return + */ + private boolean isAutoID(String collectionName) { + R response = serviceClient.describeCollection( + DescribeCollectionParam.newBuilder() + .withCollectionName(collectionName) + .build() + ); + DescCollResponseWrapper wrapper = new DescCollResponseWrapper(response.getData()); + return wrapper.getPrimaryField().isAutoID(); + } + @Override public void dropCollection(String collectionName) { try { @@ -204,6 +230,7 @@ public class MilvusClient implements VectorDBClient { if(faceVector.getVector() == null || faceVector.getVector().length == 0){ throw new VectorDBException("插入数据失败：vector不能为空"); } + //自定义ID if(config.getIdStrategy() == IdStrategy.CUSTOM){ if(StringUtils.isBlank(faceVector.getId())){ @@ -525,40 +552,22 @@ public class MilvusClient implements VectorDBClient { } } - public void releaseCollection(String collectionName) { - if (!isInit){ - throw new VectorDBException("Milvus未初始化完毕"); - } - R response = serviceClient.releaseCollection(ReleaseCollectionParam.newBuilder() - .withCollectionName(collectionName) - .build()); - - if (response.getStatus() != R.Status.Success.getCode()) { - throw new VectorDBException("Milvus releaseCollection失败，msg: " + response.getMessage()); - } - } - - @Override - public FaceSearchResult getById(String id) { + public FaceVector getFaceInfoById(String id) { try { - if (!isInit){ + if (!isInit) { throw new VectorDBException("Milvus未初始化完毕"); } - // 构造搜索参数 - SearchParam searchParam = SearchParam.newBuilder() - .withCollectionName(collectionName) - .withExpr(VectorDBConstants.FieldNames.ID_FIELD + " == " + id) - .withOutFields(Arrays.asList(VectorDBConstants.FieldNames.ID_FIELD, VectorDBConstants.FieldNames.METADATA_FIELD)) - .build(); - - + String expr = VectorDBConstants.FieldNames.ID_FIELD + " == '" + id + "'"; + if(config.getIdStrategy() == IdStrategy.AUTO){ + expr = VectorDBConstants.FieldNames.ID_FIELD + " == " + id; + } // 5. 执行查询 R response = serviceClient.query( QueryParam.newBuilder() .withCollectionName(collectionName) - .withExpr(VectorDBConstants.FieldNames.ID_FIELD + " == " + id) - .withOutFields(Arrays.asList(VectorDBConstants.FieldNames.ID_FIELD, VectorDBConstants.FieldNames.METADATA_FIELD)) + .withExpr(expr) + .withOutFields(Arrays.asList(VectorDBConstants.FieldNames.ID_FIELD, VectorDBConstants.FieldNames.VECTOR_FIELD, VectorDBConstants.FieldNames.METADATA_FIELD)) .build() ); @@ -574,12 +583,91 @@ public class MilvusClient implements VectorDBClient { } // 提取第一条记录 QueryResultsWrapper.RowRecord row = records.get(0); - return new FaceSearchResult(id, 1,(String)row.get(VectorDBConstants.FieldNames.METADATA_FIELD)); + Object vectorObj = row.get(VectorDBConstants.FieldNames.VECTOR_FIELD); + float[] vector = null; + if (vectorObj instanceof List) { + // Milvus SDK通常返回List，转成float[] + List vectorList = (List) vectorObj; + vector = new float[vectorList.size()]; + for (int i = 0; i < vectorList.size(); i++) { + vector[i] = vectorList.get(i); + } + } + return new FaceVector(id, vector, (String) row.get(VectorDBConstants.FieldNames.METADATA_FIELD)); } catch (Exception e) { - throw new VectorDBException("搜索 Milvus 向量失败", e); + throw new RuntimeException(e); } } + @Override + public List listFaces(long pageNum, long pageSize) { + try { + if (!isInit) { + throw new VectorDBException("Milvus未初始化完毕"); + } + if (pageNum < 1 || pageSize < 1) { + throw new IllegalArgumentException("pageNum和pageSize必须大于0"); + } + long offset = (pageNum - 1) * pageSize; + // 构造查询参数，使用offset和limit实现分页 + QueryParam queryParam = QueryParam.newBuilder() + .withCollectionName(collectionName) + .withOutFields(Arrays.asList( + VectorDBConstants.FieldNames.ID_FIELD, + VectorDBConstants.FieldNames.VECTOR_FIELD, + VectorDBConstants.FieldNames.METADATA_FIELD)) + .withOffset(offset) + .withLimit(pageSize) + .build(); + + R response = serviceClient.query(queryParam); + + if (response.getStatus() != R.Status.Success.getCode()) { + throw new VectorDBException("分页查询失败: " + response.getMessage()); + } + + QueryResultsWrapper wrapper = new QueryResultsWrapper(response.getData()); + List records = wrapper.getRowRecords(); + if (records.isEmpty()) { + return Collections.emptyList(); + } + + List result = new ArrayList<>(); + for (QueryResultsWrapper.RowRecord row : records) { + String id = (String) row.get(VectorDBConstants.FieldNames.ID_FIELD); + Object vectorObj = row.get(VectorDBConstants.FieldNames.VECTOR_FIELD); + float[] vector = null; + if (vectorObj instanceof List) { + List vectorList = (List) vectorObj; + vector = new float[vectorList.size()]; + for (int i = 0; i < vectorList.size(); i++) { + vector[i] = vectorList.get(i); + } + } + String metadata = (String) row.get(VectorDBConstants.FieldNames.METADATA_FIELD); + result.add(new FaceVector(id, vector, metadata)); + } + + return result; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + public void releaseCollection(String collectionName) { + if (!isInit){ + throw new VectorDBException("Milvus未初始化完毕"); + } + R response = serviceClient.releaseCollection(ReleaseCollectionParam.newBuilder() + .withCollectionName(collectionName) + .build()); + + if (response.getStatus() != R.Status.Success.getCode()) { + throw new VectorDBException("Milvus releaseCollection失败，msg: " + response.getMessage()); + } + } + + @Override public void loadFaceFeatures() { // 加载集合到内存 diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/core/SQLiteClient.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/core/SQLiteClient.java index 098520e..11e2437 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/core/SQLiteClient.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/core/SQLiteClient.java @@ -196,6 +196,42 @@ public class SQLiteClient implements VectorDBClient { } } + @Override + public FaceVector getFaceInfoById(String id) { + if (!isInit) { + throw new VectorDBException("人脸库未加载完毕"); + } + // 先从内存缓存中获取 + FaceVector faceVector = memoryIndex.get(id); + if (faceVector == null) { + // 如果内存中没有，则从数据库查询 + try { + faceVector = faceDao.findById(id); + } catch (SQLException | ClassNotFoundException e) { + throw new VectorDBException("SQLite查询异常", e); + } + } + return faceVector; + } + + @Override + public List listFaces(long pageNum, long pageSize) { + if (!isInit) { + throw new VectorDBException("人脸库未加载完毕"); + } + + if (pageNum < 1 || pageSize < 1) { + throw new IllegalArgumentException("pageNum和pageSize必须大于0"); + } + + // 从数据库中查询指定分页的数据 + try { + return faceDao.findFace((int)pageNum, (int)pageSize); + } catch (Exception e) { + throw new VectorDBException("分页查询失败", e); + } + } + // ============= 私有辅助方法 ============= private void loadAllFeaturesToMemory() { @@ -234,19 +270,6 @@ public class SQLiteClient implements VectorDBClient { } } - @Override - public FaceSearchResult getById(String id) { - try { - FaceVector faceVector = faceDao.findById(id); - if(faceVector != null){ - return new FaceSearchResult(faceVector.getId(), 1.0f, faceVector.getMetadata()); - } - return null; - } catch (SQLException | RuntimeException | ClassNotFoundException e ) { - throw new VectorDBException("SQLite查询异常", e); - } - } - @Override public void loadFaceFeatures() { // 加载所有特征到内存 diff --git a/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/core/VectorDBClient.java b/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/core/VectorDBClient.java index 91ab325..4d3a3f7 100644 --- a/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/core/VectorDBClient.java +++ b/smartjavaai-face/src/main/java/cn/smartjavaai/face/vector/core/VectorDBClient.java @@ -1,5 +1,6 @@ package cn.smartjavaai.face.vector.core; +import cn.smartjavaai.common.entity.R; import cn.smartjavaai.face.entity.FaceSearchParams; import cn.smartjavaai.face.vector.entity.FaceVector; import cn.smartjavaai.common.entity.face.FaceSearchResult; @@ -94,12 +95,21 @@ public interface VectorDBClient extends AutoCloseable { @Override void close(); - /** - * 获取指定ID的向量 + * 使用人脸ID获取人脸信息 + * @param id * @return */ - FaceSearchResult getById(String id); + FaceVector getFaceInfoById(String id); + + + /** + * 获取人脸列表 + * @param pageNum + * @param pageSize + * @return + */ + List listFaces(long pageNum, long pageSize); /** * 加载人脸特征到内存 diff --git a/smartjavaai-objectdetection/pom.xml b/smartjavaai-objectdetection/pom.xml index a3c550a..8b10643 100644 --- a/smartjavaai-objectdetection/pom.xml +++ b/smartjavaai-objectdetection/pom.xml @@ -6,11 +6,11 @@ cn.smartjavaai smartjavaai-parent - 1.0.19 + 1.0.20 smartjavaai-objectdetection - 1.0.19 + 1.0.20 smartjavaai-objectdetection SmartJavaAI https://github.com/geekwenjie/SmartJavaAI diff --git a/smartjavaai-ocr/pom.xml b/smartjavaai-ocr/pom.xml index 24fca8e..70315cd 100644 --- a/smartjavaai-ocr/pom.xml +++ b/smartjavaai-ocr/pom.xml @@ -6,7 +6,7 @@ cn.smartjavaai smartjavaai-parent - 1.0.19 + 1.0.20 smartjavaai-ocr @@ -18,9 +18,31 @@ smartjavaai-common ${project.version} + + + + + + + + + + + + + + org.apache.poi + poi + 4.0.0 + + + dom4j + dom4j + 1.6.1 + - 1.0.19 + 1.0.20 smartjavaai-ocr SmartJavaAI https://github.com/geekwenjie/SmartJavaAI diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/DirectionModelConfig.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/DirectionModelConfig.java index 4b90ba1..a9457d2 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/DirectionModelConfig.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/DirectionModelConfig.java @@ -1,8 +1,10 @@ package cn.smartjavaai.ocr.config; +import cn.smartjavaai.common.config.ModelConfig; import cn.smartjavaai.common.enums.DeviceEnum; import cn.smartjavaai.ocr.enums.CommonDetModelEnum; import cn.smartjavaai.ocr.enums.DirectionModelEnum; +import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel; import cn.smartjavaai.ocr.model.common.recognize.OcrCommonRecModel; import lombok.Data; @@ -12,33 +14,22 @@ import lombok.Data; * @date 2025/4/22 */ @Data -public class DirectionModelConfig { +public class DirectionModelConfig extends ModelConfig { /** * 模型 */ private DirectionModelEnum modelEnum; - /** - * 设备类型 - */ - private DeviceEnum device; - /** * 检测模型路径 */ private String modelPath; /** - * 检测模型 + * 文本检测模型 */ - private CommonDetModelEnum detModelEnum; - - /** - * 检测模型路径 - */ - private String detModelPath; - + private OcrCommonDetModel textDetModel; diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/OcrDetModelConfig.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/OcrDetModelConfig.java index f03bbf4..db19928 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/OcrDetModelConfig.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/OcrDetModelConfig.java @@ -1,5 +1,6 @@ package cn.smartjavaai.ocr.config; +import cn.smartjavaai.common.config.ModelConfig; import cn.smartjavaai.common.enums.DeviceEnum; import cn.smartjavaai.ocr.enums.CommonDetModelEnum; import lombok.Data; @@ -10,18 +11,13 @@ import lombok.Data; * @date 2025/4/22 */ @Data -public class OcrDetModelConfig { +public class OcrDetModelConfig extends ModelConfig { /** * 模型 */ private CommonDetModelEnum modelEnum; - /** - * 设备类型 - */ - private DeviceEnum device; - /** * 检测模型路径 */ diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/OcrRecModelConfig.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/OcrRecModelConfig.java index 0aa9b35..1850c46 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/OcrRecModelConfig.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/OcrRecModelConfig.java @@ -1,9 +1,12 @@ package cn.smartjavaai.ocr.config; +import cn.smartjavaai.common.config.ModelConfig; import cn.smartjavaai.common.enums.DeviceEnum; import cn.smartjavaai.ocr.enums.CommonDetModelEnum; import cn.smartjavaai.ocr.enums.CommonRecModelEnum; import cn.smartjavaai.ocr.enums.DirectionModelEnum; +import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel; +import cn.smartjavaai.ocr.model.common.direction.OcrDirectionModel; import lombok.Data; /** @@ -12,41 +15,26 @@ import lombok.Data; * @date 2025/4/22 */ @Data -public class OcrRecModelConfig { - - /** - * 检测模型 - */ - private CommonDetModelEnum detModelEnum; +public class OcrRecModelConfig extends ModelConfig { /** * 识别模型 */ private CommonRecModelEnum recModelEnum; - /** - * 设备类型 - */ - private DeviceEnum device; - - /** - * 检测模型路径 - */ - private String detModelPath; - /** * 识别模型路径 */ private String recModelPath; /** - * 方向检测模型 + * 文本检测模型 */ - private DirectionModelEnum directionModelEnum; + private OcrCommonDetModel textDetModel; /** - * 方向检测模型路径 + * 文本方向模型 */ - private String directionModelPath; + private OcrDirectionModel directionModel; } diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/OcrRecOptions.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/OcrRecOptions.java new file mode 100644 index 0000000..86e1b86 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/OcrRecOptions.java @@ -0,0 +1,30 @@ +package cn.smartjavaai.ocr.config; + +import lombok.Data; + +/** + * OCR 识别配置 + * @author dwj + */ +@Data +public class OcrRecOptions { + + /** + * 是否进行文本方向矫正 + */ + private boolean enableDirectionCorrect = false; + + /** + * 是否进行结果分行 + */ + private boolean enableLineSplit = true; + + + public OcrRecOptions(boolean enableDirectionCorrect, boolean enableLineSplit) { + this.enableDirectionCorrect = enableDirectionCorrect; + this.enableLineSplit = enableLineSplit; + } + + public OcrRecOptions() { + } +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/TableStructureConfig.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/TableStructureConfig.java new file mode 100644 index 0000000..6ca3e48 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/config/TableStructureConfig.java @@ -0,0 +1,27 @@ +package cn.smartjavaai.ocr.config; + +import cn.smartjavaai.common.config.ModelConfig; +import cn.smartjavaai.common.enums.DeviceEnum; +import cn.smartjavaai.ocr.enums.CommonDetModelEnum; +import cn.smartjavaai.ocr.enums.TableStructureModelEnum; +import lombok.Data; + +/** + * OCR表格结构识别模型配置 + * @author dwj + */ +@Data +public class TableStructureConfig extends ModelConfig { + + /** + * 模型 + */ + private TableStructureModelEnum modelEnum; + + /** + * 检测模型路径 + */ + private String modelPath; + + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrBox.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrBox.java index fc2b6df..53fa899 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrBox.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrBox.java @@ -1,5 +1,6 @@ package cn.smartjavaai.ocr.entity; +import cn.smartjavaai.common.entity.DetectionRectangle; import cn.smartjavaai.common.entity.Point; import lombok.Data; @@ -49,4 +50,23 @@ public class OcrBox { (float)bottomLeft.getX(), (float)bottomLeft.getY() }; } + + /** + * 转换为 DetectionRectangle，使用最小外包矩形 + */ + public DetectionRectangle toDetectionRectangle() { + float[] pts = toFloatArray(); + float minX = Math.min(Math.min(pts[0], pts[2]), Math.min(pts[4], pts[6])); + float minY = Math.min(Math.min(pts[1], pts[3]), Math.min(pts[5], pts[7])); + float maxX = Math.max(Math.max(pts[0], pts[2]), Math.max(pts[4], pts[6])); + float maxY = Math.max(Math.max(pts[1], pts[3]), Math.max(pts[5], pts[7])); + DetectionRectangle rect = new DetectionRectangle(); + rect.setX((int) minX); + rect.setY((int) minY); + rect.setWidth((int) (maxX - minX)); + rect.setHeight((int) (maxY - minY)); + return rect; + } + + } diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrInfo.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrInfo.java index f287033..9661fdc 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrInfo.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrInfo.java @@ -4,6 +4,7 @@ import lombok.Data; import java.util.ArrayList; import java.util.List; +import java.util.stream.Collectors; /** * OCR信息 @@ -15,13 +16,22 @@ public class OcrInfo { private List> lineList; + private List ocrItemList; + private String fullText; + public OcrInfo(List> lineList, String fullText) { this.lineList = lineList; this.fullText = fullText; } public OcrInfo() { } + + public List flattenLines() { + return lineList.stream() + .flatMap(List::stream) + .collect(Collectors.toList()); + } } diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrItem.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrItem.java index 9c74752..19d02c0 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrItem.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrItem.java @@ -31,7 +31,6 @@ public class OcrItem { private float score; - public OcrItem(OcrBox ocrBox, String text) { this.ocrBox = ocrBox; this.text = text; diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/TableStructureResult.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/TableStructureResult.java new file mode 100644 index 0000000..725bea2 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/entity/TableStructureResult.java @@ -0,0 +1,33 @@ +package cn.smartjavaai.ocr.entity; + +import lombok.Data; + +import java.util.List; + +/** + * @author dwj + */ +@Data +public class TableStructureResult { + + private List ocrItemList; + + private List tableTagList; + + private String html; + + + public TableStructureResult(List ocrItemList, List tableTagList) { + this.ocrItemList = ocrItemList; + this.tableTagList = tableTagList; + } + + public TableStructureResult() { + } + + public TableStructureResult(List ocrItemList, List tableTagList, String html) { + this.ocrItemList = ocrItemList; + this.tableTagList = tableTagList; + this.html = html; + } +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/CommonDetModelEnum.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/CommonDetModelEnum.java index d3aaae3..cb34e1b 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/CommonDetModelEnum.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/CommonDetModelEnum.java @@ -3,11 +3,16 @@ package cn.smartjavaai.ocr.enums; /** * OCR检测模型枚举 * @author dwj - * @date 2025/4/4 */ public enum CommonDetModelEnum { - PADDLEOCR_V5_DET_MODEL; + PP_OCR_V5_SERVER_DET_MODEL, + + PP_OCR_V5_MOBILE_DET_MODEL, + + PP_OCR_V4_SERVER_DET_MODEL, + + PP_OCR_V4_MOBILE_DET_MODEL; /** diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/CommonRecModelEnum.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/CommonRecModelEnum.java index 2348dcf..f3dc0b8 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/CommonRecModelEnum.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/CommonRecModelEnum.java @@ -3,11 +3,16 @@ package cn.smartjavaai.ocr.enums; /** * OCR识别模型枚举 * @author dwj - * @date 2025/4/4 */ public enum CommonRecModelEnum { - PADDLEOCR_V5_REC_MODEL; + PP_OCR_V5_SERVER_REC_MODEL, + + PP_OCR_V5_MOBILE_REC_MODEL, + + PP_OCR_V4_SERVER_REC_MODEL, + + PP_OCR_V4_MOBILE_REC_MODEL; /** diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/DirectionModelEnum.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/DirectionModelEnum.java index a7f50c5..17afe4c 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/DirectionModelEnum.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/DirectionModelEnum.java @@ -7,7 +7,12 @@ package cn.smartjavaai.ocr.enums; */ public enum DirectionModelEnum { - CH_PPOCR_MOBILE_V2_CLS; + CH_PPOCR_MOBILE_V2_CLS, + + PP_LCNET_X0_25, + + PP_LCNET_X1_0; + /** diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/TableStructureModelEnum.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/TableStructureModelEnum.java new file mode 100644 index 0000000..795e17a --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/enums/TableStructureModelEnum.java @@ -0,0 +1,30 @@ +package cn.smartjavaai.ocr.enums; + +/** + * OCR表格结构模型枚举 + * @author dwj + */ +public enum TableStructureModelEnum { + + SLANET, + //SLANEXT_WIRED, + SLANET_PLUS; + + + + + /** + * 根据名称获取枚举 (忽略大小写和下划线变体) + */ + public static TableStructureModelEnum fromName(String name) { + String formatted = name.trim().toUpperCase().replaceAll("[-_]", ""); + for (TableStructureModelEnum 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/factory/OcrModelFactory.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/factory/OcrModelFactory.java index 4552092..7f79255 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/factory/OcrModelFactory.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/factory/OcrModelFactory.java @@ -3,14 +3,17 @@ package cn.smartjavaai.ocr.factory; import cn.smartjavaai.common.config.Config; import cn.smartjavaai.ocr.config.DirectionModelConfig; import cn.smartjavaai.ocr.config.OcrRecModelConfig; +import cn.smartjavaai.ocr.enums.CommonDetModelEnum; +import cn.smartjavaai.ocr.enums.CommonRecModelEnum; +import cn.smartjavaai.ocr.enums.DirectionModelEnum; import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel; -import cn.smartjavaai.ocr.model.common.detect.PpOCRV5DetModel; -import cn.smartjavaai.ocr.model.common.recognize.OcrCommonRecModel; +import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModelImpl; import cn.smartjavaai.ocr.config.OcrDetModelConfig; import cn.smartjavaai.ocr.exception.OcrException; import cn.smartjavaai.ocr.model.common.direction.OcrDirectionModel; -import cn.smartjavaai.ocr.model.common.direction.PPOCRMobileV2Model; -import cn.smartjavaai.ocr.model.common.recognize.PpOCRV5RecModel; +import cn.smartjavaai.ocr.model.common.direction.PPOCRMobileV2ClsModel; +import cn.smartjavaai.ocr.model.common.recognize.OcrCommonRecModel; +import cn.smartjavaai.ocr.model.common.recognize.OcrCommonRecModelImpl; import lombok.extern.slf4j.Slf4j; import java.util.Map; @@ -27,29 +30,29 @@ public class OcrModelFactory { // 使用 volatile 和双重检查锁定来确保线程安全的单例模式 private static volatile OcrModelFactory instance; - private static final ConcurrentHashMap commonDetModelMap = new ConcurrentHashMap<>(); + private static final ConcurrentHashMap commonDetModelMap = new ConcurrentHashMap<>(); - private static final ConcurrentHashMap commonRecModelMap = new ConcurrentHashMap<>(); + private static final ConcurrentHashMap commonRecModelMap = new ConcurrentHashMap<>(); - private static final ConcurrentHashMap directionModelMap = new ConcurrentHashMap<>(); + private static final ConcurrentHashMap directionModelMap = new ConcurrentHashMap<>(); /** * 检测模型注册表 */ - private static final Map> commonDetRegistry = + private static final Map> commonDetRegistry = new ConcurrentHashMap<>(); /** * 识别模型注册表 */ - private static final Map> commonRecRegistry = + private static final Map> commonRecRegistry = new ConcurrentHashMap<>(); /** * 方向分类模型注册表 */ - private static final Map> directionRegistry = + private static final Map> directionRegistry = new ConcurrentHashMap<>(); @@ -68,29 +71,29 @@ public class OcrModelFactory { /** * 注册通用检测模型 - * @param name + * @param detModelEnum * @param clazz */ - private static void registerCommonDetModel(String name, Class clazz) { - commonDetRegistry.put(name.toLowerCase(), clazz); + private static void registerCommonDetModel(CommonDetModelEnum detModelEnum, Class clazz) { + commonDetRegistry.put(detModelEnum, clazz); } /** * 注册通用识别模型 - * @param name + * @param recModelEnum * @param clazz */ - private static void registerCommonRecModel(String name, Class clazz) { - commonRecRegistry.put(name.toLowerCase(), clazz); + private static void registerCommonRecModel(CommonRecModelEnum recModelEnum, Class clazz) { + commonRecRegistry.put(recModelEnum, clazz); } /** * 注册通用方向分类模型 - * @param name + * @param directionModelEnum * @param clazz */ - private static void registerDirectionModel(String name, Class clazz) { - directionRegistry.put(name.toLowerCase(), clazz); + private static void registerDirectionModel(DirectionModelEnum directionModelEnum, Class clazz) { + directionRegistry.put(directionModelEnum, clazz); } @@ -103,7 +106,7 @@ public class OcrModelFactory { if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){ throw new OcrException("未配置OCR模型"); } - return commonDetModelMap.computeIfAbsent(config.getModelEnum().name(), k -> { + return commonDetModelMap.computeIfAbsent(config.getModelEnum(), k -> { return createCommonDetModel(config); }); } @@ -117,7 +120,7 @@ public class OcrModelFactory { if(Objects.isNull(config) || Objects.isNull(config.getRecModelEnum())){ throw new OcrException("未配置OCR模型"); } - return commonRecModelMap.computeIfAbsent(config.getRecModelEnum().name(), k -> { + return commonRecModelMap.computeIfAbsent(config.getRecModelEnum(), k -> { return createCommonRecModel(config); }); } @@ -131,7 +134,7 @@ public class OcrModelFactory { if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){ throw new OcrException("未配置OCR模型"); } - return directionModelMap.computeIfAbsent(config.getModelEnum().name(), k -> { + return directionModelMap.computeIfAbsent(config.getModelEnum(), k -> { return createDirectionModel(config); }); } @@ -144,7 +147,7 @@ public class OcrModelFactory { * @return */ private OcrCommonDetModel createCommonDetModel(OcrDetModelConfig config) { - Class clazz = commonDetRegistry.get(config.getModelEnum().name().toLowerCase()); + Class clazz = commonDetRegistry.get(config.getModelEnum()); if(clazz == null){ throw new OcrException("Unsupported model"); } @@ -165,7 +168,7 @@ public class OcrModelFactory { * @return */ private OcrCommonRecModel createCommonRecModel(OcrRecModelConfig config) { - Class clazz = commonRecRegistry.get(config.getRecModelEnum().name().toLowerCase()); + Class clazz = commonRecRegistry.get(config.getRecModelEnum()); if(clazz == null){ throw new OcrException("Unsupported model"); } @@ -185,7 +188,7 @@ public class OcrModelFactory { * @return */ private OcrDirectionModel createDirectionModel(DirectionModelConfig config) { - Class clazz = directionRegistry.get(config.getModelEnum().name().toLowerCase()); + Class clazz = directionRegistry.get(config.getModelEnum()); if(clazz == null){ throw new OcrException("Unsupported model"); } @@ -202,9 +205,18 @@ public class OcrModelFactory { // 初始化默认算法 static { - registerCommonDetModel("PADDLEOCR_V5_DET_MODEL", PpOCRV5DetModel.class); - registerCommonRecModel("PADDLEOCR_V5_REC_MODEL", PpOCRV5RecModel.class); - registerDirectionModel("CH_PPOCR_MOBILE_V2_CLS", PPOCRMobileV2Model.class); + //通用-检测模型 + registerCommonDetModel(CommonDetModelEnum.PP_OCR_V5_SERVER_DET_MODEL, OcrCommonDetModelImpl.class); + registerCommonDetModel(CommonDetModelEnum.PP_OCR_V5_MOBILE_DET_MODEL, OcrCommonDetModelImpl.class); + registerCommonDetModel(CommonDetModelEnum.PP_OCR_V4_SERVER_DET_MODEL, OcrCommonDetModelImpl.class); + registerCommonDetModel(CommonDetModelEnum.PP_OCR_V4_MOBILE_DET_MODEL, OcrCommonDetModelImpl.class); + registerCommonRecModel(CommonRecModelEnum.PP_OCR_V5_SERVER_REC_MODEL, OcrCommonRecModelImpl.class); + registerCommonRecModel(CommonRecModelEnum.PP_OCR_V5_MOBILE_REC_MODEL, OcrCommonRecModelImpl.class); + registerCommonRecModel(CommonRecModelEnum.PP_OCR_V4_SERVER_REC_MODEL, OcrCommonRecModelImpl.class); + registerCommonRecModel(CommonRecModelEnum.PP_OCR_V4_MOBILE_REC_MODEL, OcrCommonRecModelImpl.class); + registerDirectionModel(DirectionModelEnum.CH_PPOCR_MOBILE_V2_CLS, PPOCRMobileV2ClsModel.class); + registerDirectionModel(DirectionModelEnum.PP_LCNET_X0_25, PPOCRMobileV2ClsModel.class); + registerDirectionModel(DirectionModelEnum.PP_LCNET_X1_0, PPOCRMobileV2ClsModel.class); log.debug("缓存目录：{}", Config.getCachePath()); } diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/factory/TableRecModelFactory.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/factory/TableRecModelFactory.java new file mode 100644 index 0000000..9723173 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/factory/TableRecModelFactory.java @@ -0,0 +1,119 @@ +package cn.smartjavaai.ocr.factory; + +import cn.smartjavaai.common.config.Config; +import cn.smartjavaai.ocr.config.DirectionModelConfig; +import cn.smartjavaai.ocr.config.OcrDetModelConfig; +import cn.smartjavaai.ocr.config.OcrRecModelConfig; +import cn.smartjavaai.ocr.config.TableStructureConfig; +import cn.smartjavaai.ocr.enums.CommonDetModelEnum; +import cn.smartjavaai.ocr.enums.CommonRecModelEnum; +import cn.smartjavaai.ocr.enums.DirectionModelEnum; +import cn.smartjavaai.ocr.enums.TableStructureModelEnum; +import cn.smartjavaai.ocr.exception.OcrException; +import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel; +import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModelImpl; +import cn.smartjavaai.ocr.model.common.direction.OcrDirectionModel; +import cn.smartjavaai.ocr.model.common.direction.PPOCRMobileV2ClsModel; +import cn.smartjavaai.ocr.model.common.recognize.OcrCommonRecModel; +import cn.smartjavaai.ocr.model.common.recognize.OcrCommonRecModelImpl; +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; + +/** + * OCR 表格识别模型工厂 + * @author dwj + */ +@Slf4j +public class TableRecModelFactory { + + // 使用 volatile 和双重检查锁定来确保线程安全的单例模式 + private static volatile TableRecModelFactory instance; + + /** + * 模型缓存 + */ + private static final ConcurrentHashMap tableStructureModelMap = new ConcurrentHashMap<>(); + + + /** + * 模型注册表 + */ + private static final Map> tableStructureRegistry = + new ConcurrentHashMap<>(); + + + public static TableRecModelFactory getInstance() { + if (instance == null) { + synchronized (TableRecModelFactory.class) { + if (instance == null) { + instance = new TableRecModelFactory(); + } + } + } + return instance; + } + + + + /** + * 注册模型 + * @param tableStructureModelEnum + * @param clazz + */ + private static void registerTableStructureModel(TableStructureModelEnum tableStructureModelEnum, Class clazz) { + tableStructureRegistry.put(tableStructureModelEnum, clazz); + } + + + /** + * 获取模型（通过配置） + * @param config + * @return + */ + public TableStructureModel getTableStructureModel(TableStructureConfig config) { + if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){ + throw new OcrException("未配置OCR模型"); + } + return tableStructureModelMap.computeIfAbsent(config.getModelEnum(), k -> { + return createTableStructureModel(config); + }); + } + + + + /** + * 创建模型 + * @param config + * @return + */ + private TableStructureModel createTableStructureModel(TableStructureConfig config) { + Class clazz = tableStructureRegistry.get(config.getModelEnum()); + if(clazz == null){ + throw new OcrException("Unsupported model"); + } + TableStructureModel model = null; + try { + model = (TableStructureModel) clazz.newInstance(); + } catch (InstantiationException | IllegalAccessException e) { + throw new OcrException(e); + } + model.loadModel(config); + return model; + } + + + + + // 初始化默认算法 + static { + registerTableStructureModel(TableStructureModelEnum.SLANET, CommonTableStructureModel.class); + registerTableStructureModel(TableStructureModelEnum.SLANET_PLUS, CommonTableStructureModel.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 b28f232..aa29177 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 @@ -29,8 +29,6 @@ public interface OcrCommonDetModel extends AutoCloseable{ } - - /** * 文本检测 * @param image BufferedImage @@ -77,4 +75,23 @@ public interface OcrCommonDetModel extends AutoCloseable{ throw new UnsupportedOperationException("默认不支持该功能"); } + /** + * 文本检测（批量） + * @param imageList BufferedImage + * @return + */ + default List> batchDetect(List imageList) { + throw new UnsupportedOperationException("默认不支持该功能"); + } + + /** + * 文本检测（批量） + * @param imageList DJL Image + * @return + */ + default List> batchDetectDJLImage(List imageList){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + + } diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/PpOCRV5DetModel.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/OcrCommonDetModelImpl.java similarity index 81% rename from smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/PpOCRV5DetModel.java rename to smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/OcrCommonDetModelImpl.java index a09fee4..648c68c 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/PpOCRV5DetModel.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/OcrCommonDetModelImpl.java @@ -1,20 +1,16 @@ package cn.smartjavaai.ocr.model.common.detect; -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; -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 ai.djl.training.util.ProgressBar; -import cn.smartjavaai.common.enums.DeviceEnum; import cn.smartjavaai.common.pool.PredictorFactory; import cn.smartjavaai.common.utils.FileUtils; import cn.smartjavaai.common.utils.ImageUtils; @@ -22,7 +18,7 @@ import cn.smartjavaai.common.utils.OpenCVUtils; import cn.smartjavaai.ocr.config.OcrDetModelConfig; import cn.smartjavaai.ocr.entity.OcrBox; import cn.smartjavaai.ocr.exception.OcrException; -import cn.smartjavaai.ocr.model.common.detect.translator.PPOCRV5DetTranslator; +import cn.smartjavaai.ocr.model.common.detect.criteria.OcrCommonDetCriterialFactory; import cn.smartjavaai.ocr.utils.OcrUtils; import lombok.extern.slf4j.Slf4j; import org.apache.commons.lang3.StringUtils; @@ -38,18 +34,14 @@ import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.Paths; -import java.util.List; -import java.util.Objects; -import java.util.concurrent.ConcurrentHashMap; +import java.util.*; /** - * PPOCRV5 检测模型 + * ocr通用检测模型实现类 * @author dwj - * @date 2025/4/21 */ @Slf4j -public class PpOCRV5DetModel implements OcrCommonDetModel { - +public class OcrCommonDetModelImpl implements OcrCommonDetModel{ private ObjectPool> detPredictorPool; @@ -62,22 +54,9 @@ public class PpOCRV5DetModel implements OcrCommonDetModel { if(StringUtils.isBlank(config.getDetModelPath())){ throw new OcrException("modelPath is null"); } - Device device = null; - if(!Objects.isNull(config.getDevice())){ - device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(); - } this.config = config; //初始化 检测Criteria - Criteria detCriteria = - Criteria.builder() - .optEngine("OnnxRuntime") - .setTypes(Image.class, NDList.class) - .optModelPath(Paths.get(config.getDetModelPath())) - .optTranslator(new PPOCRV5DetTranslator(new ConcurrentHashMap())) - .optDevice(device) - .optProgress(new ProgressBar()) - .build(); - + Criteria detCriteria = OcrCommonDetCriterialFactory.createCriteria(config); try{ detectionModel = ModelZoo.loadModel(detCriteria); // 创建池子：每个线程独享 Predictor @@ -107,28 +86,9 @@ public class PpOCRV5DetModel implements OcrCommonDetModel { @Override public List detect(Image image){ - Predictor predictor = null; - try (NDManager manager = NDManager.newBaseManager()) { - predictor = detPredictorPool.borrowObject(); - NDList result = predictor.predict(image); - result.attach(manager); - return OcrUtils.convertToOcrBox(result, image); - } catch (Exception e) { - throw new OcrException("OCR检测错误", 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); - } - } - } - } + List imageList = Collections.singletonList(image); + List> result = batchDetectDJLImage(imageList); + return result.get(0); } @Override @@ -136,7 +96,7 @@ public class PpOCRV5DetModel implements OcrCommonDetModel { if(!FileUtils.isFileExists(imagePath)){ throw new OcrException("图像文件不存在"); } - try (NDManager manager = NDManager.newBaseManager()) { + try { Image img = ImageFactory.getInstance().fromFile(Paths.get(imagePath)); List boxList = detect(img); if(Objects.isNull(boxList) || boxList.isEmpty()){ @@ -194,10 +154,58 @@ public class PpOCRV5DetModel implements OcrCommonDetModel { img.save(outputStream, "png"); // 将字节流转换为 BufferedImage byte[] imageBytes = outputStream.toByteArray(); - ((Mat) img.getWrappedImage()).release(); return ImageIO.read(new ByteArrayInputStream(imageBytes)); } catch (IOException e) { throw new OcrException("导出图片失败", e); + } finally { + if (img != null){ + ((Mat) img.getWrappedImage()).release(); + } + } + } + + @Override + public List> batchDetect(List imageList) { + List djlImageList = new ArrayList<>(imageList.size()); + try { + for (BufferedImage bufferedImage : imageList) { + djlImageList.add(ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(bufferedImage))); + } + return batchDetectDJLImage(djlImageList); + } catch (Exception e) { + throw new OcrException(e); + } finally { + djlImageList.forEach(image -> ((Mat)image.getWrappedImage()).release()); + } + } + + @Override + public List> batchDetectDJLImage(List imageList) { + if(!ImageUtils.isAllImageSizeEqual(imageList)){ + throw new OcrException("图片尺寸不一致"); + } + Predictor predictor = null; + try (NDManager manager = NDManager.newBaseManager()) { + predictor = detPredictorPool.borrowObject(); + List result = predictor.batchPredict(imageList); + result.forEach(ndList -> ndList.attach(manager)); + return OcrUtils.convertToOcrBox(result); + } catch (Exception e) { + throw new OcrException("OCR检测错误", 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); + } + } + } + } } @@ -218,4 +226,6 @@ public class PpOCRV5DetModel implements OcrCommonDetModel { log.warn("关闭 model 失败", e); } } + + } 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 new file mode 100644 index 0000000..74a6059 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/criteria/OcrCommonDetCriterialFactory.java @@ -0,0 +1,53 @@ +package cn.smartjavaai.ocr.model.common.detect.criteria; + +import ai.djl.Device; +import ai.djl.modality.cv.Image; +import ai.djl.ndarray.NDList; +import ai.djl.repository.zoo.Criteria; +import ai.djl.training.util.ProgressBar; +import cn.smartjavaai.common.enums.DeviceEnum; +import cn.smartjavaai.ocr.config.OcrDetModelConfig; +import cn.smartjavaai.ocr.enums.CommonDetModelEnum; +import cn.smartjavaai.ocr.model.common.detect.translator.PPOCRDetTranslator; +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 OcrCommonDetCriterialFactory { + + + public static Criteria createCriteria(OcrDetModelConfig config) { + Device device = null; + if(!Objects.isNull(config.getDevice())){ + device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(); + } + Criteria criteria = null; + ConcurrentHashMap params = new ConcurrentHashMap(); + params.putAll(config.getCustomParams()); + if(StringUtils.isNotBlank(config.getBatchifier())){ + params.put("batchifier", config.getBatchifier()); + } + if(config.getModelEnum() == CommonDetModelEnum.PP_OCR_V5_SERVER_DET_MODEL || + config.getModelEnum() == CommonDetModelEnum.PP_OCR_V5_MOBILE_DET_MODEL || + config.getModelEnum() == CommonDetModelEnum.PP_OCR_V4_SERVER_DET_MODEL || + config.getModelEnum() == CommonDetModelEnum.PP_OCR_V4_MOBILE_DET_MODEL + ){ + criteria = + Criteria.builder() + .optEngine("OnnxRuntime") + .setTypes(Image.class, NDList.class) + .optModelPath(Paths.get(config.getDetModelPath())) + .optTranslator(new PPOCRDetTranslator(params)) + .optDevice(device) + .optProgress(new ProgressBar()) + .build(); + } + return criteria; + } +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/translator/PPOCRV5DetTranslator.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/translator/PPOCRDetTranslator.java similarity index 95% rename from smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/translator/PPOCRV5DetTranslator.java rename to smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/translator/PPOCRDetTranslator.java index 9725561..23476d0 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/translator/PPOCRV5DetTranslator.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/detect/translator/PPOCRDetTranslator.java @@ -27,24 +27,35 @@ import java.util.Map; * @mail 179209347@qq.com * @website www.aias.top */ -public class PPOCRV5DetTranslator implements Translator { +public class PPOCRDetTranslator implements Translator { // det_algorithm == "DB" private final float thresh = 0.3f; private final boolean use_dilation = false; private final String score_mode = "fast"; private final String box_type = "quad"; + //检测的图像边长限制 private final int limit_side_len; + //输出的最大文本框数量 private final int max_candidates; + //文本框最小尺寸阈值 private final int min_size; + //文本框的分数阈值 private final float box_thresh; + + /** + * 这个参数是检测后处理时控制文本框大小的，默认1.6，可以尝试改成2.5或者更大，反之，如果觉得文本框不够紧凑，也可以把该参数调小。 + * 检测框大小过于紧贴文字或检测框过大，可以调整db_unclip_ratio这个参数，加大参数可以扩大检测框，减小参数可以减小检测框大小； + */ private final float unclip_ratio; private float ratio_h; private float ratio_w; private int img_height; private int img_width; - public PPOCRV5DetTranslator(Map arguments) { + private String batchifier; + + public PPOCRDetTranslator(Map arguments) { limit_side_len = arguments.containsKey("limit_side_len") ? Integer.parseInt(arguments.get("limit_side_len").toString()) @@ -65,6 +76,10 @@ public class PPOCRV5DetTranslator implements Translator { arguments.containsKey("unclip_ratio") ? Float.parseFloat(arguments.get("unclip_ratio").toString()) : 1.6f; + + batchifier = arguments.containsKey("batchifier") + ? arguments.get("batchifier").toString() + : "stack"; } @Override @@ -509,13 +524,13 @@ public class PPOCRV5DetTranslator implements Translator { new float[]{0.485f, 0.456f, 0.406f}, new float[]{0.229f, 0.224f, 0.225f}); - img = img.expandDims(0); +// img = img.expandDims(0); return new NDList(img); } @Override public Batchifier getBatchifier() { - return null; + return Batchifier.fromString(batchifier); } } 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 c2f45c2..ef2e8a4 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 @@ -8,6 +8,7 @@ import cn.smartjavaai.ocr.entity.DirectionInfo; 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.opencv.core.Mat; import java.awt.image.BufferedImage; @@ -19,6 +20,14 @@ import java.util.List; */ public interface OcrDirectionModel extends AutoCloseable{ + default void setTextDetModel(OcrCommonDetModel detModel){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + + default OcrCommonDetModel getTextDetModel(){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + /** * 加载模型 * @param config @@ -71,7 +80,11 @@ public interface OcrDirectionModel extends AutoCloseable{ * @param manager * @return */ - default List detect(List boxList, Mat srcMat, NDManager manager) { + default List detect(List boxList, Mat srcMat) { + throw new UnsupportedOperationException("默认不支持该功能"); + } + + default List> batchDetect(List> boxList, List srcMatList) { throw new UnsupportedOperationException("默认不支持该功能"); } diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/PPOCRMobileV2Model.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/PPOCRMobileV2ClsModel.java similarity index 51% rename from smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/PPOCRMobileV2Model.java rename to smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/PPOCRMobileV2ClsModel.java index 37c4e2d..483ada0 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/PPOCRMobileV2Model.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/PPOCRMobileV2ClsModel.java @@ -6,9 +6,7 @@ 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.Point; import ai.djl.ndarray.NDManager; -import ai.djl.opencv.OpenCVImageFactory; import ai.djl.repository.zoo.Criteria; import ai.djl.repository.zoo.ModelNotFoundException; import ai.djl.repository.zoo.ModelZoo; @@ -20,17 +18,15 @@ import cn.smartjavaai.common.utils.FileUtils; import cn.smartjavaai.common.utils.ImageUtils; import cn.smartjavaai.common.utils.OpenCVUtils; import cn.smartjavaai.ocr.config.DirectionModelConfig; -import cn.smartjavaai.ocr.config.OcrDetModelConfig; import cn.smartjavaai.ocr.entity.*; import cn.smartjavaai.ocr.enums.AngleEnum; import cn.smartjavaai.ocr.exception.OcrException; -import cn.smartjavaai.ocr.factory.OcrModelFactory; import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel; +import cn.smartjavaai.ocr.model.common.direction.criteria.DirectionCriteriaFactory; import cn.smartjavaai.ocr.model.common.direction.translator.PpWordRotateTranslator; -import cn.smartjavaai.ocr.opencv.OcrNDArrayUtils; -import cn.smartjavaai.ocr.opencv.OcrOpenCVUtils; 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; @@ -45,9 +41,10 @@ import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.Paths; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Objects; -import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; /** * PPOCRMobileV2Model 方向分类模型 @@ -55,38 +52,34 @@ import java.util.UUID; * @date 2025/4/21 */ @Slf4j -public class PPOCRMobileV2Model implements OcrDirectionModel { +public class PPOCRMobileV2ClsModel implements OcrDirectionModel { private ObjectPool> predictorPool; private DirectionModelConfig config; - private OcrCommonDetModel detModel; - private ZooModel model; + private OcrCommonDetModel textDetModel; + @Override public void loadModel(DirectionModelConfig config){ if(StringUtils.isBlank(config.getModelPath())){ throw new OcrException("modelPath is null"); } - this.config = config; + this.textDetModel = config.getTextDetModel(); Device device = null; if(!Objects.isNull(config.getDevice())){ device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(); } - Criteria criteria = - Criteria.builder() - .optEngine("OnnxRuntime") - .setTypes(Image.class, DirectionInfo.class) - .optModelPath(Paths.get(config.getModelPath())) - .optDevice(device) - .optTranslator(new PpWordRotateTranslator()) - .optProgress(new ProgressBar()) - .build(); + ConcurrentHashMap params = new ConcurrentHashMap(); + if(StringUtils.isNotBlank(config.getBatchifier())){ + params.put("batchifier", config.getBatchifier()); + } + Criteria criteria = DirectionCriteriaFactory.createCriteria(config); try{ model = ModelZoo.loadModel(criteria); // 创建池子：每个线程独享 Predictor @@ -96,15 +89,6 @@ public class PPOCRMobileV2Model implements OcrDirectionModel { } catch (IOException | ModelNotFoundException | MalformedModelException e) { throw new OcrException("模型加载失败", e); } - - //获取检测模型 - if(StringUtils.isNotBlank(config.getDetModelPath()) && Objects.nonNull(config.getDetModelEnum())){ - OcrDetModelConfig detModelConfig = new OcrDetModelConfig(); - detModelConfig.setModelEnum(config.getDetModelEnum()); - detModelConfig.setDetModelPath(config.getDetModelPath()); - detModelConfig.setDevice(config.getDevice()); - detModel = OcrModelFactory.getInstance().getDetModel(detModelConfig); - } } @Override @@ -115,124 +99,85 @@ public class PPOCRMobileV2Model implements OcrDirectionModel { Image img = null; try { img = ImageFactory.getInstance().fromFile(Paths.get(imagePath)); + return detect(img); } catch (IOException e) { throw new OcrException("无效的图片", e); + }finally { + if(img != null){ + ((Mat)img.getWrappedImage()).release(); + } } - List ocrItemList = detect(img); - ((Mat)img.getWrappedImage()).release(); - return ocrItemList; } @Override public List detect(Image image){ + if(Objects.isNull(textDetModel)){ + throw new OcrException("textDetModel is null"); + } //检测文本 - List boxeList = detModel.detect(image); + List boxeList = textDetModel.detect(image); if(Objects.isNull(boxeList) || boxeList.isEmpty()){ throw new OcrException("未检测到文本"); } - Predictor predictor = null; - List ocrItemList = new ArrayList<>(); - try (NDManager manager = NDManager.newBaseManager()) { - Mat srcMat = (Mat) image.getWrappedImage(); - predictor = predictorPool.borrowObject(); - for (OcrBox box : boxeList){ - OcrItem ocrItem = detect(box, srcMat, predictor, manager); - ocrItemList.add(ocrItem); - } - return ocrItemList; - } catch (Exception e) { - throw new OcrException("OCR检测错误", e); - }finally { - if (predictor != null) { - try { - predictorPool.returnObject(predictor); //归还 - } catch (Exception e) { - log.warn("归还Predictor失败", e); - try { - predictor.close(); // 归还失败才销毁 - } catch (Exception ex) { - log.error("关闭Predictor失败", ex); - } - } - } - } + Mat srcMat = (Mat) image.getWrappedImage(); + return detect(boxeList, srcMat); } - /** - * 基于文本框检测方向 - * @param box - * @param srcMat - * @param predictor - * @param manager - * @return - */ - private OcrItem detect(OcrBox box, Mat srcMat, Predictor predictor, NDManager manager){ - if(Objects.isNull(box)){ - throw new OcrException("box参数为空"); - } - try { - //透视变换及裁剪 - Image subImg = OcrUtils.transformAndCrop(srcMat, box); - DirectionInfo directionInfo = null; - String angle; - //高宽比 > 1.5 纵向 - if (subImg.getHeight() * 1.0 / subImg.getWidth() > 1.5) { - //旋转图片90度 - subImg = OcrUtils.rotateImg(manager, subImg); - //检测方向 - directionInfo = predictor.predict(subImg); - if (directionInfo.getName().equalsIgnoreCase("Rotate")) { - angle = "270"; - } else { - angle = "90"; - } - }else{ //横向 - directionInfo = predictor.predict(subImg); - if (directionInfo.getName().equalsIgnoreCase("No Rotate")) { - angle = "0"; - } else { - angle = "180"; - } - } - ((Mat)subImg.getWrappedImage()).release(); - return new OcrItem(box, AngleEnum.fromValue(angle), directionInfo.getProb().floatValue()); - } catch (Exception e) { - throw new OcrException("OCR检测错误", e); - } - } +// /** +// * 基于文本框检测方向 +// * @param box +// * @param srcMat +// * @param predictor +// * @param manager +// * @return +// */ +// private OcrItem detect(OcrBox box, Mat srcMat, Predictor predictor, NDManager manager){ +// if(Objects.isNull(box)){ +// throw new OcrException("box参数为空"); +// } +// try { +// //透视变换及裁剪 +// Image subImg = OcrUtils.transformAndCrop(srcMat, box); +// DirectionInfo directionInfo = null; +// String angle; +// //高宽比 > 1.5 纵向 +// if (subImg.getHeight() * 1.0 / subImg.getWidth() > 1.5) { +// //旋转图片90度 +// subImg = OcrUtils.rotateImg(manager, subImg); +// //检测方向 +// directionInfo = predictor.predict(subImg); +// if (directionInfo.getName().equalsIgnoreCase("Rotate")) { +// angle = "270"; +// } else { +// angle = "90"; +// } +// }else{ //横向 +// directionInfo = predictor.predict(subImg); +// if (directionInfo.getName().equalsIgnoreCase("No Rotate")) { +// angle = "0"; +// } else { +// angle = "180"; +// } +// } +// ((Mat)subImg.getWrappedImage()).release(); +// return new OcrItem(box, AngleEnum.fromValue(angle), directionInfo.getProb().floatValue()); +// } catch (Exception e) { +// throw new OcrException("OCR检测错误", e); +// } +// } @Override - public List detect(List boxList,Mat srcMat,NDManager manager){ + public List detect(List boxList,Mat srcMat){ if(Objects.isNull(boxList) || boxList.isEmpty()){ throw new OcrException("boxList为空"); } - Predictor predictor = null; - List ocrItemList = new ArrayList<>(); - try { - predictor = predictorPool.borrowObject(); - for (OcrBox box : boxList){ - OcrItem ocrItem = detect(box, srcMat, predictor, manager); - ocrItemList.add(ocrItem); - } - } catch (Exception e) { - throw new OcrException("OCR检测错误", e); - }finally { - if (predictor != null) { - try { - predictorPool.returnObject(predictor); //归还 - } catch (Exception e) { - log.warn("归还Predictor失败", e); - try { - predictor.close(); // 归还失败才销毁 - } catch (Exception ex) { - log.error("关闭Predictor失败", ex); - } - } - } + List> ocrItemList = batchDetect(Collections.singletonList(boxList), Collections.singletonList(srcMat)); + if(Objects.isNull(ocrItemList) || ocrItemList.isEmpty()){ + throw new OcrException("方向检测失败"); } - return ocrItemList; + return ocrItemList.get(0); } @Override @@ -240,8 +185,9 @@ public class PPOCRMobileV2Model implements OcrDirectionModel { if(!FileUtils.isFileExists(imagePath)){ throw new OcrException("图像文件不存在"); } - try (NDManager manager = NDManager.newBaseManager()) { - Image img = ImageFactory.getInstance().fromFile(Paths.get(imagePath)); + Image img = null; + try { + img = ImageFactory.getInstance().fromFile(Paths.get(imagePath)); List itemList = detect(img); if(Objects.isNull(itemList) || itemList.isEmpty()){ throw new OcrException("未检测到文字"); @@ -250,9 +196,12 @@ public class PPOCRMobileV2Model implements OcrDirectionModel { Path output = Paths.get(outputPath); log.debug("Saving to {}", output.toAbsolutePath().toString()); img.save(Files.newOutputStream(output), "png"); - ((Mat) img.getWrappedImage()).release(); } catch (IOException e) { throw new OcrException(e); + } finally { + if (img != null){ + ((Mat)img.getWrappedImage()).release(); + } } } @@ -297,13 +246,129 @@ public class PPOCRMobileV2Model implements OcrDirectionModel { img.save(outputStream, "png"); // 将字节流转换为 BufferedImage byte[] imageBytes = outputStream.toByteArray(); - ((Mat) img.getWrappedImage()).release(); return ImageIO.read(new ByteArrayInputStream(imageBytes)); } catch (IOException e) { throw new OcrException("导出图片失败", e); + } finally { + if (img != null){ + ((Mat) img.getWrappedImage()).release(); + } } } + @Override + public List> batchDetect(List> boxList, List srcMatList) { + if(CollectionUtils.isEmpty(boxList)){ + throw new OcrException("boxList 不能为空"); + } + if(CollectionUtils.isEmpty(srcMatList)){ + throw new OcrException("srcMatList 不能为空"); + } + //检查参数 + for (int i = 0; i < srcMatList.size(); i++) { + List ocrBoxes = boxList.get(i); + Mat mat = srcMatList.get(i); + if (ocrBoxes == null) { + throw new OcrException("第 " + i + " 个 boxList 为 null"); + } + if (ocrBoxes.isEmpty()) { + throw new OcrException("第 " + i + " 个 boxList 没有检测结果"); + } + if (mat.empty()) { + throw new OcrException("第 " + i + " 张图片为空 Mat"); + } + } + List imageList = new ArrayList(); + List isRotatedList = new ArrayList(); + int index = 0; + try (NDManager manager = model.getNDManager().newSubManager()){ + for(int i = 0; i < srcMatList.size(); i++){ + for (int j = 0; j < boxList.get(i).size(); j++){ + //透视变换及裁剪 + Image subImg = OcrUtils.transformAndCrop(srcMatList.get(i), boxList.get(i).get(j)); + //高宽比 > 1.5 纵向 + if (subImg.getHeight() * 1.0 / subImg.getWidth() > 1.5) { + //旋转图片90度 + subImg = OcrUtils.rotateImg(manager, subImg); + isRotatedList.add(true); + imageList.add(subImg); + }else{ + isRotatedList.add(false); + imageList.add(subImg); + } + index++; + } + } + List> result = new ArrayList<>(); + List directionInfos = batchDetect(imageList); + if(CollectionUtils.isEmpty(directionInfos)){ + throw new OcrException("方向检测失败"); + } + index = 0; + for(int i = 0; i < srcMatList.size(); i++){ + List ocrItemList = new ArrayList<>(); + for (int j = 0; j < boxList.get(i).size(); j++){ + DirectionInfo directionInfo = directionInfos.get(index); + if(Objects.isNull(directionInfo)){ + throw new OcrException("方向检测失败: 第" + i + "张图片, 第" + j + "个文本块，未检测到方向"); + } + String angle; + if(isRotatedList.get(index)){ + if (directionInfo.getName().equalsIgnoreCase("Rotate")) { + angle = "270"; + } else { + angle = "90"; + } + }else{ + if (directionInfo.getName().equalsIgnoreCase("No Rotate")) { + angle = "0"; + } else { + angle = "180"; + } + } + OcrItem ocrItem = new OcrItem(boxList.get(i).get(j), AngleEnum.fromValue(angle), directionInfo.getProb().floatValue()); + ocrItemList.add(ocrItem); + index++; + } + result.add(ocrItemList); + } + return result; + } + } + + private List batchDetect(List imageList) { + Predictor predictor = null; + try { + predictor = predictorPool.borrowObject(); + return predictor.batchPredict(imageList); + } catch (Exception e) { + throw new OcrException("OCR检测错误", e); + }finally { + if (predictor != null) { + try { + predictorPool.returnObject(predictor); //归还 + } catch (Exception e) { + log.warn("归还Predictor失败", e); + try { + predictor.close(); // 归还失败才销毁 + } catch (Exception ex) { + log.error("关闭Predictor失败", ex); + } + } + } + } + } + + @Override + public void setTextDetModel(OcrCommonDetModel detModel) { + this.textDetModel = detModel; + } + + @Override + public OcrCommonDetModel getTextDetModel() { + return textDetModel; + } + @Override public void close() throws Exception { try { @@ -320,12 +385,5 @@ public class PPOCRMobileV2Model implements OcrDirectionModel { } catch (Exception e) { log.warn("关闭 model 失败", e); } - try { - if (detModel != null) { - detModel.close(); - } - } catch (Exception e) { - log.warn("关闭 model 失败", e); - } } } 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 new file mode 100644 index 0000000..540b698 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/criteria/DirectionCriteriaFactory.java @@ -0,0 +1,61 @@ +package cn.smartjavaai.ocr.model.common.direction.criteria; + +import ai.djl.Device; +import ai.djl.modality.cv.Image; +import ai.djl.ndarray.NDList; +import ai.djl.repository.zoo.Criteria; +import ai.djl.training.util.ProgressBar; +import cn.smartjavaai.common.enums.DeviceEnum; +import cn.smartjavaai.ocr.config.DirectionModelConfig; +import cn.smartjavaai.ocr.config.OcrDetModelConfig; +import cn.smartjavaai.ocr.entity.DirectionInfo; +import cn.smartjavaai.ocr.enums.CommonDetModelEnum; +import cn.smartjavaai.ocr.enums.DirectionModelEnum; +import cn.smartjavaai.ocr.model.common.detect.translator.PPOCRDetTranslator; +import cn.smartjavaai.ocr.model.common.direction.translator.PpWordRotateTranslator; +import org.apache.commons.lang3.StringUtils; + +import java.nio.file.Paths; +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; + +/** + * 行方向分类 + * @author dwj + */ +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(); + } + Criteria criteria = null; + ConcurrentHashMap params = new ConcurrentHashMap(); + params.putAll(config.getCustomParams()); + if(StringUtils.isNotBlank(config.getBatchifier())){ + params.put("batchifier", config.getBatchifier()); + } + if(config.getModelEnum() == DirectionModelEnum.CH_PPOCR_MOBILE_V2_CLS){ + params.put("resizeWidth", 192); + params.put("resizeHeight", 48); + }else if (config.getModelEnum() == DirectionModelEnum.PP_LCNET_X0_25){ + params.put("resizeWidth", 160); + params.put("resizeHeight", 80); + }else if (config.getModelEnum() == DirectionModelEnum.PP_LCNET_X1_0){ + params.put("resizeWidth", 160); + params.put("resizeHeight", 80); + } + criteria = + Criteria.builder() + .optEngine("OnnxRuntime") + .setTypes(Image.class, DirectionInfo.class) + .optModelPath(Paths.get(config.getModelPath())) + .optDevice(device) + .optTranslator(new PpWordRotateTranslator(params)) + .optProgress(new ProgressBar()) + .build(); + return criteria; + } + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/translator/PpWordRotateTranslator.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/translator/PpWordRotateTranslator.java index c0e4e3a..3229f2f 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/translator/PpWordRotateTranslator.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/direction/translator/PpWordRotateTranslator.java @@ -13,6 +13,7 @@ import cn.smartjavaai.ocr.entity.DirectionInfo; import java.util.Arrays; import java.util.List; +import java.util.Map; /** * 方向检测 @@ -24,7 +25,24 @@ import java.util.List; public class PpWordRotateTranslator implements Translator { List classes = Arrays.asList("No Rotate", "Rotate"); - public PpWordRotateTranslator() { + private String batchifier; + + private int resizeHeight; + + private int resizeWidth; + + public PpWordRotateTranslator(Map arguments) { + batchifier = arguments.containsKey("batchifier") + ? arguments.get("batchifier").toString() + : "padding"; + + resizeWidth = arguments.containsKey("resizeWidth") + ? (Integer) arguments.get("resizeWidth") + : 192; + + resizeHeight = arguments.containsKey("resizeHeight") + ? (Integer) arguments.get("resizeHeight") + : 48; } @Override @@ -51,8 +69,8 @@ public class PpWordRotateTranslator implements Translator public NDList processInput(TranslatorContext ctx, Image input) { NDArray img = input.toNDArray(ctx.getNDManager()); int imgC = 3; - int imgH = 48; - int imgW = 192; + int imgH = resizeHeight; + int imgW = resizeWidth; NDArray array = ctx.getNDManager().zeros(new Shape(imgC, imgH, imgW)); @@ -74,13 +92,14 @@ public class PpWordRotateTranslator implements Translator array.set(new NDIndex(":,:,0:" + resized_w), img); - array = array.expandDims(0); +// array = array.expandDims(0); return new NDList(new NDArray[]{array}); } @Override public Batchifier getBatchifier() { - return null; + return Batchifier.fromString(batchifier); } + } 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 73dd78d..e9ac8b2 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 @@ -3,8 +3,11 @@ package cn.smartjavaai.ocr.model.common.recognize; import ai.djl.modality.cv.Image; import cn.smartjavaai.ocr.config.OcrDetModelConfig; import cn.smartjavaai.ocr.config.OcrRecModelConfig; +import cn.smartjavaai.ocr.config.OcrRecOptions; 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 java.awt.image.BufferedImage; import java.util.List; @@ -15,6 +18,22 @@ import java.util.List; */ public interface OcrCommonRecModel extends AutoCloseable{ + default void setTextDetModel(OcrCommonDetModel detModel){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + + default OcrCommonDetModel getTextDetModel(){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + + default void setDirectionModel(OcrDirectionModel directionModel){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + + default OcrDirectionModel getDirectionModel(){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + /** * 加载模型 * @param config @@ -26,7 +45,16 @@ public interface OcrCommonRecModel extends AutoCloseable{ * @param imagePath 图片路径 * @return */ - default OcrInfo recognize(String imagePath) { + default OcrInfo recognize(String imagePath, OcrRecOptions options) { + throw new UnsupportedOperationException("默认不支持该功能"); + } + + /** + * 文本识别 + * @param image + * @return + */ + default OcrInfo recognize(Image image, OcrRecOptions options) { throw new UnsupportedOperationException("默认不支持该功能"); } @@ -36,7 +64,7 @@ public interface OcrCommonRecModel extends AutoCloseable{ * @param image BufferedImage * @return */ - default OcrInfo recognize(BufferedImage image) { + default OcrInfo recognize(BufferedImage image, OcrRecOptions options) { throw new UnsupportedOperationException("默认不支持该功能"); } @@ -46,7 +74,7 @@ public interface OcrCommonRecModel extends AutoCloseable{ * @param imageData 图片字节数组 * @return */ - default OcrInfo recognize(byte[] imageData) { + default OcrInfo recognize(byte[] imageData, OcrRecOptions options) { throw new UnsupportedOperationException("默认不支持该功能"); } @@ -56,7 +84,7 @@ public interface OcrCommonRecModel extends AutoCloseable{ * @param imagePath * @param outputPath */ - default void recognizeAndDraw(String imagePath, String outputPath, int fontSize) { + default void recognizeAndDraw(String imagePath, String outputPath, int fontSize, OcrRecOptions options) { throw new UnsupportedOperationException("默认不支持该功能"); } @@ -65,7 +93,16 @@ public interface OcrCommonRecModel extends AutoCloseable{ * @param sourceImage * @return */ - default BufferedImage recognizeAndDraw(BufferedImage sourceImage, int fontSize){ + default BufferedImage recognizeAndDraw(BufferedImage sourceImage, int fontSize, OcrRecOptions options){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + + + default List batchRecognize(List imageList, OcrRecOptions options) { + throw new UnsupportedOperationException("默认不支持该功能"); + } + + default List batchRecognizeDJLImage(List imageList, OcrRecOptions options) { throw new UnsupportedOperationException("默认不支持该功能"); } diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/PpOCRV5RecModel.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/OcrCommonRecModelImpl.java similarity index 50% rename from smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/PpOCRV5RecModel.java rename to smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/OcrCommonRecModelImpl.java index 423b0a8..775efd0 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/PpOCRV5RecModel.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/OcrCommonRecModelImpl.java @@ -6,35 +6,28 @@ 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.Point; import ai.djl.ndarray.NDArray; -import ai.djl.ndarray.NDList; import ai.djl.ndarray.NDManager; -import ai.djl.opencv.OpenCVImageFactory; 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 ai.djl.training.util.ProgressBar; import cn.smartjavaai.common.enums.DeviceEnum; import cn.smartjavaai.common.pool.PredictorFactory; import cn.smartjavaai.common.utils.FileUtils; import cn.smartjavaai.common.utils.ImageUtils; import cn.smartjavaai.common.utils.OpenCVUtils; -import cn.smartjavaai.ocr.config.DirectionModelConfig; -import cn.smartjavaai.ocr.config.OcrDetModelConfig; import cn.smartjavaai.ocr.config.OcrRecModelConfig; +import cn.smartjavaai.ocr.config.OcrRecOptions; import cn.smartjavaai.ocr.entity.*; import cn.smartjavaai.ocr.exception.OcrException; -import cn.smartjavaai.ocr.factory.OcrModelFactory; import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel; -import cn.smartjavaai.ocr.model.common.detect.translator.PPOCRV5DetTranslator; import cn.smartjavaai.ocr.model.common.direction.OcrDirectionModel; -import cn.smartjavaai.ocr.model.common.recognize.translator.PPOCRV5RecTranslator; -import cn.smartjavaai.ocr.opencv.OcrNDArrayUtils; +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; import org.apache.commons.lang3.StringUtils; import org.apache.commons.pool2.ObjectPool; import org.apache.commons.pool2.impl.GenericObjectPool; @@ -45,50 +38,37 @@ import java.awt.image.BufferedImage; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; -import java.nio.file.Files; -import java.nio.file.Path; import java.nio.file.Paths; import java.util.*; -import java.util.concurrent.ConcurrentHashMap; +import java.util.stream.Collectors; /** * PPOCRV5 识别模型 * @author dwj - * @date 2025/4/21 */ @Slf4j -public class PpOCRV5RecModel implements OcrCommonRecModel { +public class OcrCommonRecModelImpl implements OcrCommonRecModel { private ObjectPool> recPredictorPool; private OcrRecModelConfig config; - private OcrCommonDetModel detModel; + private ZooModel recognitionModel; private OcrDirectionModel directionModel; - private ZooModel recognitionModel; + private OcrCommonDetModel textDetModel; @Override public void loadModel(OcrRecModelConfig config){ if(StringUtils.isBlank(config.getRecModelPath())){ throw new OcrException("recModelPath is null"); } - Device device = null; - if(!Objects.isNull(config.getDevice())){ - device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(); - } this.config = config; + this.directionModel = config.getDirectionModel(); + this.textDetModel = config.getTextDetModel(); //初始化 识别Criteria - Criteria recCriteria = - Criteria.builder() - .optEngine("OnnxRuntime") - .setTypes(Image.class, String.class) - .optModelPath(Paths.get(config.getRecModelPath())) - .optTranslator(new PPOCRV5RecTranslator(new ConcurrentHashMap())) - .optProgress(new ProgressBar()) - .optDevice(device) - .build(); + Criteria recCriteria = OcrCommonRecCriterialFactory.createCriteria(config); try{ recognitionModel = ModelZoo.loadModel(recCriteria); this.recPredictorPool = new GenericObjectPool<>(new PredictorFactory<>(recognitionModel)); @@ -98,29 +78,11 @@ public class PpOCRV5RecModel implements OcrCommonRecModel { throw new OcrException("识别模型加载失败", e); } - - //获取检测模型 - if(StringUtils.isNotBlank(config.getDetModelPath()) && Objects.nonNull(config.getDetModelEnum())){ - OcrDetModelConfig detModelConfig = new OcrDetModelConfig(); - detModelConfig.setModelEnum(config.getDetModelEnum()); - detModelConfig.setDetModelPath(config.getDetModelPath()); - detModelConfig.setDevice(config.getDevice()); - detModel = OcrModelFactory.getInstance().getDetModel(detModelConfig); - } - - //获取方向检测模型 - if(StringUtils.isNotBlank(config.getDirectionModelPath()) && Objects.nonNull(config.getDirectionModelEnum())){ - DirectionModelConfig directionModelConfig = new DirectionModelConfig(); - directionModelConfig.setModelEnum(config.getDirectionModelEnum()); - directionModelConfig.setModelPath(config.getDirectionModelPath()); - directionModelConfig.setDevice(config.getDevice()); - directionModel = OcrModelFactory.getInstance().getDirectionModel(directionModelConfig); - } } @Override - public OcrInfo recognize(String imagePath) { + public OcrInfo recognize(String imagePath, OcrRecOptions options) { if(StringUtils.isBlank(config.getRecModelPath())){ throw new OcrException("recModelPath为空，无法识别"); } @@ -130,78 +92,44 @@ public class PpOCRV5RecModel implements OcrCommonRecModel { Image img = null; try { img = ImageFactory.getInstance().fromFile(Paths.get(imagePath)); + return recognize(img, options); } catch (IOException e) { throw new OcrException("无效的图片", e); - } - OcrInfo ocrInfo = recognize(img); - ((Mat)img.getWrappedImage()).release(); - return ocrInfo; - } - - private OcrInfo recognize(Image image) { - //检测文本 - List boxeList = detModel.detect(image); - if(Objects.isNull(boxeList) || boxeList.isEmpty()){ - throw new OcrException("未检测到文本"); - } - Predictor predictor = null; - List rotatedBoxes = new ArrayList<>(); - List ocrItemList = new ArrayList<>(); - try (NDManager manager = NDManager.newBaseManager()) { - Mat srcMat = (Mat) image.getWrappedImage(); - predictor = recPredictorPool.borrowObject(); - //检测方向 - if(directionModel != null){ - ocrItemList = directionModel.detect(boxeList, srcMat, manager); - if(Objects.isNull(ocrItemList) || ocrItemList.isEmpty()){ - throw new OcrException("方向检测失败"); - } - for (OcrItem ocrItem : ocrItemList){ - //放射变换+裁剪 - Image subImage = OcrUtils.transformAndCrop(srcMat, ocrItem.getOcrBox()); - //ImageUtils.saveImage(subImage, UUID.randomUUID().toString() + "_aaa.png", "build/output"); - //纠正文本框 - subImage = OcrUtils.rotateImg(subImage, ocrItem.getAngle()); - //ImageUtils.saveImage(subImage, UUID.randomUUID().toString() + "_bbb.png", "build/output"); - //识别 - String name = predictor.predict(subImage); - ocrItem.setText(name); - NDArray ndArray = manager.create(ocrItem.getOcrBox().toFloatArray()); - rotatedBoxes.add(new RotatedBox(ndArray, ocrItem.getText())); - ((Mat)subImage.getWrappedImage()).release(); - } - }else{ - for (OcrBox box : boxeList){ - RotatedBox rotatedBox = recognize(box, srcMat, predictor, manager); - rotatedBoxes.add(rotatedBox); - } - } - //后处理 - return postProcessOcrResult(rotatedBoxes); - } catch (Exception e) { - throw new OcrException("OCR检测错误", 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); - } - } + } finally { + if(img != null){ + ((Mat)img.getWrappedImage()).release(); } } } + /** + * + * @param image + * @param options + * @return + */ + @Override + public OcrInfo recognize(Image image, OcrRecOptions options) { + List result = batchRecognizeDJLImage(Collections.singletonList(image), options); + if(CollectionUtils.isEmpty(result)){ + throw new OcrException("OCR识别结果为空"); + } + return result.get(0); + } - private RotatedBox recognize(OcrBox box,Mat srcMat,Predictor recPredictor,NDManager manager){ - try { + /** + * 批量矫正文本框 + * @param boxList + * @param srcMat + * @param manager + * @return + */ + private List batchAlign(List boxList, Mat srcMat,NDManager manager){ + List imageList = new ArrayList<>(boxList.size()); + for (int i = 0; i < boxList.size(); i++) { //透视变换 + 裁剪 - Image subImg = OcrUtils.transformAndCrop(srcMat, box); + Image subImg = OcrUtils.transformAndCrop(srcMat, boxList.get(i)); //ImageUtils.saveImage(subImg, i + "crop.png", "build/output"); //高宽比 > 1.5 if (subImg.getHeight() * 1.0 / subImg.getWidth() > 1.5) { @@ -209,21 +137,63 @@ public class PpOCRV5RecModel implements OcrCommonRecModel { subImg = OcrUtils.rotateImg(manager, subImg); //ImageUtils.saveImage(subImg, i + "rotate.png", "build/output"); } - String name = recPredictor.predict(subImg); - ((Mat)subImg.getWrappedImage()).release(); - NDArray pointsArray = manager.create(box.toFloatArray()); - return new RotatedBox(pointsArray, name); - } catch (Exception e) { - throw new OcrException("OCR检测错误", e); + imageList.add(subImg); } + return imageList; } + /** + * 批量矫正文本框 + * @param itemList + * @param srcMat + * @param manager + * @return + */ + private List batchAlignWithDirection(List itemList, Mat srcMat,NDManager manager){ + List imageList = new ArrayList<>(itemList.size()); + for (OcrItem ocrItem : itemList) { + //放射变换+裁剪 + Image subImage = OcrUtils.transformAndCrop(srcMat, ocrItem.getOcrBox()); + //ImageUtils.saveImage(subImage, UUID.randomUUID().toString() + "_aaa.png", "build/output"); + //纠正文本框 + subImage = OcrUtils.rotateImg(subImage, ocrItem.getAngle()); + imageList.add(subImage); + } + return imageList; + } + + + +// private RotatedBox recognize(OcrBox box,Mat srcMat,Predictor recPredictor,NDManager manager){ +// try { +// //透视变换 + 裁剪 +// Image subImg = OcrUtils.transformAndCrop(srcMat, box); +// //ImageUtils.saveImage(subImg, i + "crop.png", "build/output"); +// //高宽比 > 1.5 +// if (subImg.getHeight() * 1.0 / subImg.getWidth() > 1.5) { +// //旋转图片90度 +// subImg = OcrUtils.rotateImg(manager, subImg); +// //ImageUtils.saveImage(subImg, i + "rotate.png", "build/output"); +// } +// String name = recPredictor.predict(subImg); +// ((Mat)subImg.getWrappedImage()).release(); +// NDArray pointsArray = manager.create(box.toFloatArray()); +// return new RotatedBox(pointsArray, name); +// } catch (Exception e) { +// throw new OcrException("OCR检测错误", e); +// } +// } + /** * 后处理：排序,分行 * @param rotatedBoxes */ - private OcrInfo postProcessOcrResult(List rotatedBoxes){ + private OcrInfo postProcessOcrResult(List rotatedBoxes, OcrRecOptions ocrRecOptions){ + //不分行 + if(!ocrRecOptions.isEnableLineSplit()){ + return OcrUtils.convertRotatedBoxesToOcrItems(rotatedBoxes); + } //Y坐标升序排序 List initList = new ArrayList<>(); for (RotatedBox result : rotatedBoxes) { @@ -257,13 +227,13 @@ public class PpOCRV5RecModel implements OcrCommonRecModel { @Override - public void recognizeAndDraw(String imagePath, String outputPath, int fontSize) { + public void recognizeAndDraw(String imagePath, String outputPath, int fontSize, OcrRecOptions options) { if(!FileUtils.isFileExists(imagePath)){ throw new OcrException("图像文件不存在"); } try { Image img = ImageFactory.getInstance().fromFile(Paths.get(imagePath)); - OcrInfo ocrInfo = recognize(img); + OcrInfo ocrInfo = recognize(img, options); if(Objects.isNull(ocrInfo) || Objects.isNull(ocrInfo.getLineList()) || ocrInfo.getLineList().isEmpty()){ throw new OcrException("未检测到文字"); } @@ -278,36 +248,36 @@ public class PpOCRV5RecModel implements OcrCommonRecModel { } @Override - public OcrInfo recognize(BufferedImage image) { + public OcrInfo recognize(BufferedImage image, OcrRecOptions options) { if(!ImageUtils.isImageValid(image)){ throw new OcrException("图像无效"); } Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image)); - OcrInfo ocrInfo = recognize(img); + OcrInfo ocrInfo = recognize(img, options); ((Mat)img.getWrappedImage()).release(); return ocrInfo; } @Override - public OcrInfo recognize(byte[] imageData) { + public OcrInfo recognize(byte[] imageData, OcrRecOptions options) { if(Objects.isNull(imageData)){ throw new OcrException("图像无效"); } try { BufferedImage image = ImageIO.read(new ByteArrayInputStream(imageData)); - return recognize(image); + return recognize(image, options); } catch (IOException e) { throw new OcrException("错误的图像", e); } } @Override - public BufferedImage recognizeAndDraw(BufferedImage sourceImage, int fontSize) { + public BufferedImage recognizeAndDraw(BufferedImage sourceImage, int fontSize, OcrRecOptions options) { if(!ImageUtils.isImageValid(sourceImage)){ throw new OcrException("图像无效"); } Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(sourceImage)); - OcrInfo ocrInfo = recognize(img); + OcrInfo ocrInfo = recognize(img, options); if(Objects.isNull(ocrInfo) || Objects.isNull(ocrInfo.getLineList()) || ocrInfo.getLineList().isEmpty()){ throw new OcrException("未检测到文字"); } @@ -325,6 +295,154 @@ public class PpOCRV5RecModel implements OcrCommonRecModel { } } + @Override + public List batchRecognize(List imageList, OcrRecOptions options) { + List djlImageList = new ArrayList<>(imageList.size()); + try { + for (BufferedImage bufferedImage : imageList) { + djlImageList.add(ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(bufferedImage))); + } + return batchRecognizeDJLImage(djlImageList, options); + } catch (Exception e) { + throw new OcrException(e); + } finally { + djlImageList.forEach(image -> ((Mat)image.getWrappedImage()).release()); + } + } + + @Override + public List batchRecognizeDJLImage(List imageList, OcrRecOptions options) { + if(Objects.isNull(textDetModel)){ + throw new OcrException("textDetModel is null"); + } + OcrRecOptions ocrRecOptions = options; + if(Objects.isNull(options)){ + ocrRecOptions = new OcrRecOptions(); + } + if(CollectionUtils.isEmpty(imageList)){ + throw new OcrException("imageList is empty"); + } + //检测文本 + List> boxeList = textDetModel.batchDetectDJLImage(imageList); + if(CollectionUtils.isEmpty(boxeList) || boxeList.size() != imageList.size()){ + throw new OcrException("未检测到文本"); + } + Predictor predictor = null; + List ocrInfoList = new ArrayList(); + try (NDManager manager = NDManager.newBaseManager()) { + predictor = recPredictorPool.borrowObject(); + List allImageAlignList = new ArrayList(); + //检测方向 + if(ocrRecOptions.isEnableDirectionCorrect()){ + if(Objects.isNull(directionModel)){ + throw new OcrException("请配置方向模型"); + } + List matList = imageList.stream() + .map(image -> (Mat)image.getWrappedImage()) + .collect(Collectors.toList()); + List> ocrItemList = directionModel.batchDetect(boxeList, matList); + if(CollectionUtils.isEmpty(ocrItemList) || ocrItemList.size() != imageList.size()){ + throw new OcrException("方向检测失败"); + } + allImageAlignList = new ArrayList(); + for (int i = 0; i < ocrItemList.size(); i++) { + Mat srcMat = (Mat) imageList.get(i).getWrappedImage(); + List imageAlignList = batchAlignWithDirection(ocrItemList.get(i), srcMat, manager); +// for(int j = 0; j < imageAlignList.size(); j++){ +// ImageUtils.saveImage(imageAlignList.get(j),"dir-"+i+"-"+j+".png","/Users/xxx/Downloads/testing33"); +// } + allImageAlignList.addAll(imageAlignList); + } + }else{ + for (int i = 0; i < boxeList.size(); i++) { + Mat srcMat = (Mat) imageList.get(i).getWrappedImage(); + List imageAlignList = batchAlign(boxeList.get(i), srcMat, manager); +// for(int j = 0; j < imageAlignList.size(); j++){ +// ImageUtils.saveImage(imageAlignList.get(j),i+"-"+j+".png","/Users/xxx/Downloads/testing33"); +// } + allImageAlignList.addAll(imageAlignList); + } + } + List textList = batchRecognize(allImageAlignList); + int textIndex = 0; + for (int i = 0; i < boxeList.size(); i++) { + List rotatedBoxes = new ArrayList<>(); + for (int j = 0; j < boxeList.get(i).size(); j++){ + if(textIndex >= textList.size()){ + throw new OcrException("识别失败: 第" + i + "张图片, 第" + j + "个文本块，未识别到文本"); + } + OcrBox box = boxeList.get(i).get(j); + NDArray pointsArray = manager.create(box.toFloatArray()); + rotatedBoxes.add(new RotatedBox(pointsArray, textList.get(textIndex))); + textIndex++; + } + OcrInfo ocrInfo = postProcessOcrResult(rotatedBoxes, ocrRecOptions); + ocrInfoList.add(ocrInfo); + } + return ocrInfoList; + } catch (Exception e) { + throw new OcrException("OCR检测错误", 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); + } + } + } + } + } + + private List batchRecognize(List imageAlignList){ + Predictor predictor = null; + try { + predictor = recPredictorPool.borrowObject(); + List textList = predictor.batchPredict(imageAlignList); + imageAlignList.forEach(subImg -> ((Mat)subImg.getWrappedImage()).release()); + return textList; + } catch (Exception e) { + throw new OcrException("OCR检测错误", 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 void setTextDetModel(OcrCommonDetModel detModel) { + this.textDetModel = detModel; + } + + @Override + public OcrCommonDetModel getTextDetModel() { + return textDetModel; + } + + @Override + public void setDirectionModel(OcrDirectionModel directionModel) { + this.directionModel = directionModel; + } + + @Override + public OcrDirectionModel getDirectionModel() { + return directionModel; + } + @Override public void close() throws Exception { try { @@ -334,20 +452,6 @@ public class PpOCRV5RecModel implements OcrCommonRecModel { } catch (Exception e) { log.warn("关闭 predictorPool 失败", e); } - try { - if (detModel != null) { - detModel.close(); - } - } catch (Exception e) { - log.warn("关闭 model 失败", e); - } - try { - if (directionModel != null) { - directionModel.close(); - } - } catch (Exception e) { - log.warn("关闭 model 失败", e); - } try { if (recognitionModel != null) { recognitionModel.close(); 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 new file mode 100644 index 0000000..6eb61e5 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/criteria/OcrCommonRecCriterialFactory.java @@ -0,0 +1,51 @@ +package cn.smartjavaai.ocr.model.common.recognize.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.OcrRecModelConfig; +import cn.smartjavaai.ocr.enums.CommonRecModelEnum; +import cn.smartjavaai.ocr.model.common.recognize.translator.PPOCRRecTranslator; +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 OcrCommonRecCriterialFactory { + + + public static Criteria createCriteria(OcrRecModelConfig config) { + Device device = null; + if(!Objects.isNull(config.getDevice())){ + device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(); + } + Criteria criteria = null; + ConcurrentHashMap params = new ConcurrentHashMap(); + params.putAll(config.getCustomParams()); + if(StringUtils.isNotBlank(config.getBatchifier())){ + params.put("batchifier", config.getBatchifier()); + } + if(config.getRecModelEnum() == CommonRecModelEnum.PP_OCR_V5_SERVER_REC_MODEL || + config.getRecModelEnum() == CommonRecModelEnum.PP_OCR_V5_MOBILE_REC_MODEL || + config.getRecModelEnum() == CommonRecModelEnum.PP_OCR_V4_SERVER_REC_MODEL || + config.getRecModelEnum() == CommonRecModelEnum.PP_OCR_V4_MOBILE_REC_MODEL ){ + criteria = + Criteria.builder() + .optEngine("OnnxRuntime") + .setTypes(Image.class, String.class) + .optModelPath(Paths.get(config.getRecModelPath())) + .optTranslator(new PPOCRRecTranslator(params)) + .optProgress(new ProgressBar()) + .optDevice(device) + .build(); + } + return criteria; + } +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/translator/PPOCRV5RecTranslator.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/translator/PPOCRRecTranslator.java similarity index 87% rename from smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/translator/PPOCRV5RecTranslator.java rename to smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/translator/PPOCRRecTranslator.java index dab81d9..5f2df56 100644 --- a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/translator/PPOCRV5RecTranslator.java +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/common/recognize/translator/PPOCRRecTranslator.java @@ -23,15 +23,20 @@ import java.util.Map; * 文字识别前后处理 * */ -public class PPOCRV5RecTranslator implements Translator { +public class PPOCRRecTranslator implements Translator { private List table; private final boolean use_space_char; - public PPOCRV5RecTranslator(Map arguments) { + private String batchifier; + + public PPOCRRecTranslator(Map arguments) { use_space_char = arguments.containsKey("use_space_char") ? Boolean.parseBoolean(arguments.get("use_space_char").toString()) : true; + batchifier = arguments.containsKey("batchifier") + ? arguments.get("batchifier").toString() + : "padding"; } @Override @@ -57,7 +62,8 @@ public class PPOCRV5RecTranslator implements Translator { StringBuilder sb = new StringBuilder(); NDArray tokens = list.singletonOrThrow(); - long[] indices = tokens.get(0).argMax(1).toLongArray(); +// long[] indices = tokens.get(0).argMax(1).toLongArray(); + long[] indices = tokens.argMax(1).toLongArray(); boolean[] selection = new boolean[indices.length]; Arrays.fill(selection, true); for (int i = 1; i < indices.length; i++) { @@ -111,13 +117,13 @@ public class PPOCRV5RecTranslator implements Translator { padding_im.set(new NDIndex(":,:,0:" + resized_w), resized_image); padding_im = padding_im.flip(0); - padding_im = padding_im.expandDims(0); +// padding_im = padding_im.expandDims(0); return new NDList(padding_im); } @Override public Batchifier getBatchifier() { - return null; + return Batchifier.fromString(batchifier); } } 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 new file mode 100644 index 0000000..a8a506b --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/CommonTableStructureModel.java @@ -0,0 +1,160 @@ +package cn.smartjavaai.ocr.model.table; + +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.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.FileUtils; +import cn.smartjavaai.common.utils.ImageUtils; +import cn.smartjavaai.common.utils.OpenCVUtils; +import cn.smartjavaai.ocr.config.OcrDetModelConfig; +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 cn.smartjavaai.ocr.exception.OcrException; +import cn.smartjavaai.ocr.model.table.criteria.StructureCriteriaFactory; +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.IOException; +import java.nio.file.Paths; +import java.util.List; +import java.util.Objects; + +/** + * 表格结构模型 + * @author dwj + */ +@Slf4j +public class CommonTableStructureModel implements TableStructureModel{ + + private ZooModel model; + + private ObjectPool> predictorPool; + + @Override + public void loadModel(TableStructureConfig config) { + if(StringUtils.isBlank(config.getModelPath())){ + throw new OcrException("modelPath is null"); + } + Criteria criteria = StructureCriteriaFactory.createCriteria(config); + try{ + model = ModelZoo.loadModel(criteria); + // 创建池子：每个线程独享 Predictor + this.predictorPool = new GenericObjectPool<>(new PredictorFactory<>(model)); + log.debug("当前设备: " + model.getNDManager().getDevice()); + log.debug("当前引擎: " + Engine.getInstance().getEngineName()); + } catch (IOException | ModelNotFoundException | MalformedModelException e) { + throw new OcrException("表格结构识别模型加载失败", e); + } + } + + @Override + public R detect(BufferedImage image) { + if(!ImageUtils.isImageValid(image)){ + return R.fail(R.Status.INVALID_IMAGE); + } + Image img = null; + try { + img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image)); + return detect(img); + } catch (Exception e) { + throw new OcrException(e); + } finally { + if(Objects.nonNull(img)){ + ((Mat)img.getWrappedImage()).release(); + } + } + } + + @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)); + return detect(img); + } catch (IOException e) { + throw new OcrException("无效的图片", e); + } finally { + if (Objects.nonNull(img)){ + ((Mat)img.getWrappedImage()).release(); + } + } + } + + @Override + public R detect(byte[] imageData) { + if(Objects.isNull(imageData)){ + return R.fail(R.Status.INVALID_IMAGE); + } + try { + BufferedImage image = ImageIO.read(new ByteArrayInputStream(imageData)); + return detect(image); + } catch (IOException e) { + throw new OcrException("错误的图像", e); + } + } + + @Override + public R detect(Image image) { + Predictor predictor = null; + try { + predictor = predictorPool.borrowObject(); + TableStructureResult result = predictor.predict(image); + return R.ok(result); + } catch (Exception e) { + throw new OcrException("OCR检测错误", e); + }finally { + if (predictor != null) { + try { + predictorPool.returnObject(predictor); //归还 + } catch (Exception e) { + log.warn("归还Predictor失败", e); + try { + predictor.close(); // 归还失败才销毁 + } catch (Exception ex) { + log.error("关闭Predictor失败", ex); + } + } + } + } + } + + @Override + public void close() throws Exception { + try { + if (predictorPool != null) { + predictorPool.close(); + } + } catch (Exception e) { + log.warn("关闭 predictorPool 失败", e); + } + try { + if (model != null) { + model.close(); + } + } catch (Exception e) { + log.warn("关闭 model 失败", e); + } + } +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/TableRecognizer.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/TableRecognizer.java new file mode 100644 index 0000000..faf76a6 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/TableRecognizer.java @@ -0,0 +1,454 @@ +package cn.smartjavaai.ocr.model.table; + +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.output.Rectangle; +import ai.djl.translate.TranslateException; +import ai.djl.util.JsonUtils; +import cn.smartjavaai.common.entity.DetectionRectangle; +import cn.smartjavaai.common.entity.R; +import cn.smartjavaai.common.utils.FileUtils; +import cn.smartjavaai.common.utils.ImageUtils; +import cn.smartjavaai.common.utils.OpenCVUtils; +import cn.smartjavaai.ocr.config.OcrRecModelConfig; +import cn.smartjavaai.ocr.config.OcrRecOptions; +import cn.smartjavaai.ocr.entity.OcrBox; +import cn.smartjavaai.ocr.entity.OcrInfo; +import cn.smartjavaai.ocr.entity.OcrItem; +import cn.smartjavaai.ocr.entity.TableStructureResult; +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.OcrCommonRecModel; +import cn.smartjavaai.ocr.utils.ConvertHtml2Excel; +import lombok.extern.slf4j.Slf4j; +import org.apache.commons.collections.CollectionUtils; +import org.apache.commons.lang3.StringUtils; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.poi.hssf.usermodel.HSSFWorkbook; +import org.opencv.core.Mat; + +import javax.imageio.ImageIO; +import java.awt.*; +import java.awt.image.BufferedImage; +import java.io.ByteArrayInputStream; +import java.io.File; +import java.io.IOException; +import java.nio.file.Paths; +import java.util.*; +import java.util.List; +import java.util.concurrent.ConcurrentHashMap; + +/** + * 表格内容识别器 + * @author dwj + */ +@Slf4j +public class TableRecognizer { + + private OcrCommonDetModel textDetector; + private TableStructureModel tableStructureModel; + private OcrCommonRecModel textRecModel; + private OcrDirectionModel directionModel; + + private TableRecognizer(Builder builder) { + this.tableStructureModel = builder.tableStructureModel; + this.textRecModel = builder.textRecModel; + this.directionModel = builder.directionModel; + this.textDetector = builder.textDetector; + textRecModel.setTextDetModel(textDetector); + textRecModel.setDirectionModel(directionModel); + } + + public static Builder builder() { + return new Builder(); + } + + // 链式设置文本识别模型 + public TableRecognizer withTextRecModel(OcrCommonRecModel textRecModel) { + this.textRecModel = textRecModel; + return this; + } + + // 链式设置表格结构模型 + public TableRecognizer withStructureModel(TableStructureModel tableStructureModel) { + this.tableStructureModel = tableStructureModel; + return this; + } + + /** + * 表格识别 + * @param image + * @return + */ + public R recognize(Image image) { + //表格结构识别 + R result = tableStructureModel.detect(image); + if(!result.isSuccess()){ + return R.fail(result.getCode(), result.getMessage()); + } + //文本检测+文字识别 + boolean enableDirectionCorrect = directionModel == null ? false : true; + OcrRecOptions options = new OcrRecOptions(enableDirectionCorrect, false); + OcrInfo ocrInfo = textRecModel.recognize(image, options); + List tableContentList = buildTable(result.getData(), ocrInfo); + String html = convertHtml(result.getData().getTableTagList(), tableContentList); + result.getData().setHtml(html); + return result; + } + + + /** + * 表格识别 + * @param image + * @return + */ + public R recognize(BufferedImage image) { + if(!ImageUtils.isImageValid(image)){ + return R.fail(R.Status.INVALID_IMAGE); + } + Image img = null; + try { + img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image)); + return recognize(img); + } catch (Exception e) { + throw new OcrException(e); + } finally { + if(Objects.nonNull(img)){ + ((Mat)img.getWrappedImage()).release(); + } + } + } + + /** + * 表格识别 + * @param imagePath + * @return + */ + 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)); + return recognize(img); + } catch (IOException e) { + throw new OcrException("无效的图片", e); + } finally { + if (Objects.nonNull(img)){ + ((Mat)img.getWrappedImage()).release(); + } + } + } + + + /** + * 表格识别 + * @param imageData + * @return + */ + public R recognize(byte[] imageData) { + if(Objects.isNull(imageData)){ + return R.fail(R.Status.INVALID_IMAGE); + } + try { + BufferedImage image = ImageIO.read(new ByteArrayInputStream(imageData)); + return recognize(image); + } catch (IOException e) { + throw new OcrException("错误的图像", e); + } + } + + /** + * 绘制表格 + * @param tableStructureResult + * @param image + * @param savePath + */ + public void drawTable(TableStructureResult tableStructureResult, BufferedImage image, String savePath){ + if(Objects.isNull(tableStructureResult) || CollectionUtils.isEmpty(tableStructureResult.getTableTagList())){ + throw new OcrException("表格结构为空"); + } + for (int i = 0; i < tableStructureResult.getOcrItemList().size(); i++){ + OcrItem item = tableStructureResult.getOcrItemList().get(i); + DetectionRectangle detectionRectangle = item.getOcrBox().toDetectionRectangle(); + ImageUtils.drawImageRectWithText(image, detectionRectangle, i + "", Color.RED); + } + ImageUtils.saveImage(image, savePath); + } + + /** + * 删除 HTML 中第一个 段落 + * @param html 原始 HTML + * @return 去掉 ", styleStart); + if (styleEnd == -1) { + return html; // 没闭合标签，不处理 + } + styleEnd += "".length(); + // 去掉 style 块 + return html.substring(0, styleStart) + html.substring(styleEnd); + } + + /** + * 导出 Excel + * @param html + * @param savePath + */ + public void exportExcel(String html, String savePath){ + try { + String content = removeStyleBlock(html); + content = content.replace("", ""); + content = content.replace("", ""); + HSSFWorkbook workbook = ConvertHtml2Excel.table2Excel(content); + workbook.write(new File(savePath)); + } catch (Exception e) { + throw new OcrException("导出excel失败，请检查表结构是否识别正确"); + } + } + + + /** + * 构建表格 + * @param tableStructureResult + * @param ocrInfo + * @return + */ + public List buildTable(TableStructureResult tableStructureResult, OcrInfo ocrInfo) { + // 获取 Cell 与 文本检测框 的对应关系(1:N)。 + Map> matched = new ConcurrentHashMap<>(); + List ocrItems = ocrInfo.getOcrItemList(); + + for (int i = 0; i < ocrItems.size(); i++) { + OcrBox ocrBox = ocrItems.get(i).getOcrBox(); + int[] box_1 = { + (int)ocrBox.getTopLeft().getX(), + (int)ocrBox.getTopLeft().getY(), + (int)ocrBox.getBottomRight().getX(), + (int)ocrBox.getBottomRight().getY() + }; + // 获取两两cell之间的L1距离和 1- IOU + List> distances = new ArrayList<>(); + for (OcrItem cell : tableStructureResult.getOcrItemList()) { + OcrBox cellBox = cell.getOcrBox(); + int[] box_2 = { + (int)cellBox.getTopLeft().getX(), + (int)cellBox.getTopLeft().getY(), + (int)cellBox.getBottomRight().getX(), + (int)cellBox.getBottomRight().getY() + }; + float distance = distance(box_1, box_2); + float iou = 1 - computeIou(box_1, box_2); + distances.add(Pair.of(distance, iou)); + } + // 根据距离和IOU挑选最"近"的cell + Pair nearest = sorted(distances); + + // 获取最小距离对应的下标id，也等价于cell的下标id （distances列表是根据遍历cells生成的） + int id = 0; + for (int idx = 0; idx < distances.size(); idx++) { + Pair current = distances.get(idx); + if (current.getLeft().floatValue() == nearest.getLeft().floatValue() + && current.getRight().floatValue() == nearest.getRight().floatValue()) { + id = idx; + break; + } + } + if (!matched.containsKey(id)) { + List textIds = new ArrayList<>(); + textIds.add(i); + // cell id, text id list (dt_boxes index list) + matched.put(id, textIds); + } else { + matched.get(id).add(i); + } + } + + List cell_contents = new ArrayList<>(); + List probs = new ArrayList<>(); + for (int i = 0; i < tableStructureResult.getOcrItemList().size(); i++) { + List textIds = matched.get(i); + List contents = new ArrayList<>(); + String content = ""; + if (textIds != null) { + for (Integer id : textIds) { + contents.add(ocrItems.get(id).getText()); + } + content = StringUtils.join(contents, " "); + } + cell_contents.add(content); + probs.add(-1.0); + } + return cell_contents; + } + + /** + * 计算欧式距离 + * Calculate L1 distance + * + * @param box_1 + * @param box_2 + * @return + */ + private int distance(int[] box_1, int[] box_2) { + int x1 = box_1[0]; + int y1 = box_1[1]; + int x2 = box_1[2]; + int y2 = box_1[3]; + int x3 = box_2[0]; + int y3 = box_2[1]; + int x4 = box_2[2]; + int y4 = box_2[3]; + int dis = Math.abs(x3 - x1) + Math.abs(y3 - y1) + Math.abs(x4 - x2) + Math.abs(y4 - y2); + int dis_2 = Math.abs(x3 - x1) + Math.abs(y3 - y1); + int dis_3 = Math.abs(x4 - x2) + Math.abs(y4 - y2); + return dis + Math.min(dis_2, dis_3); + } + + /** + * 计算交并比 + * computing IoU + * + * @param rec1: (y0, x0, y1, x1), which reflects (top, left, bottom, right) + * @param rec2: (y0, x0, y1, x1) + * @return scala value of IoU + */ + private float computeIou(int[] rec1, int[] rec2) { + // computing area of each rectangles + int S_rec1 = (rec1[2] - rec1[0]) * (rec1[3] - rec1[1]); + int S_rec2 = (rec2[2] - rec2[0]) * (rec2[3] - rec2[1]); + + // computing the sum_area + int sum_area = S_rec1 + S_rec2; + + // find the each edge of intersect rectangle + int left_line = Math.max(rec1[1], rec2[1]); + int right_line = Math.min(rec1[3], rec2[3]); + int top_line = Math.max(rec1[0], rec2[0]); + int bottom_line = Math.min(rec1[2], rec2[2]); + + // judge if there is an intersect + if (left_line >= right_line || top_line >= bottom_line) { + return 0.0f; + } else { + float intersect = (right_line - left_line) * (bottom_line - top_line); + return (intersect / (sum_area - intersect)) * 1.0f; + } + } + + /** + * 距离排序 + * Distance sorted + * + * @param distances + * @return + */ + private Pair sorted(List> distances) { + Comparator> comparator = + new Comparator>() { + @Override + public int compare(Pair a1, Pair a2) { + // 首先根据IoU排序 + if (a1.getRight().floatValue() > a2.getRight().floatValue()) { + return 1; + } else if (a1.getRight().floatValue() == a2.getRight().floatValue()) { + // 然后根据L1距离排序 + if (a1.getLeft().floatValue() > a2.getLeft().floatValue()) { + return 1; + } + return -1; + } + return -1; + } + }; + + // 距离排序 + List> newDistances = new ArrayList<>(); + CollectionUtils.addAll(newDistances, new Object[distances.size()]); + Collections.copy(newDistances, distances); + Collections.sort(newDistances, comparator); + return newDistances.get(0); + } + + /** + * 生成表格html + * Generate table html + * + * @param pred_structures + * @param cell_contents + * @return + */ + public String convertHtml(List pred_structures, List cell_contents) { + StringBuffer html = new StringBuffer(); + // 添加统一的样式（可选放到中） + html.append("\n"); + int td_index = 0; + for (String tag : pred_structures) { + if (tag.contains("")) { + String content = cell_contents.get(td_index); + html.append(""); + html.append(content); + html.append(""); + td_index++; + continue; + } + html.append(tag); + } + return html.toString(); + } + + + public static class Builder { + private TableStructureModel tableStructureModel; + private OcrCommonRecModel textRecModel; + private OcrDirectionModel directionModel; + private OcrCommonDetModel textDetector; + + public Builder withStructureModel(TableStructureModel model) { + this.tableStructureModel = model; + return this; + } + + public Builder withTextRecModel(OcrCommonRecModel model) { + this.textRecModel = model; + return this; + } + + public Builder withDirectionModel(OcrDirectionModel model) { + this.directionModel = model; + return this; + } + + public Builder withTextDetModel(OcrCommonDetModel model) { + this.textDetector = model; + return this; + } + + public TableRecognizer build() { + if (this.tableStructureModel == null) { + throw new IllegalStateException("tableStructureModel 未设置"); + } + if (this.textDetector == null) { + throw new IllegalStateException("textDetector 未设置"); + } + if (this.textRecModel == null) { + throw new IllegalStateException("textRecModel 未设置"); + } + return new TableRecognizer(this); + } + } + +} 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 new file mode 100644 index 0000000..a1ea4e3 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/TableStructureModel.java @@ -0,0 +1,63 @@ +package cn.smartjavaai.ocr.model.table; + +import ai.djl.modality.cv.Image; +import cn.smartjavaai.common.entity.R; +import cn.smartjavaai.ocr.config.OcrDetModelConfig; +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 java.awt.image.BufferedImage; +import java.util.List; + +/** + * 表格结构识别模型 + * @author dwj + */ +public interface TableStructureModel extends AutoCloseable{ + + /** + * 加载模型 + * @param config + */ + void loadModel(TableStructureConfig config); + + + /** + * 表格结构检测 + * @param image + * @return + */ + default R detect(BufferedImage image){ + throw new UnsupportedOperationException("默认不支持该功能"); + } + + /** + * 表格结构检测 + * @param imagePath 图片路径 + * @return + */ + default R detect(String imagePath) { + throw new UnsupportedOperationException("默认不支持该功能"); + } + + /** + * 表格结构检测 + * @param imageData 图片字节数组 + * @return + */ + default R detect(byte[] imageData) { + throw new UnsupportedOperationException("默认不支持该功能"); + } + + + /** + * 表格结构检测 + * @param image DJL Image + * @return + */ + default R detect(Image image){ + throw new UnsupportedOperationException("默认不支持该功能"); + } +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/criteria/StructureCriteriaFactory.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/criteria/StructureCriteriaFactory.java new file mode 100644 index 0000000..b54afe8 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/criteria/StructureCriteriaFactory.java @@ -0,0 +1,60 @@ +package cn.smartjavaai.ocr.model.table.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.TableStructureConfig; +import cn.smartjavaai.ocr.entity.OcrItem; +import cn.smartjavaai.ocr.entity.TableStructureResult; +import cn.smartjavaai.ocr.enums.TableStructureModelEnum; +import cn.smartjavaai.ocr.model.table.translator.TableStructTranslator; + +import java.nio.file.Paths; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; + +/** + * @author dwj + * @date 2025/7/10 + */ +public class StructureCriteriaFactory { + + + public static Criteria createCriteria(TableStructureConfig 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() == TableStructureModelEnum.SLANET){ + criteria = + Criteria.builder() + .optEngine("OnnxRuntime") + .setTypes(Image.class, TableStructureResult.class) + .optModelPath(Paths.get(config.getModelPath())) + .optOption("removePass", "repeated_fc_relu_fuse_pass") + .optDevice(device) + .optTranslator(new TableStructTranslator()) + .optProgress(new ProgressBar()) + .build(); + }else if(config.getModelEnum() == TableStructureModelEnum.SLANET_PLUS){ + criteria = + Criteria.builder() + .optEngine("OnnxRuntime") + .setTypes(Image.class, TableStructureResult.class) + .optModelPath(Paths.get(config.getModelPath())) + .optOption("removePass", "repeated_fc_relu_fuse_pass") + .optDevice(device) + .optTranslator(new TableStructTranslator()) + .optProgress(new ProgressBar()) + .build(); + } + return criteria; + } + + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/translator/TableStructTranslator.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/translator/TableStructTranslator.java new file mode 100644 index 0000000..08041ef --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/model/table/translator/TableStructTranslator.java @@ -0,0 +1,198 @@ +package cn.smartjavaai.ocr.model.table.translator; + +import ai.djl.Model; +import ai.djl.modality.cv.Image; +import ai.djl.modality.cv.output.BoundingBox; +import ai.djl.modality.cv.output.Rectangle; +import ai.djl.modality.cv.util.NDImageUtils; +import ai.djl.ndarray.NDArray; +import ai.djl.ndarray.NDList; +import ai.djl.ndarray.index.NDIndex; +import ai.djl.ndarray.types.DataType; +import ai.djl.ndarray.types.Shape; +import ai.djl.translate.Batchifier; +import ai.djl.translate.Translator; +import ai.djl.translate.TranslatorContext; +import ai.djl.util.Utils; +import cn.smartjavaai.common.entity.Point; +import cn.smartjavaai.ocr.entity.OcrBox; +import cn.smartjavaai.ocr.entity.OcrItem; +import cn.smartjavaai.ocr.entity.TableStructureResult; + +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.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +/** + * 表格识别的前后处理 + */ +public class TableStructTranslator implements Translator { + + private final int maxLength = 488; + private int height; + private int width; + private float scale = 1.0f; + private float xScale; + private float yScale; + private List dict; + private String beg_str = "sos"; + private String end_str = "eos"; + private List td_token = new ArrayList<>(); + + @Override + public void prepare(TranslatorContext ctx) throws IOException { + Model model = ctx.getModel(); + try (InputStream is = model.getArtifact("table_structure_dict_ch.txt").openStream()) { + dict = Utils.readLines(is, false); + dict.add(0,beg_str); + if(dict.contains("")) + dict.remove(""); + if(!dict.contains("")) + dict.add(""); + dict.add(end_str); + } + + td_token.add(""); + td_token.add(""); + } + + @Override + public NDList processInput(TranslatorContext ctx, Image input) { + NDArray img = input.toNDArray(ctx.getNDManager(), Image.Flag.COLOR); + height = input.getHeight(); + width = input.getWidth(); + + img = ResizeTableImage(img, height, width, maxLength); + img = PaddingTableImage(ctx, img, maxLength); + + img = img.transpose(2, 0, 1).div(255).flip(0); + img = NDImageUtils.normalize( + img, new float[]{0.485f, 0.456f, 0.406f}, new float[]{0.229f, 0.224f, 0.225f}); + img = img.expandDims(0); + return new NDList(img); + } + + @Override + public TableStructureResult processOutput(TranslatorContext ctx, NDList list) { + NDArray bbox_preds = list.get(0); + NDArray structure_probs = list.get(1); + + NDArray structure_idx = structure_probs.argMax(2); + structure_probs = structure_probs.max(new int[]{2}); + + List> structure_batch_list = new ArrayList<>(); + List> bbox_batch_list = new ArrayList<>(); + List> result_score_list = new ArrayList<>(); + + // get ignored tokens + int beg_idx = dict.indexOf(beg_str); + int end_idx = dict.indexOf(end_str); + + long batch_size = structure_idx.size(0); + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + List structure_list = new ArrayList<>(); + List bbox_list = new ArrayList<>(); + List score_list = new ArrayList<>(); + + long len = structure_idx.get(batch_idx).size(); + for (int idx = 0; idx < len; idx++) { + int char_idx = (int) structure_idx.get(batch_idx).get(idx).toLongArray()[0]; + if (idx > 0 && char_idx == end_idx) { + break; + } +// if (char_idx == beg_idx || char_idx == end_idx) { +// continue; +// } + String text = dict.get(char_idx); + if(td_token.indexOf(text)>-1){ + NDArray bbox = bbox_preds.get(batch_idx, idx); +// bbox.set(new NDIndex("0::2"), bbox.get(new NDIndex("0::2"))); +// bbox.set(new NDIndex("1::2"), bbox.get(new NDIndex("1::2"))); + bbox_list.add(bbox); + } + structure_list.add(text); + score_list.add(structure_probs.get(batch_idx, idx)); + } + + structure_batch_list.add(structure_list); // structure_str + bbox_batch_list.add(bbox_list); + result_score_list.add(score_list); + } + List structure_str_list =structure_batch_list.get(0); + List bbox_list = bbox_batch_list.get(0); + List score_list = result_score_list.get(0); + + structure_str_list.add(0,""); + structure_str_list.add(1,""); + structure_str_list.add(2,""); + structure_str_list.add("
"); + structure_str_list.add(""); + structure_str_list.add(""); + + List ocrItemList = new ArrayList<>(); + + for (int i = 0; i < bbox_list.size(); i++) { + NDArray box = bbox_list.get(i); + float[] arr = new float[4]; + arr[0] = box.get(new NDIndex("0::2")).min().toFloatArray()[0]; + arr[1] = box.get(new NDIndex("1::2")).min().toFloatArray()[0]; + arr[2] = box.get(new NDIndex("0::2")).max().toFloatArray()[0]; + arr[3] = box.get(new NDIndex("1::2")).max().toFloatArray()[0]; + + Point topLeft = new Point(arr[0] * xScale * width, arr[1] * yScale * height); + Point topRight = new Point(arr[2] * xScale * width, arr[1] * yScale * height); + Point bottomRight = new Point(arr[2] * xScale * width, arr[3] * yScale * height); + Point bottomLeft = new Point(arr[0] * xScale * width, arr[3] * yScale * height); + + + OcrBox ocrBox = new OcrBox(topLeft, topRight, bottomRight, bottomLeft); + //String tag = structure_str_list.get(i + 3); // 前面加了 所以偏移+3 + float score = score_list.get(i).toFloatArray()[0]; // 获取每个结构token的得分 + OcrItem item = new OcrItem(); + item.setOcrBox(ocrBox); + item.setScore(score); + //item.setTableTag(tag); + ocrItemList.add(item); + } + return new TableStructureResult(ocrItemList, structure_str_list); + } + + @Override + public Batchifier getBatchifier() { + return null; + } + + private NDArray ResizeTableImage(NDArray img, int height, int width, int maxLen) { + int localMax = Math.max(height, width); + float ratio = maxLen * 1.0f / localMax; + int resize_h = (int) (height * ratio); + int resize_w = (int) (width * ratio); + scale = ratio; + + if(width > height){ + xScale = 1f; + yScale = (float)width /(float)height; + } else{ + xScale = (float)height /(float)width; + yScale = 1f; + } + + img = NDImageUtils.resize(img, resize_w, resize_h); + return img; + } + + private NDArray PaddingTableImage(TranslatorContext ctx, NDArray img, int maxLen) { + NDArray paddingImg = ctx.getNDManager().zeros(new Shape(maxLen, maxLen, 3), DataType.UINT8); + paddingImg.set( + new NDIndex("0:" + img.getShape().get(0) + ",0:" + img.getShape().get(1) + ",:"), img); + return paddingImg; + } + +} diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/utils/ConvertHtml2Excel.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/utils/ConvertHtml2Excel.java new file mode 100644 index 0000000..e868545 --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/utils/ConvertHtml2Excel.java @@ -0,0 +1,233 @@ +package cn.smartjavaai.ocr.utils; + +import org.apache.commons.lang3.StringUtils; +import org.apache.commons.lang3.math.NumberUtils; +import org.apache.poi.hssf.usermodel.*; +import org.apache.poi.ss.usermodel.*; +import org.apache.poi.ss.util.CellRangeAddress; +import org.dom4j.Document; +import org.dom4j.DocumentException; +import org.dom4j.DocumentHelper; +import org.dom4j.Element; + +import java.util.ArrayList; +import java.util.List; + +/** + * @Auther: xiaoqiang + * @Date: 2020/12/9 9:16 + * @Description: + */ +public class ConvertHtml2Excel { + + /** + * html表格转excel + * + * @param tableHtml 如 + *
+ * .. + *
+ * @return + */ + public static HSSFWorkbook table2Excel(String tableHtml) { + HSSFWorkbook wb = new HSSFWorkbook(); + HSSFSheet sheet = wb.createSheet(); + List crossRowEleMetaLs = new ArrayList<>(); + int rowIndex = 0; + try { + Document data = DocumentHelper.parseText(tableHtml); + // 生成表头 + Element thead = data.getRootElement().element("thead"); + HSSFCellStyle titleStyle = getTitleStyle(wb); + int ls=0;//列数 + if (thead != null) { + List trLs = thead.elements("tr"); + for (Element trEle : trLs) { + HSSFRow row = sheet.createRow(rowIndex); + List thLs = trEle.elements("td"); + ls=thLs.size(); + makeRowCell(thLs, rowIndex, row, 0, titleStyle, crossRowEleMetaLs); + rowIndex++; + } + } + // 生成表体 + Element tbody = data.getRootElement().element("tbody"); + HSSFCellStyle contentStyle = getContentStyle(wb); + if (tbody != null) { + List trLs = tbody.elements("tr"); + for (Element trEle : trLs) { + HSSFRow row = sheet.createRow(rowIndex); + List thLs = trEle.elements("th"); + int cellIndex = makeRowCell(thLs, rowIndex, row, 0, titleStyle, crossRowEleMetaLs); + List tdLs = trEle.elements("td"); + makeRowCell(tdLs, rowIndex, row, cellIndex, contentStyle, crossRowEleMetaLs); + rowIndex++; + } + } + // 合并表头 + for (CrossRangeCellMeta crcm : crossRowEleMetaLs) { + sheet.addMergedRegion(new CellRangeAddress(crcm.getFirstRow(), crcm.getLastRow(), crcm.getFirstCol(), crcm.getLastCol())); + setRegionStyle(sheet, new CellRangeAddress(crcm.getFirstRow(), crcm.getLastRow(), crcm.getFirstCol(), crcm.getLastCol()),titleStyle); + } + for(int i=0;i tdLs, int rowIndex, HSSFRow row, int startCellIndex, HSSFCellStyle cellStyle, + List crossRowEleMetaLs) { + int i = startCellIndex; + for (int eleIndex = 0; eleIndex < tdLs.size(); i++, eleIndex++) { + int captureCellSize = getCaptureCellSize(rowIndex, i, crossRowEleMetaLs); + while (captureCellSize > 0) { + for (int j = 0; j < captureCellSize; j++) {// 当前行跨列处理（补单元格） + row.createCell(i); + i++; + } + captureCellSize = getCaptureCellSize(rowIndex, i, crossRowEleMetaLs); + } + Element thEle = tdLs.get(eleIndex); + String val = thEle.getTextTrim(); + if (StringUtils.isBlank(val)) { + Element e = thEle.element("a"); + if (e != null) { + val = e.getTextTrim(); + } + } + HSSFCell c = row.createCell(i); + if (NumberUtils.isNumber(val)) { + c.setCellValue(Double.parseDouble(val)); + c.setCellType(CellType.NUMERIC); + } else { + c.setCellValue(val); + } + int rowSpan = NumberUtils.toInt(thEle.attributeValue("rowspan"), 1); + int colSpan = NumberUtils.toInt(thEle.attributeValue("colspan"), 1); + c.setCellStyle(cellStyle); + if (rowSpan > 1 || colSpan > 1) { // 存在跨行或跨列 + crossRowEleMetaLs.add(new CrossRangeCellMeta(rowIndex, i, rowSpan, colSpan)); + } + if (colSpan > 1) {// 当前行跨列处理（补单元格） + for (int j = 1; j < colSpan; j++) { + i++; + row.createCell(i); + } + } + } + return i; + } + + /** + * 设置合并单元格的边框样式 + * + * @param sheet + * @param region + * @param cs + */ + public static void setRegionStyle(HSSFSheet sheet, CellRangeAddress region, HSSFCellStyle cs) { + for (int i = region.getFirstRow(); i <= region.getLastRow(); i++) { + HSSFRow row = sheet.getRow(i); + for (int j = region.getFirstColumn(); j <= region.getLastColumn(); j++) { + HSSFCell cell = row.getCell(j); + cell.setCellStyle(cs); + } + } + } + + /** + * 获得因rowSpan占据的单元格 + * + * @param rowIndex 行号 + * @param colIndex 列号 + * @param crossRowEleMetaLs 跨行列元数据 + * @return 当前行在某列需要占据单元格 + */ + private static int getCaptureCellSize(int rowIndex, int colIndex, List crossRowEleMetaLs) { + int captureCellSize = 0; + for (CrossRangeCellMeta crossRangeCellMeta : crossRowEleMetaLs) { + if (crossRangeCellMeta.getFirstRow() < rowIndex && crossRangeCellMeta.getLastRow() >= rowIndex) { + if (crossRangeCellMeta.getFirstCol() <= colIndex && crossRangeCellMeta.getLastCol() >= colIndex) { + captureCellSize = crossRangeCellMeta.getLastCol() - colIndex + 1; + } + } + } + return captureCellSize; + } + + /** + * 获得标题样式 + * + * @param workbook + * @return + */ + private static HSSFCellStyle getTitleStyle(HSSFWorkbook workbook) { + //short titlebackgroundcolor = IndexedColors.GREY_25_PERCENT.index; + short fontSize = 12; + String fontName = "宋体"; + HSSFCellStyle style = workbook.createCellStyle(); + style.setVerticalAlignment(VerticalAlignment.CENTER); + style.setAlignment(HorizontalAlignment.CENTER); + style.setBorderBottom(BorderStyle.THIN); //下边框 + style.setBorderLeft(BorderStyle.THIN);//左边框 + style.setBorderTop(BorderStyle.THIN);//上边框 + style.setBorderRight(BorderStyle.THIN);//右边框 + //style.setFillPattern(FillPatternType.SOLID_FOREGROUND); + //style.setFillForegroundColor(titlebackgroundcolor);// 背景色 + + HSSFFont font = workbook.createFont(); + font.setFontName(fontName); + font.setFontHeightInPoints(fontSize); + font.setBold(true); + style.setFont(font); + return style; + } + + /** + * 获得内容样式 + * + * @param wb + * @return + */ + private static HSSFCellStyle getContentStyle(HSSFWorkbook wb) { + short fontSize = 12; + String fontName = "宋体"; + HSSFCellStyle style = wb.createCellStyle(); + style.setBorderBottom(BorderStyle.THIN); //下边框 + style.setBorderLeft(BorderStyle.THIN);//左边框 + style.setBorderTop(BorderStyle.THIN);//上边框 + style.setBorderRight(BorderStyle.THIN);//右边框 + HSSFFont font = wb.createFont(); + font.setFontName(fontName); + font.setFontHeightInPoints(fontSize); + style.setFont(font); + style.setAlignment(HorizontalAlignment.CENTER);//水平居中 + style.setVerticalAlignment(VerticalAlignment.CENTER);//垂直居中 + style.setWrapText(true); + return style; + } +} + diff --git a/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/utils/CrossRangeCellMeta.java b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/utils/CrossRangeCellMeta.java new file mode 100644 index 0000000..5e3150c --- /dev/null +++ b/smartjavaai-ocr/src/main/java/cn/smartjavaai/ocr/utils/CrossRangeCellMeta.java @@ -0,0 +1,42 @@ +package cn.smartjavaai.ocr.utils; + +/** + * @Auther: xiaoqiang + * @Date: 2020/12/9 9:17 + * @Description: + */ +public class CrossRangeCellMeta { + public CrossRangeCellMeta(int firstRowIndex, int firstColIndex, int rowSpan, int colSpan) { + super(); + this.firstRowIndex = firstRowIndex; + this.firstColIndex = firstColIndex; + this.rowSpan = rowSpan; + this.colSpan = colSpan; + } + + private int firstRowIndex; + private int firstColIndex; + private int rowSpan;// 跨越行数 + private int colSpan;// 跨越列数 + + public int getFirstRow() { + return firstRowIndex; + } + + public int getLastRow() { + return firstRowIndex + rowSpan - 1; + } + + public int getFirstCol() { + return firstColIndex; + } + + public int getLastCol() { + return firstColIndex + colSpan - 1; + } + + public int getColSpan(){ + return colSpan; + } +} + 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 6039b4f..797181b 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,14 +11,12 @@ import ai.djl.ndarray.NDManager; import ai.djl.opencv.OpenCVImageFactory; import cn.smartjavaai.common.entity.*; import cn.smartjavaai.common.entity.Point; -import cn.smartjavaai.ocr.entity.OcrBox; -import cn.smartjavaai.ocr.entity.OcrInfo; -import cn.smartjavaai.ocr.entity.OcrItem; -import cn.smartjavaai.ocr.entity.RotatedBoxCompX; +import cn.smartjavaai.ocr.entity.*; import cn.smartjavaai.ocr.enums.AngleEnum; 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; import org.opencv.core.Scalar; import org.opencv.imgproc.Imgproc; @@ -26,10 +24,8 @@ import org.opencv.imgproc.Imgproc; import java.awt.*; import java.awt.image.BufferedImage; import java.math.BigDecimal; -import java.util.ArrayList; -import java.util.Iterator; +import java.util.*; import java.util.List; -import java.util.Objects; /** * @author dwj @@ -42,27 +38,41 @@ public class OcrUtils { /** * 转换为OcrBox * @param dt_boxes - * @param img * @return */ - public static List convertToOcrBox(NDList dt_boxes, Image img){ - if(Objects.isNull(dt_boxes) || dt_boxes.size() == 0){ - return null; - } - List boxList = new ArrayList(); - for(NDArray box : dt_boxes){ + public static List convertToOcrBox(NDList dt_boxes) { + List boxList = new ArrayList<>(); + for (NDArray box : dt_boxes) { float[] pointsArr = box.toFloatArray(); - //log.debug("points: {}", pointsArr); - float[] lt = java.util.Arrays.copyOfRange(pointsArr, 0, 2); - 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); - OcrBox ocrBox = new OcrBox(new Point(lt[0], lt[1]), new Point(rt[0], rt[1]), new Point(rb[0], rb[1]), new Point(lb[0], lb[1])); + OcrBox ocrBox = new OcrBox( + new Point(pointsArr[0], pointsArr[1]), + new Point(pointsArr[2], pointsArr[3]), + new Point(pointsArr[4], pointsArr[5]), + new Point(pointsArr[6], pointsArr[7]) + ); boxList.add(ocrBox); } return boxList; } + /** + * 转换为OcrBox + * @param dt_boxes + * @return + */ + public static List> convertToOcrBox(List ndLists) { + if (ndLists == null || ndLists.isEmpty()) { + return Collections.emptyList(); + } + List> boxLists = new ArrayList<>(); + for (NDList dt_boxes : ndLists) { + boxLists.add(convertToOcrBox(dt_boxes)); + } + return boxLists; + } + + + /** * 欧式距离计算 * @@ -140,7 +150,6 @@ public class OcrUtils { if(Objects.isNull(lines) || lines.size() == 0){ return null; } - List detectionInfoList = new ArrayList(); List> lineList = new ArrayList>(); String fullText = ""; for(ArrayList boxList : lines){ @@ -165,6 +174,36 @@ public class OcrUtils { return new OcrInfo(lineList, fullText); } + public static OcrInfo convertRotatedBoxesToOcrItems(List rotatedBoxes) { + OcrInfo ocrInfo = new OcrInfo(); + List ocrItems = new ArrayList<>(); + StringBuilder fullText = new StringBuilder(); + for (RotatedBox rotatedBox : rotatedBoxes) { + NDArray box = rotatedBox.getBox(); + float[] points = box.toFloatArray(); + Point topLeft = new Point(points[0], points[1]); + Point topRight = new Point(points[2], points[3]); + Point bottomRight = new Point(points[4], points[5]); + Point bottomLeft = new Point(points[6], points[7]); + + OcrBox ocrBox = new OcrBox(topLeft, topRight, bottomRight, bottomLeft); + String text = rotatedBox.getText(); + + OcrItem item = new OcrItem(); + item.setOcrBox(ocrBox); + item.setText(text); + ocrItems.add(item); + fullText.append(text + " "); + } + if (fullText.length() > 0) { + fullText.deleteCharAt(fullText.length() - 1); + } + ocrInfo.setOcrItemList(ocrItems); + ocrInfo.setFullText(fullText.toString()); + return ocrInfo; + } + + /** * 放射变换+裁剪 @@ -235,26 +274,28 @@ public class OcrUtils { // 声明画笔属性 ：粗 细（单位像素）末端无修饰 折线处呈尖角 BasicStroke bStroke = new BasicStroke(2, BasicStroke.CAP_BUTT, BasicStroke.JOIN_MITER); g.setStroke(bStroke); - for(List ocrItemList : ocrInfo.getLineList()){ - for(OcrItem item : ocrItemList){ - OcrBox box = item.getOcrBox(); - int[] xPoints = { - (int)box.getTopLeft().getX(), - (int)box.getTopRight().getX(), - (int)box.getBottomRight().getX(), - (int)box.getBottomLeft().getX(), - (int)box.getTopLeft().getX() - }; - int[] yPoints = { - (int)box.getTopLeft().getY(), - (int)box.getTopRight().getY(), - (int)box.getBottomRight().getY(), - (int)box.getBottomLeft().getY(), - (int)box.getTopLeft().getY() - }; - g.drawPolyline(xPoints, yPoints, 5); - g.drawString(item.getText(), xPoints[0], yPoints[0]); - } + List ocrItemList = ocrInfo.getOcrItemList(); + if(CollectionUtils.isNotEmpty(ocrInfo.getLineList())){ + ocrItemList = ocrInfo.flattenLines(); + } + for(OcrItem item : ocrItemList){ + OcrBox box = item.getOcrBox(); + int[] xPoints = { + (int)box.getTopLeft().getX(), + (int)box.getTopRight().getX(), + (int)box.getBottomRight().getX(), + (int)box.getBottomLeft().getX(), + (int)box.getTopLeft().getX() + }; + int[] yPoints = { + (int)box.getTopLeft().getY(), + (int)box.getTopRight().getY(), + (int)box.getBottomRight().getY(), + (int)box.getBottomLeft().getY(), + (int)box.getTopLeft().getY() + }; + g.drawPolyline(xPoints, yPoints, 5); + g.drawString(item.getText(), xPoints[0], yPoints[0]); } } finally { g.dispose(); diff --git a/smartjavaai-translate/pom.xml b/smartjavaai-translate/pom.xml index 7438e4b..d856263 100644 --- a/smartjavaai-translate/pom.xml +++ b/smartjavaai-translate/pom.xml @@ -20,7 +20,7 @@ - 1.0.19 + 1.0.20 smartjavaai-translate SmartJavaAI https://github.com/geekwenjie/SmartJavaAI