mirror of
https://github.com/geekwenjie/SmartJavaAI.git
synced 2026-09-10 03:28:49 +00:00
临时提交
This commit is contained in:
150
ocr/pom.xml
Normal file
150
ocr/pom.xml
Normal file
@@ -0,0 +1,150 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0"
|
||||
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
|
||||
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
<parent>
|
||||
<groupId>cn.smartjavaai</groupId>
|
||||
<artifactId>smartjavaai-parent</artifactId>
|
||||
<version>1.0.24</version>
|
||||
</parent>
|
||||
|
||||
<artifactId>ocr</artifactId>
|
||||
|
||||
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>cn.smartjavaai</groupId>
|
||||
<artifactId>common</artifactId>
|
||||
<version>${project.version}</version>
|
||||
</dependency>
|
||||
|
||||
<!-- <dependency>-->
|
||||
<!-- <groupId>ai.djl.paddlepaddle</groupId>-->
|
||||
<!-- <artifactId>paddlepaddle-engine</artifactId>-->
|
||||
<!-- <version>0.22.1</version>-->
|
||||
<!-- </dependency>-->
|
||||
<!-- <dependency>-->
|
||||
<!-- <groupId>ai.djl.paddlepaddle</groupId>-->
|
||||
<!-- <artifactId>paddlepaddle-model-zoo</artifactId>-->
|
||||
<!-- <version>0.22.1</version>-->
|
||||
<!-- </dependency>-->
|
||||
|
||||
<dependency>
|
||||
<groupId>org.apache.poi</groupId>
|
||||
<artifactId>poi</artifactId>
|
||||
<version>4.0.0</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>dom4j</groupId>
|
||||
<artifactId>dom4j</artifactId>
|
||||
<version>1.6.1</version>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
<version>1.0.24</version>
|
||||
<name>ocr</name>
|
||||
<description>SmartJavaAI</description>
|
||||
<url>https://github.com/geekwenjie/SmartJavaAI</url>
|
||||
<licenses>
|
||||
<license>
|
||||
<name>MIT License</name>
|
||||
<url>https://opensource.org/licenses/MIT</url>
|
||||
</license>
|
||||
</licenses>
|
||||
|
||||
<build>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<groupId>org.sonatype.central</groupId>
|
||||
<artifactId>central-publishing-maven-plugin</artifactId>
|
||||
<version>0.4.0</version>
|
||||
<extensions>true</extensions>
|
||||
<configuration>
|
||||
<publishingServerId>dengwenjie</publishingServerId>
|
||||
<tokenAuth>true</tokenAuth>
|
||||
<deploymentName>${project.groupId}:${project.artifactId}:${project.version}</deploymentName>
|
||||
</configuration>
|
||||
</plugin>
|
||||
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-source-plugin</artifactId>
|
||||
<version>3.1.0</version>
|
||||
<executions>
|
||||
<execution>
|
||||
<id>attach-sources</id>
|
||||
<goals>
|
||||
<goal>jar-no-fork</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-javadoc-plugin</artifactId>
|
||||
<version>3.1.0</version>
|
||||
<configuration>
|
||||
<!-- <javadocExecutable>${java.home}/bin/javadoc</javadocExecutable>-->
|
||||
<doclint>none</doclint>
|
||||
<additionalJOptions>
|
||||
<additionalJOption>-Xdoclint:none</additionalJOption>
|
||||
</additionalJOptions>
|
||||
</configuration>
|
||||
<executions>
|
||||
<execution>
|
||||
<id>attach-javadocs</id>
|
||||
<goals>
|
||||
<goal>jar</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-gpg-plugin</artifactId>
|
||||
<version>3.1.0</version>
|
||||
<executions>
|
||||
<execution>
|
||||
<id>sign-artifacts</id>
|
||||
<phase>verify</phase>
|
||||
<goals>
|
||||
<goal>sign</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
</plugins>
|
||||
</build>
|
||||
|
||||
<!-- 必须添加:SCM信息 -->
|
||||
<scm>
|
||||
<connection>scm:git:git://github.com/geekwenjie/SmartJavaAI.git</connection>
|
||||
<developerConnection>scm:git:ssh://github.com/geekwenjie/SmartJavaAI.git</developerConnection>
|
||||
<url>http://github.com/geekwenjie/SmartJavaAI/tree/master</url>
|
||||
</scm>
|
||||
|
||||
|
||||
<distributionManagement>
|
||||
<snapshotRepository>
|
||||
<id>dengwenjie</id>
|
||||
<url>https://s01.oss.sonatype.org/content/repositories/snapshots</url>
|
||||
</snapshotRepository>
|
||||
<repository>
|
||||
<id>dengwenjie</id>
|
||||
<url>https://s01.oss.sonatype.org/service/local/staging/deploy/maven2/</url>
|
||||
</repository>
|
||||
</distributionManagement>
|
||||
|
||||
<developers>
|
||||
<developer>
|
||||
<name>dengwenjie</name>
|
||||
<email>775747758@qq.com</email>
|
||||
<roles>
|
||||
<role>Project Manager</role>
|
||||
<role>Architect</role>
|
||||
</roles>
|
||||
</developer>
|
||||
</developers>
|
||||
|
||||
</project>
|
||||
@@ -0,0 +1,36 @@
|
||||
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;
|
||||
|
||||
/**
|
||||
* 文本方向分类模型配置
|
||||
* @author dwj
|
||||
* @date 2025/4/22
|
||||
*/
|
||||
@Data
|
||||
public class DirectionModelConfig extends ModelConfig {
|
||||
|
||||
/**
|
||||
* 模型
|
||||
*/
|
||||
private DirectionModelEnum modelEnum;
|
||||
|
||||
/**
|
||||
* 检测模型路径
|
||||
*/
|
||||
private String modelPath;
|
||||
|
||||
/**
|
||||
* 文本检测模型
|
||||
*/
|
||||
private OcrCommonDetModel textDetModel;
|
||||
|
||||
|
||||
|
||||
}
|
||||
@@ -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 lombok.Data;
|
||||
|
||||
/**
|
||||
* OCR检测模型配置
|
||||
* @author dwj
|
||||
* @date 2025/4/22
|
||||
*/
|
||||
@Data
|
||||
public class OcrDetModelConfig extends ModelConfig {
|
||||
|
||||
/**
|
||||
* 模型
|
||||
*/
|
||||
private CommonDetModelEnum modelEnum;
|
||||
|
||||
/**
|
||||
* 检测模型路径
|
||||
*/
|
||||
private String detModelPath;
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
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;
|
||||
|
||||
/**
|
||||
* OCR识别模型配置
|
||||
* @author dwj
|
||||
* @date 2025/4/22
|
||||
*/
|
||||
@Data
|
||||
public class OcrRecModelConfig extends ModelConfig {
|
||||
|
||||
/**
|
||||
* 识别模型
|
||||
*/
|
||||
private CommonRecModelEnum recModelEnum;
|
||||
|
||||
/**
|
||||
* 识别模型路径
|
||||
*/
|
||||
private String recModelPath;
|
||||
|
||||
/**
|
||||
* 文本检测模型
|
||||
*/
|
||||
private OcrCommonDetModel textDetModel;
|
||||
|
||||
/**
|
||||
* 文本方向模型
|
||||
*/
|
||||
private OcrDirectionModel directionModel;
|
||||
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
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() {
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
package cn.smartjavaai.ocr.config;
|
||||
|
||||
import cn.smartjavaai.common.config.ModelConfig;
|
||||
import cn.smartjavaai.ocr.enums.CommonDetModelEnum;
|
||||
import cn.smartjavaai.ocr.enums.PlateDetModelEnum;
|
||||
import lombok.Data;
|
||||
|
||||
/**
|
||||
* 车牌检测模型配置
|
||||
* @author dwj
|
||||
*/
|
||||
@Data
|
||||
public class PlateDetModelConfig extends ModelConfig {
|
||||
|
||||
/**
|
||||
* 模型
|
||||
*/
|
||||
private PlateDetModelEnum modelEnum;
|
||||
|
||||
/**
|
||||
* 检测模型路径
|
||||
*/
|
||||
private String modelPath;
|
||||
|
||||
/**
|
||||
* 置信度阈值
|
||||
*/
|
||||
private float confidenceThreshold;
|
||||
|
||||
/**
|
||||
* iou阈值
|
||||
*/
|
||||
private float iouThreshold;
|
||||
|
||||
/**
|
||||
* 检测结果数量
|
||||
*/
|
||||
private int topK;
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
package cn.smartjavaai.ocr.config;
|
||||
|
||||
import cn.smartjavaai.common.config.ModelConfig;
|
||||
import cn.smartjavaai.ocr.enums.PlateDetModelEnum;
|
||||
import cn.smartjavaai.ocr.enums.PlateRecModelEnum;
|
||||
import cn.smartjavaai.ocr.model.plate.PlateDetModel;
|
||||
import lombok.Data;
|
||||
|
||||
/**
|
||||
* 车牌识别模型配置
|
||||
* @author dwj
|
||||
*/
|
||||
@Data
|
||||
public class PlateRecModelConfig extends ModelConfig {
|
||||
|
||||
/**
|
||||
* 模型
|
||||
*/
|
||||
private PlateRecModelEnum modelEnum;
|
||||
|
||||
/**
|
||||
* 检测模型路径
|
||||
*/
|
||||
private String modelPath;
|
||||
|
||||
/**
|
||||
* 车牌检测模型
|
||||
*/
|
||||
private PlateDetModel plateDetModel;
|
||||
|
||||
|
||||
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
package cn.smartjavaai.ocr.entity;
|
||||
|
||||
/**
|
||||
* 方向检测结果
|
||||
* @author Calvin
|
||||
* @mail 179209347@qq.com
|
||||
* @website www.aias.top
|
||||
*/
|
||||
public class DirectionInfo {
|
||||
|
||||
/**
|
||||
* 方向 0 90 180 270
|
||||
*/
|
||||
private String name;
|
||||
|
||||
/**
|
||||
* 置信度
|
||||
*/
|
||||
private Double prob;
|
||||
|
||||
public DirectionInfo(String name, Double prob) {
|
||||
this.name = name;
|
||||
this.prob = prob;
|
||||
}
|
||||
|
||||
public String getName() {
|
||||
return name;
|
||||
}
|
||||
|
||||
public void setName(String name) {
|
||||
this.name = name;
|
||||
}
|
||||
|
||||
public Double getProb() {
|
||||
return prob;
|
||||
}
|
||||
|
||||
public void setProb(Double prob) {
|
||||
this.prob = prob;
|
||||
}
|
||||
}
|
||||
12
ocr/src/main/java/cn/smartjavaai/ocr/entity/IdCardInfo.java
Normal file
12
ocr/src/main/java/cn/smartjavaai/ocr/entity/IdCardInfo.java
Normal file
@@ -0,0 +1,12 @@
|
||||
package cn.smartjavaai.ocr.entity;
|
||||
|
||||
/**
|
||||
* 身份证信息
|
||||
* @author dwj
|
||||
* @date 2025/5/22
|
||||
*/
|
||||
public class IdCardInfo {
|
||||
|
||||
|
||||
|
||||
}
|
||||
50
ocr/src/main/java/cn/smartjavaai/ocr/entity/ImageInfo.java
Normal file
50
ocr/src/main/java/cn/smartjavaai/ocr/entity/ImageInfo.java
Normal file
@@ -0,0 +1,50 @@
|
||||
package cn.smartjavaai.ocr.entity;
|
||||
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
/**
|
||||
* 图像信息
|
||||
*/
|
||||
public class ImageInfo {
|
||||
private String name;
|
||||
private Double prob;
|
||||
private Image image;
|
||||
private NDArray box;
|
||||
|
||||
public ImageInfo(Image image, NDArray box) {
|
||||
this.image = image;
|
||||
this.box = box;
|
||||
}
|
||||
|
||||
public String getName() {
|
||||
return name;
|
||||
}
|
||||
|
||||
public void setName(String name) {
|
||||
this.name = name;
|
||||
}
|
||||
|
||||
public Double getProb() {
|
||||
return prob;
|
||||
}
|
||||
|
||||
public void setProb(Double prob) {
|
||||
this.prob = prob;
|
||||
}
|
||||
|
||||
public Image getImage() {
|
||||
return image;
|
||||
}
|
||||
|
||||
public void setImage(Image image) {
|
||||
this.image = image;
|
||||
}
|
||||
|
||||
public NDArray getBox() {
|
||||
return box;
|
||||
}
|
||||
|
||||
public void setBox(NDArray box) {
|
||||
this.box = box;
|
||||
}
|
||||
}
|
||||
72
ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrBox.java
Normal file
72
ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrBox.java
Normal file
@@ -0,0 +1,72 @@
|
||||
package cn.smartjavaai.ocr.entity;
|
||||
|
||||
import cn.smartjavaai.common.entity.DetectionRectangle;
|
||||
import cn.smartjavaai.common.entity.Point;
|
||||
import lombok.Data;
|
||||
|
||||
/**
|
||||
* OCR 检测框
|
||||
* @author dwj
|
||||
* @date 2025/5/20
|
||||
*/
|
||||
@Data
|
||||
public class OcrBox {
|
||||
|
||||
/**
|
||||
* 左上角
|
||||
*/
|
||||
private Point topLeft;
|
||||
|
||||
/**
|
||||
* 右上角
|
||||
*/
|
||||
private Point topRight;
|
||||
|
||||
/**
|
||||
* 右下角
|
||||
*/
|
||||
private Point bottomRight;
|
||||
|
||||
/**
|
||||
* 左下角
|
||||
*/
|
||||
private Point bottomLeft;
|
||||
|
||||
public OcrBox(Point topLeft, Point topRight, Point bottomRight, Point bottomLeft) {
|
||||
this.topLeft = topLeft;
|
||||
this.topRight = topRight;
|
||||
this.bottomRight = bottomRight;
|
||||
this.bottomLeft = bottomLeft;
|
||||
}
|
||||
|
||||
public OcrBox() {
|
||||
}
|
||||
|
||||
public float[] toFloatArray() {
|
||||
return new float[]{
|
||||
(float)topLeft.getX(), (float)topLeft.getY(),
|
||||
(float)topRight.getX(), (float)topRight.getY(),
|
||||
(float)bottomRight.getX(), (float)bottomRight.getY(),
|
||||
(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;
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
38
ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrInfo.java
Normal file
38
ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrInfo.java
Normal file
@@ -0,0 +1,38 @@
|
||||
package cn.smartjavaai.ocr.entity;
|
||||
|
||||
import lombok.Data;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
/**
|
||||
* OCR信息
|
||||
* @author dwj
|
||||
* @date 2025/5/20
|
||||
*/
|
||||
@Data
|
||||
public class OcrInfo {
|
||||
|
||||
private List<List<OcrItem>> lineList;
|
||||
|
||||
private List<OcrItem> ocrItemList;
|
||||
|
||||
private String fullText;
|
||||
|
||||
private String base64Img;
|
||||
|
||||
|
||||
public OcrInfo(List<List<OcrItem>> lineList, String fullText) {
|
||||
this.lineList = lineList;
|
||||
this.fullText = fullText;
|
||||
}
|
||||
public OcrInfo() {
|
||||
}
|
||||
|
||||
public List<OcrItem> flattenLines() {
|
||||
return lineList.stream()
|
||||
.flatMap(List::stream)
|
||||
.collect(Collectors.toList());
|
||||
}
|
||||
}
|
||||
54
ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrItem.java
Normal file
54
ocr/src/main/java/cn/smartjavaai/ocr/entity/OcrItem.java
Normal file
@@ -0,0 +1,54 @@
|
||||
package cn.smartjavaai.ocr.entity;
|
||||
|
||||
import cn.smartjavaai.ocr.enums.AngleEnum;
|
||||
import lombok.Data;
|
||||
|
||||
/**
|
||||
* @author dwj
|
||||
* @date 2025/5/20
|
||||
*/
|
||||
@Data
|
||||
public class OcrItem {
|
||||
|
||||
/**
|
||||
* 识别框
|
||||
*/
|
||||
private OcrBox ocrBox;
|
||||
|
||||
/**
|
||||
* 文本
|
||||
*/
|
||||
private String text;
|
||||
|
||||
/**
|
||||
* 方向
|
||||
*/
|
||||
private AngleEnum angle;
|
||||
|
||||
/**
|
||||
* 检测得分
|
||||
*/
|
||||
private float score;
|
||||
|
||||
|
||||
public OcrItem(OcrBox ocrBox, String text) {
|
||||
this.ocrBox = ocrBox;
|
||||
this.text = text;
|
||||
}
|
||||
|
||||
public OcrItem() {
|
||||
}
|
||||
|
||||
public OcrItem(OcrBox ocrBox, String text, AngleEnum angle) {
|
||||
this.ocrBox = ocrBox;
|
||||
this.text = text;
|
||||
this.angle = angle;
|
||||
}
|
||||
|
||||
public OcrItem(OcrBox ocrBox, AngleEnum angle, float score) {
|
||||
this.ocrBox = ocrBox;
|
||||
this.angle = angle;
|
||||
this.score = score;
|
||||
}
|
||||
|
||||
}
|
||||
45
ocr/src/main/java/cn/smartjavaai/ocr/entity/PlateInfo.java
Normal file
45
ocr/src/main/java/cn/smartjavaai/ocr/entity/PlateInfo.java
Normal file
@@ -0,0 +1,45 @@
|
||||
package cn.smartjavaai.ocr.entity;
|
||||
|
||||
import cn.smartjavaai.common.entity.DetectionRectangle;
|
||||
import cn.smartjavaai.ocr.enums.PlateType;
|
||||
import lombok.Data;
|
||||
|
||||
/**
|
||||
* 车牌识别信息
|
||||
* @author dwj
|
||||
*/
|
||||
@Data
|
||||
public class PlateInfo {
|
||||
|
||||
/**
|
||||
* 车牌类型
|
||||
*/
|
||||
private PlateType plateType;
|
||||
|
||||
/**
|
||||
* 车牌号码
|
||||
*/
|
||||
private String plateNumber;
|
||||
|
||||
/**
|
||||
* 车牌颜色
|
||||
*/
|
||||
private String plateColor;
|
||||
|
||||
/**
|
||||
* 检测位置信息
|
||||
*/
|
||||
private DetectionRectangle detectionRectangle;
|
||||
|
||||
/**
|
||||
* 车牌4角坐标
|
||||
*/
|
||||
private OcrBox box;
|
||||
|
||||
/**
|
||||
* 检测得分
|
||||
*/
|
||||
private float score;
|
||||
|
||||
|
||||
}
|
||||
35
ocr/src/main/java/cn/smartjavaai/ocr/entity/PlateResult.java
Normal file
35
ocr/src/main/java/cn/smartjavaai/ocr/entity/PlateResult.java
Normal file
@@ -0,0 +1,35 @@
|
||||
package cn.smartjavaai.ocr.entity;
|
||||
|
||||
import lombok.Data;
|
||||
|
||||
/**
|
||||
* @author dwj
|
||||
*/
|
||||
@Data
|
||||
public class PlateResult {
|
||||
|
||||
/**
|
||||
* 车牌号码
|
||||
*/
|
||||
private String plateNo;
|
||||
|
||||
/**
|
||||
* 车牌颜色
|
||||
*/
|
||||
private String plateColor;
|
||||
|
||||
public PlateResult(String plateNo, String plateColor) {
|
||||
this.plateNo = plateNo;
|
||||
this.plateColor = plateColor;
|
||||
}
|
||||
|
||||
|
||||
@Override
|
||||
public String toString() {
|
||||
return "PlateResult{" +
|
||||
"plateNo='" + plateNo + '\'' +
|
||||
", plateColor='" + plateColor + '\'' +
|
||||
'}';
|
||||
}
|
||||
}
|
||||
|
||||
46
ocr/src/main/java/cn/smartjavaai/ocr/entity/RotatedBox.java
Normal file
46
ocr/src/main/java/cn/smartjavaai/ocr/entity/RotatedBox.java
Normal file
@@ -0,0 +1,46 @@
|
||||
package cn.smartjavaai.ocr.entity;
|
||||
|
||||
import ai.djl.ndarray.NDArray;
|
||||
/**
|
||||
* 旋转检测框
|
||||
*/
|
||||
public class RotatedBox implements Comparable<RotatedBox> {
|
||||
private NDArray box;
|
||||
private String text;
|
||||
|
||||
public RotatedBox(NDArray box, String text) {
|
||||
this.box = box;
|
||||
this.text = text;
|
||||
}
|
||||
|
||||
/**
|
||||
* 将左上角 Y 坐标升序排序
|
||||
*
|
||||
* @param o
|
||||
* @return
|
||||
*/
|
||||
@Override
|
||||
public int compareTo(RotatedBox o) {
|
||||
NDArray lowBox = this.getBox();
|
||||
NDArray highBox = o.getBox();
|
||||
float lowY = lowBox.toFloatArray()[1];
|
||||
float highY = highBox.toFloatArray()[1];
|
||||
return (lowY < highY) ? -1 : 1;
|
||||
}
|
||||
|
||||
public NDArray getBox() {
|
||||
return box;
|
||||
}
|
||||
|
||||
public void setBox(NDArray box) {
|
||||
this.box = box;
|
||||
}
|
||||
|
||||
public String getText() {
|
||||
return text;
|
||||
}
|
||||
|
||||
public void setText(String text) {
|
||||
this.text = text;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package cn.smartjavaai.ocr.entity;
|
||||
|
||||
import ai.djl.ndarray.NDArray;
|
||||
|
||||
/**
|
||||
* 旋转检测框 - 支持左上角 X 坐标升序排序
|
||||
*/
|
||||
public class RotatedBoxCompX implements Comparable<RotatedBoxCompX> {
|
||||
private NDArray box;
|
||||
private String text;
|
||||
|
||||
public RotatedBoxCompX(NDArray box, String text) {
|
||||
this.box = box;
|
||||
this.text = text;
|
||||
}
|
||||
|
||||
/**
|
||||
* 将左上角 X 坐标升序排序
|
||||
*
|
||||
* @param o
|
||||
* @return
|
||||
*/
|
||||
@Override
|
||||
public int compareTo(RotatedBoxCompX o) {
|
||||
NDArray leftBox = this.getBox();
|
||||
NDArray rightBox = o.getBox();
|
||||
float leftX = leftBox.toFloatArray()[0];
|
||||
float rightX = rightBox.toFloatArray()[0];
|
||||
return (leftX < rightX) ? -1 : 1;
|
||||
}
|
||||
|
||||
public NDArray getBox() {
|
||||
return box;
|
||||
}
|
||||
|
||||
public void setBox(NDArray box) {
|
||||
this.box = box;
|
||||
}
|
||||
|
||||
public String getText() {
|
||||
return text;
|
||||
}
|
||||
|
||||
public void setText(String text) {
|
||||
this.text = text;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
package cn.smartjavaai.ocr.entity;
|
||||
|
||||
import lombok.Data;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* @author dwj
|
||||
*/
|
||||
@Data
|
||||
public class TableStructureResult {
|
||||
|
||||
private List<OcrItem> ocrItemList;
|
||||
|
||||
private List<String> tableTagList;
|
||||
|
||||
private String html;
|
||||
|
||||
|
||||
public TableStructureResult(List<OcrItem> ocrItemList, List<String> tableTagList) {
|
||||
this.ocrItemList = ocrItemList;
|
||||
this.tableTagList = tableTagList;
|
||||
}
|
||||
|
||||
public TableStructureResult() {
|
||||
}
|
||||
|
||||
public TableStructureResult(List<OcrItem> ocrItemList, List<String> tableTagList, String html) {
|
||||
this.ocrItemList = ocrItemList;
|
||||
this.tableTagList = tableTagList;
|
||||
this.html = html;
|
||||
}
|
||||
}
|
||||
38
ocr/src/main/java/cn/smartjavaai/ocr/enums/AngleEnum.java
Normal file
38
ocr/src/main/java/cn/smartjavaai/ocr/enums/AngleEnum.java
Normal file
@@ -0,0 +1,38 @@
|
||||
package cn.smartjavaai.ocr.enums;
|
||||
|
||||
/**
|
||||
* 文本方向
|
||||
* @author dwj
|
||||
* @date 2025/5/23
|
||||
*/
|
||||
public enum AngleEnum {
|
||||
|
||||
ANGLE_0("0"),
|
||||
ANGLE_90("90"),
|
||||
ANGLE_180("180"),
|
||||
ANGLE_270("270");
|
||||
|
||||
private final String value;
|
||||
|
||||
AngleEnum(String value) {
|
||||
this.value = value;
|
||||
}
|
||||
|
||||
public String getValue() {
|
||||
return value;
|
||||
}
|
||||
|
||||
public static AngleEnum fromValue(String value) {
|
||||
for (AngleEnum angle : values()) {
|
||||
if (angle.value.equals(value)) {
|
||||
return angle;
|
||||
}
|
||||
}
|
||||
throw new IllegalArgumentException("Invalid angle value: " + value);
|
||||
}
|
||||
|
||||
@Override
|
||||
public String toString() {
|
||||
return value + "°";
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package cn.smartjavaai.ocr.enums;
|
||||
|
||||
/**
|
||||
* OCR检测模型枚举
|
||||
* @author dwj
|
||||
*/
|
||||
public enum CommonDetModelEnum {
|
||||
|
||||
PP_OCR_V5_SERVER_DET_MODEL,
|
||||
|
||||
PP_OCR_V5_MOBILE_DET_MODEL,
|
||||
|
||||
PP_OCR_V4_SERVER_DET_MODEL,
|
||||
|
||||
PP_OCR_V4_MOBILE_DET_MODEL;
|
||||
|
||||
|
||||
/**
|
||||
* 根据名称获取枚举 (忽略大小写和下划线变体)
|
||||
*/
|
||||
public static CommonDetModelEnum fromName(String name) {
|
||||
String formatted = name.trim().toUpperCase().replaceAll("[-_]", "");
|
||||
for (CommonDetModelEnum model : values()) {
|
||||
if (model.name().replaceAll("_", "").equals(formatted)) {
|
||||
return model;
|
||||
}
|
||||
}
|
||||
throw new IllegalArgumentException("未知模型名称: " + name);
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package cn.smartjavaai.ocr.enums;
|
||||
|
||||
/**
|
||||
* OCR识别模型枚举
|
||||
* @author dwj
|
||||
*/
|
||||
public enum CommonRecModelEnum {
|
||||
|
||||
PP_OCR_V5_SERVER_REC_MODEL,
|
||||
|
||||
PP_OCR_V5_MOBILE_REC_MODEL,
|
||||
|
||||
PP_OCR_V4_SERVER_REC_MODEL,
|
||||
|
||||
PP_OCR_V4_MOBILE_REC_MODEL;
|
||||
|
||||
|
||||
/**
|
||||
* 根据名称获取枚举 (忽略大小写和下划线变体)
|
||||
*/
|
||||
public static CommonRecModelEnum fromName(String name) {
|
||||
String formatted = name.trim().toUpperCase().replaceAll("[-_]", "");
|
||||
for (CommonRecModelEnum model : values()) {
|
||||
if (model.name().replaceAll("_", "").equals(formatted)) {
|
||||
return model;
|
||||
}
|
||||
}
|
||||
throw new IllegalArgumentException("未知模型名称: " + name);
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package cn.smartjavaai.ocr.enums;
|
||||
|
||||
/**
|
||||
* OCR文本方向分类模型枚举
|
||||
* @author dwj
|
||||
* @date 2025/4/4
|
||||
*/
|
||||
public enum DirectionModelEnum {
|
||||
|
||||
CH_PPOCR_MOBILE_V2_CLS,
|
||||
|
||||
PP_LCNET_X0_25,
|
||||
|
||||
PP_LCNET_X1_0;
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* 根据名称获取枚举 (忽略大小写和下划线变体)
|
||||
*/
|
||||
public static DirectionModelEnum fromName(String name) {
|
||||
String formatted = name.trim().toUpperCase().replaceAll("[-_]", "");
|
||||
for (DirectionModelEnum model : values()) {
|
||||
if (model.name().replaceAll("_", "").equals(formatted)) {
|
||||
return model;
|
||||
}
|
||||
}
|
||||
throw new IllegalArgumentException("未知模型名称: " + name);
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package cn.smartjavaai.ocr.enums;
|
||||
|
||||
/**
|
||||
* 车牌检测模型枚举
|
||||
* @author dwj
|
||||
*/
|
||||
public enum PlateDetModelEnum {
|
||||
|
||||
YOLOV5,
|
||||
|
||||
YOLOV7;
|
||||
|
||||
|
||||
/**
|
||||
* 根据名称获取枚举 (忽略大小写和下划线变体)
|
||||
*/
|
||||
public static PlateDetModelEnum fromName(String name) {
|
||||
String formatted = name.trim().toUpperCase().replaceAll("[-_]", "");
|
||||
for (PlateDetModelEnum model : values()) {
|
||||
if (model.name().replaceAll("_", "").equals(formatted)) {
|
||||
return model;
|
||||
}
|
||||
}
|
||||
throw new IllegalArgumentException("未知模型名称: " + name);
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
package cn.smartjavaai.ocr.enums;
|
||||
|
||||
/**
|
||||
* 车牌识别模型枚举
|
||||
* @author dwj
|
||||
*/
|
||||
public enum PlateRecModelEnum {
|
||||
|
||||
PLATE_REC_CRNN;
|
||||
|
||||
|
||||
/**
|
||||
* 根据名称获取枚举 (忽略大小写和下划线变体)
|
||||
*/
|
||||
public static PlateRecModelEnum fromName(String name) {
|
||||
String formatted = name.trim().toUpperCase().replaceAll("[-_]", "");
|
||||
for (PlateRecModelEnum model : values()) {
|
||||
if (model.name().replaceAll("_", "").equals(formatted)) {
|
||||
return model;
|
||||
}
|
||||
}
|
||||
throw new IllegalArgumentException("未知模型名称: " + name);
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
45
ocr/src/main/java/cn/smartjavaai/ocr/enums/PlateType.java
Normal file
45
ocr/src/main/java/cn/smartjavaai/ocr/enums/PlateType.java
Normal file
@@ -0,0 +1,45 @@
|
||||
package cn.smartjavaai.ocr.enums;
|
||||
|
||||
/**
|
||||
* @author dwj
|
||||
*/
|
||||
public enum PlateType {
|
||||
|
||||
SINGLE("single", "单层"),
|
||||
DOUBLE("double", "双层"),
|
||||
UNKNOWN("unknown", "未知");
|
||||
|
||||
private final String className;
|
||||
private final String description;
|
||||
|
||||
PlateType(String className, String description) {
|
||||
this.className = className;
|
||||
this.description = description;
|
||||
}
|
||||
|
||||
public String getClassName() {
|
||||
return className;
|
||||
}
|
||||
|
||||
public String getDescription() {
|
||||
return description;
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* 根据value获取对应的PlateType
|
||||
* @param className
|
||||
* @return PlateType
|
||||
*/
|
||||
public static PlateType fromClassName(String className) {
|
||||
for (PlateType type : values()) {
|
||||
if (type.className.equals(className)) {
|
||||
return type;
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
package cn.smartjavaai.ocr.exception;
|
||||
|
||||
/**
|
||||
* OCR异常
|
||||
* @author dwj
|
||||
* @date 2025/4/4
|
||||
*/
|
||||
public class OcrException extends RuntimeException{
|
||||
|
||||
public OcrException() {
|
||||
super();
|
||||
}
|
||||
|
||||
public OcrException(String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) {
|
||||
super(message, cause, enableSuppression, writableStackTrace);
|
||||
}
|
||||
|
||||
public OcrException(String message, Throwable cause) {
|
||||
super(message, cause);
|
||||
}
|
||||
|
||||
public OcrException(String message) {
|
||||
super(message);
|
||||
}
|
||||
|
||||
public OcrException(Throwable cause) {
|
||||
super(cause);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,223 @@
|
||||
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.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.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;
|
||||
import java.util.Objects;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
|
||||
/**
|
||||
* OCR模型工厂
|
||||
* @author dwj
|
||||
*/
|
||||
@Slf4j
|
||||
public class OcrModelFactory {
|
||||
|
||||
// 使用 volatile 和双重检查锁定来确保线程安全的单例模式
|
||||
private static volatile OcrModelFactory instance;
|
||||
|
||||
private static final ConcurrentHashMap<CommonDetModelEnum, OcrCommonDetModel> commonDetModelMap = new ConcurrentHashMap<>();
|
||||
|
||||
|
||||
private static final ConcurrentHashMap<CommonRecModelEnum, OcrCommonRecModel> commonRecModelMap = new ConcurrentHashMap<>();
|
||||
|
||||
private static final ConcurrentHashMap<DirectionModelEnum, OcrDirectionModel> directionModelMap = new ConcurrentHashMap<>();
|
||||
|
||||
/**
|
||||
* 检测模型注册表
|
||||
*/
|
||||
private static final Map<CommonDetModelEnum, Class<? extends OcrCommonDetModel>> commonDetRegistry =
|
||||
new ConcurrentHashMap<>();
|
||||
|
||||
/**
|
||||
* 识别模型注册表
|
||||
*/
|
||||
private static final Map<CommonRecModelEnum, Class<? extends OcrCommonRecModel>> commonRecRegistry =
|
||||
new ConcurrentHashMap<>();
|
||||
|
||||
/**
|
||||
* 方向分类模型注册表
|
||||
*/
|
||||
private static final Map<DirectionModelEnum, Class<? extends OcrDirectionModel>> directionRegistry =
|
||||
new ConcurrentHashMap<>();
|
||||
|
||||
|
||||
public static OcrModelFactory getInstance() {
|
||||
if (instance == null) {
|
||||
synchronized (OcrModelFactory.class) {
|
||||
if (instance == null) {
|
||||
instance = new OcrModelFactory();
|
||||
}
|
||||
}
|
||||
}
|
||||
return instance;
|
||||
}
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* 注册通用检测模型
|
||||
* @param detModelEnum
|
||||
* @param clazz
|
||||
*/
|
||||
private static void registerCommonDetModel(CommonDetModelEnum detModelEnum, Class<? extends OcrCommonDetModel> clazz) {
|
||||
commonDetRegistry.put(detModelEnum, clazz);
|
||||
}
|
||||
|
||||
/**
|
||||
* 注册通用识别模型
|
||||
* @param recModelEnum
|
||||
* @param clazz
|
||||
*/
|
||||
private static void registerCommonRecModel(CommonRecModelEnum recModelEnum, Class<? extends OcrCommonRecModel> clazz) {
|
||||
commonRecRegistry.put(recModelEnum, clazz);
|
||||
}
|
||||
|
||||
/**
|
||||
* 注册通用方向分类模型
|
||||
* @param directionModelEnum
|
||||
* @param clazz
|
||||
*/
|
||||
private static void registerDirectionModel(DirectionModelEnum directionModelEnum, Class<? extends OcrDirectionModel> clazz) {
|
||||
directionRegistry.put(directionModelEnum, clazz);
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 获取检测模型(通过配置)
|
||||
* @param config
|
||||
* @return
|
||||
*/
|
||||
public OcrCommonDetModel getDetModel(OcrDetModelConfig config) {
|
||||
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
|
||||
throw new OcrException("未配置OCR模型");
|
||||
}
|
||||
return commonDetModelMap.computeIfAbsent(config.getModelEnum(), k -> {
|
||||
return createCommonDetModel(config);
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取识别模型(通过配置)
|
||||
* @param config
|
||||
* @return
|
||||
*/
|
||||
public OcrCommonRecModel getRecModel(OcrRecModelConfig config) {
|
||||
if(Objects.isNull(config) || Objects.isNull(config.getRecModelEnum())){
|
||||
throw new OcrException("未配置OCR模型");
|
||||
}
|
||||
return commonRecModelMap.computeIfAbsent(config.getRecModelEnum(), k -> {
|
||||
return createCommonRecModel(config);
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取模型(通过配置)
|
||||
* @param config
|
||||
* @return
|
||||
*/
|
||||
public OcrDirectionModel getDirectionModel(DirectionModelConfig config) {
|
||||
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
|
||||
throw new OcrException("未配置OCR模型");
|
||||
}
|
||||
return directionModelMap.computeIfAbsent(config.getModelEnum(), k -> {
|
||||
return createDirectionModel(config);
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* 创建OCR通用检测模型
|
||||
* @param config
|
||||
* @return
|
||||
*/
|
||||
private OcrCommonDetModel createCommonDetModel(OcrDetModelConfig config) {
|
||||
Class<?> clazz = commonDetRegistry.get(config.getModelEnum());
|
||||
if(clazz == null){
|
||||
throw new OcrException("Unsupported model");
|
||||
}
|
||||
OcrCommonDetModel model = null;
|
||||
try {
|
||||
model = (OcrCommonDetModel) clazz.newInstance();
|
||||
} catch (InstantiationException | IllegalAccessException e) {
|
||||
throw new OcrException(e);
|
||||
}
|
||||
model.loadModel(config);
|
||||
return model;
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 创建OCR通用识别模型
|
||||
* @param config
|
||||
* @return
|
||||
*/
|
||||
private OcrCommonRecModel createCommonRecModel(OcrRecModelConfig config) {
|
||||
Class<?> clazz = commonRecRegistry.get(config.getRecModelEnum());
|
||||
if(clazz == null){
|
||||
throw new OcrException("Unsupported model");
|
||||
}
|
||||
OcrCommonRecModel model = null;
|
||||
try {
|
||||
model = (OcrCommonRecModel) clazz.newInstance();
|
||||
} catch (InstantiationException | IllegalAccessException e) {
|
||||
throw new OcrException(e);
|
||||
}
|
||||
model.loadModel(config);
|
||||
return model;
|
||||
}
|
||||
|
||||
/**
|
||||
* 创建OCR方向分类模型
|
||||
* @param config
|
||||
* @return
|
||||
*/
|
||||
private OcrDirectionModel createDirectionModel(DirectionModelConfig config) {
|
||||
Class<?> clazz = directionRegistry.get(config.getModelEnum());
|
||||
if(clazz == null){
|
||||
throw new OcrException("Unsupported model");
|
||||
}
|
||||
OcrDirectionModel model = null;
|
||||
try {
|
||||
model = (OcrDirectionModel) clazz.newInstance();
|
||||
} catch (InstantiationException | IllegalAccessException e) {
|
||||
throw new OcrException(e);
|
||||
}
|
||||
model.loadModel(config);
|
||||
return model;
|
||||
}
|
||||
|
||||
|
||||
// 初始化默认算法
|
||||
static {
|
||||
//通用-检测模型
|
||||
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());
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
package cn.smartjavaai.ocr.factory;
|
||||
|
||||
import cn.smartjavaai.common.config.Config;
|
||||
import cn.smartjavaai.ocr.config.PlateDetModelConfig;
|
||||
import cn.smartjavaai.ocr.config.PlateRecModelConfig;
|
||||
import cn.smartjavaai.ocr.config.TableStructureConfig;
|
||||
import cn.smartjavaai.ocr.enums.PlateDetModelEnum;
|
||||
import cn.smartjavaai.ocr.enums.PlateRecModelEnum;
|
||||
import cn.smartjavaai.ocr.enums.TableStructureModelEnum;
|
||||
import cn.smartjavaai.ocr.exception.OcrException;
|
||||
import cn.smartjavaai.ocr.model.plate.CRNNPlateRecModel;
|
||||
import cn.smartjavaai.ocr.model.plate.PlateDetModel;
|
||||
import cn.smartjavaai.ocr.model.plate.PlateRecModel;
|
||||
import cn.smartjavaai.ocr.model.plate.Yolov5PlateDetModel;
|
||||
import cn.smartjavaai.ocr.model.table.CommonTableStructureModel;
|
||||
import cn.smartjavaai.ocr.model.table.TableStructureModel;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
|
||||
import java.util.Map;
|
||||
import java.util.Objects;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
|
||||
/**
|
||||
* 车牌识别模型工厂
|
||||
* @author dwj
|
||||
*/
|
||||
@Slf4j
|
||||
public class PlateModelFactory {
|
||||
|
||||
// 使用 volatile 和双重检查锁定来确保线程安全的单例模式
|
||||
private static volatile PlateModelFactory instance;
|
||||
|
||||
/**
|
||||
* 模型缓存
|
||||
*/
|
||||
private static final ConcurrentHashMap<PlateDetModelEnum, PlateDetModel> detModelMap = new ConcurrentHashMap<>();
|
||||
|
||||
/**
|
||||
* 模型缓存
|
||||
*/
|
||||
private static final ConcurrentHashMap<PlateRecModelEnum, PlateRecModel> recModelMap = new ConcurrentHashMap<>();
|
||||
|
||||
|
||||
/**
|
||||
* 模型注册表
|
||||
*/
|
||||
private static final Map<PlateDetModelEnum, Class<? extends PlateDetModel>> detModelRegistry =
|
||||
new ConcurrentHashMap<>();
|
||||
|
||||
/**
|
||||
* 模型注册表
|
||||
*/
|
||||
private static final Map<PlateRecModelEnum, Class<? extends PlateRecModel>> recModelRegistry =
|
||||
new ConcurrentHashMap<>();
|
||||
|
||||
|
||||
public static PlateModelFactory getInstance() {
|
||||
if (instance == null) {
|
||||
synchronized (PlateModelFactory.class) {
|
||||
if (instance == null) {
|
||||
instance = new PlateModelFactory();
|
||||
}
|
||||
}
|
||||
}
|
||||
return instance;
|
||||
}
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* 注册模型
|
||||
* @param plateDetModelEnum
|
||||
* @param clazz
|
||||
*/
|
||||
private static void registerDetModel(PlateDetModelEnum plateDetModelEnum, Class<? extends PlateDetModel> clazz) {
|
||||
detModelRegistry.put(plateDetModelEnum, clazz);
|
||||
}
|
||||
|
||||
/**
|
||||
* 注册模型
|
||||
* @param plateRecModelEnum
|
||||
* @param clazz
|
||||
*/
|
||||
private static void registerRecModel(PlateRecModelEnum plateRecModelEnum, Class<? extends PlateRecModel> clazz) {
|
||||
recModelRegistry.put(plateRecModelEnum, clazz);
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 获取模型
|
||||
* @param config
|
||||
* @return
|
||||
*/
|
||||
public PlateDetModel getDetModel(PlateDetModelConfig config) {
|
||||
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
|
||||
throw new OcrException("未配置OCR模型");
|
||||
}
|
||||
return detModelMap.computeIfAbsent(config.getModelEnum(), k -> {
|
||||
return createDetModel(config);
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取模型
|
||||
* @param config
|
||||
* @return
|
||||
*/
|
||||
public PlateRecModel getRecModel(PlateRecModelConfig config) {
|
||||
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
|
||||
throw new OcrException("未配置OCR模型");
|
||||
}
|
||||
return recModelMap.computeIfAbsent(config.getModelEnum(), k -> {
|
||||
return createRecModel(config);
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* 创建检测模型
|
||||
* @param config
|
||||
* @return
|
||||
*/
|
||||
private PlateDetModel createDetModel(PlateDetModelConfig config) {
|
||||
Class<?> clazz = detModelRegistry.get(config.getModelEnum());
|
||||
if(clazz == null){
|
||||
throw new OcrException("Unsupported model");
|
||||
}
|
||||
PlateDetModel model = null;
|
||||
try {
|
||||
model = (PlateDetModel) clazz.newInstance();
|
||||
} catch (InstantiationException | IllegalAccessException e) {
|
||||
throw new OcrException(e);
|
||||
}
|
||||
model.loadModel(config);
|
||||
return model;
|
||||
}
|
||||
|
||||
/**
|
||||
* 创建识别模型
|
||||
* @param config
|
||||
* @return
|
||||
*/
|
||||
private PlateRecModel createRecModel(PlateRecModelConfig config) {
|
||||
Class<?> clazz = recModelRegistry.get(config.getModelEnum());
|
||||
if(clazz == null){
|
||||
throw new OcrException("Unsupported model");
|
||||
}
|
||||
PlateRecModel model = null;
|
||||
try {
|
||||
model = (PlateRecModel) clazz.newInstance();
|
||||
} catch (InstantiationException | IllegalAccessException e) {
|
||||
throw new OcrException(e);
|
||||
}
|
||||
model.loadModel(config);
|
||||
return model;
|
||||
}
|
||||
|
||||
|
||||
// 初始化默认算法
|
||||
static {
|
||||
registerDetModel(PlateDetModelEnum.YOLOV5, Yolov5PlateDetModel.class);
|
||||
registerDetModel(PlateDetModelEnum.YOLOV7, Yolov5PlateDetModel.class);
|
||||
registerRecModel(PlateRecModelEnum.PLATE_REC_CRNN, CRNNPlateRecModel.class);
|
||||
log.debug("缓存目录:{}", Config.getCachePath());
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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<TableStructureModelEnum, TableStructureModel> tableStructureModelMap = new ConcurrentHashMap<>();
|
||||
|
||||
|
||||
/**
|
||||
* 模型注册表
|
||||
*/
|
||||
private static final Map<TableStructureModelEnum, Class<? extends TableStructureModel>> 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<? extends TableStructureModel> 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());
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
package cn.smartjavaai.ocr.model.common.detect;
|
||||
|
||||
import ai.djl.inference.Predictor;
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import cn.smartjavaai.ocr.config.OcrDetModelConfig;
|
||||
import cn.smartjavaai.ocr.entity.OcrBox;
|
||||
import org.apache.commons.pool2.impl.GenericObjectPool;
|
||||
|
||||
import java.awt.image.BufferedImage;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* OCR 通用检测模型
|
||||
* @author dwj
|
||||
*/
|
||||
public interface OcrCommonDetModel extends AutoCloseable{
|
||||
|
||||
/**
|
||||
* 加载模型
|
||||
* @param config
|
||||
*/
|
||||
void loadModel(OcrDetModelConfig config); // 加载模型
|
||||
|
||||
/**
|
||||
* 文本检测
|
||||
* @param imagePath 图片路径
|
||||
* @return
|
||||
*/
|
||||
default List<OcrBox> detect(String imagePath) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 文本检测
|
||||
* @param image BufferedImage
|
||||
* @return
|
||||
*/
|
||||
default List<OcrBox> detect(BufferedImage image) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 文本检测
|
||||
* @param imageData 图片字节数组
|
||||
* @return
|
||||
*/
|
||||
default List<OcrBox> detect(byte[] imageData) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 文本检测
|
||||
* @param image DJL Image
|
||||
* @return
|
||||
*/
|
||||
default List<OcrBox> detect(Image image){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 检测并绘制结果
|
||||
* @param imagePath 图片输入路径(包含文件名称)
|
||||
* @param outputPath 图片输出路径(包含文件名称)
|
||||
*/
|
||||
default void detectAndDraw(String imagePath, String outputPath) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 检测并绘制结果
|
||||
* @param sourceImage
|
||||
* @return
|
||||
*/
|
||||
default BufferedImage detectAndDraw(BufferedImage sourceImage){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 文本检测(批量)
|
||||
* @param imageList BufferedImage
|
||||
* @return
|
||||
*/
|
||||
default List<List<OcrBox>> batchDetect(List<BufferedImage> imageList) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 文本检测(批量)
|
||||
* @param imageList DJL Image
|
||||
* @return
|
||||
*/
|
||||
default List<List<OcrBox>> batchDetectDJLImage(List<Image> imageList){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
default GenericObjectPool<Predictor<Image, NDList>> getPool(){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,242 @@
|
||||
package cn.smartjavaai.ocr.model.common.detect;
|
||||
|
||||
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.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.entity.OcrBox;
|
||||
import cn.smartjavaai.ocr.exception.OcrException;
|
||||
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;
|
||||
import org.apache.commons.pool2.ObjectPool;
|
||||
import org.apache.commons.pool2.impl.GenericObjectPool;
|
||||
import org.opencv.core.Mat;
|
||||
|
||||
import javax.imageio.ImageIO;
|
||||
import java.awt.image.BufferedImage;
|
||||
import java.io.ByteArrayInputStream;
|
||||
import java.io.ByteArrayOutputStream;
|
||||
import java.io.IOException;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.nio.file.Paths;
|
||||
import java.util.*;
|
||||
|
||||
/**
|
||||
* ocr通用检测模型实现类
|
||||
* @author dwj
|
||||
*/
|
||||
@Slf4j
|
||||
public class OcrCommonDetModelImpl implements OcrCommonDetModel{
|
||||
|
||||
private GenericObjectPool<Predictor<Image, NDList>> detPredictorPool;
|
||||
|
||||
private ZooModel<Image, NDList> detectionModel;
|
||||
|
||||
private OcrDetModelConfig config;
|
||||
|
||||
@Override
|
||||
public void loadModel(OcrDetModelConfig config){
|
||||
if(StringUtils.isBlank(config.getDetModelPath())){
|
||||
throw new OcrException("modelPath is null");
|
||||
}
|
||||
this.config = config;
|
||||
//初始化 检测Criteria
|
||||
Criteria<Image, NDList> detCriteria = OcrCommonDetCriterialFactory.createCriteria(config);
|
||||
try{
|
||||
detectionModel = ModelZoo.loadModel(detCriteria);
|
||||
// 创建池子:每个线程独享 Predictor
|
||||
this.detPredictorPool = new GenericObjectPool<>(new PredictorFactory<>(detectionModel));
|
||||
int predictorPoolSize = config.getPredictorPoolSize();
|
||||
if(config.getPredictorPoolSize() <= 0){
|
||||
predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
|
||||
}
|
||||
detPredictorPool.setMaxTotal(predictorPoolSize);
|
||||
log.debug("当前设备: " + detectionModel.getNDManager().getDevice());
|
||||
log.debug("当前引擎: " + Engine.getInstance().getEngineName());
|
||||
log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
|
||||
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
|
||||
throw new OcrException("检测模型加载失败", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<OcrBox> detect(String imagePath){
|
||||
if(!FileUtils.isFileExists(imagePath)){
|
||||
throw new OcrException("图像文件不存在");
|
||||
}
|
||||
Image img = null;
|
||||
try {
|
||||
img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
|
||||
} catch (IOException e) {
|
||||
throw new OcrException("无效的图片", e);
|
||||
}
|
||||
List<OcrBox> ocrBoxList = detect(img);
|
||||
((Mat)img.getWrappedImage()).release();
|
||||
return ocrBoxList;
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<OcrBox> detect(Image image){
|
||||
List<Image> imageList = Collections.singletonList(image);
|
||||
List<List<OcrBox>> result = batchDetectDJLImage(imageList);
|
||||
return result.get(0);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void detectAndDraw(String imagePath, String outputPath) {
|
||||
if(!FileUtils.isFileExists(imagePath)){
|
||||
throw new OcrException("图像文件不存在");
|
||||
}
|
||||
try {
|
||||
Image img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
|
||||
List<OcrBox> boxList = detect(img);
|
||||
if(Objects.isNull(boxList) || boxList.isEmpty()){
|
||||
throw new OcrException("未检测到文字");
|
||||
}
|
||||
OcrUtils.drawRect((Mat)img.getWrappedImage(), boxList);
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@Override
|
||||
public List<OcrBox> detect(BufferedImage image) {
|
||||
if(!ImageUtils.isImageValid(image)){
|
||||
throw new OcrException("图像无效");
|
||||
}
|
||||
Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image));
|
||||
List<OcrBox> ocrBoxList = detect(img);
|
||||
((Mat)img.getWrappedImage()).release();
|
||||
return ocrBoxList;
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<OcrBox> detect(byte[] imageData) {
|
||||
if(Objects.isNull(imageData)){
|
||||
throw new OcrException("图像无效");
|
||||
}
|
||||
try {
|
||||
BufferedImage image = ImageIO.read(new ByteArrayInputStream(imageData));
|
||||
return detect(image);
|
||||
} catch (IOException e) {
|
||||
throw new OcrException("错误的图像", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public BufferedImage detectAndDraw(BufferedImage sourceImage) {
|
||||
if(!ImageUtils.isImageValid(sourceImage)){
|
||||
throw new OcrException("图像无效");
|
||||
}
|
||||
Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(sourceImage));
|
||||
List<OcrBox> ocrBoxList = detect(img);
|
||||
if(Objects.isNull(ocrBoxList) || ocrBoxList.isEmpty()){
|
||||
throw new OcrException("未检测到文字");
|
||||
}
|
||||
OcrUtils.drawRect((Mat)img.getWrappedImage(), ocrBoxList);
|
||||
try {
|
||||
ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
|
||||
// 调用 save 方法将 Image 写入字节流
|
||||
img.save(outputStream, "png");
|
||||
// 将字节流转换为 BufferedImage
|
||||
byte[] imageBytes = outputStream.toByteArray();
|
||||
return ImageIO.read(new ByteArrayInputStream(imageBytes));
|
||||
} catch (IOException e) {
|
||||
throw new OcrException("导出图片失败", e);
|
||||
} finally {
|
||||
if (img != null){
|
||||
((Mat) img.getWrappedImage()).release();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<List<OcrBox>> batchDetect(List<BufferedImage> imageList) {
|
||||
List<Image> 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<List<OcrBox>> batchDetectDJLImage(List<Image> imageList) {
|
||||
if(!ImageUtils.isAllImageSizeEqual(imageList)){
|
||||
throw new OcrException("图片尺寸不一致");
|
||||
}
|
||||
Predictor<Image, NDList> predictor = null;
|
||||
try (NDManager manager = NDManager.newBaseManager()) {
|
||||
predictor = detPredictorPool.borrowObject();
|
||||
List<NDList> 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public GenericObjectPool<Predictor<Image, NDList>> getPool() {
|
||||
return detPredictorPool;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() throws Exception {
|
||||
try {
|
||||
if (detPredictorPool != null) {
|
||||
detPredictorPool.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 predictorPool 失败", e);
|
||||
}
|
||||
try {
|
||||
if (detectionModel != null) {
|
||||
detectionModel.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 model 失败", e);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
@@ -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<Image, NDList> createCriteria(OcrDetModelConfig config) {
|
||||
Device device = null;
|
||||
if(!Objects.isNull(config.getDevice())){
|
||||
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
|
||||
}
|
||||
Criteria<Image, NDList> criteria = null;
|
||||
ConcurrentHashMap params = new ConcurrentHashMap<String, String>();
|
||||
params.putAll(config.getCustomParams());
|
||||
if(StringUtils.isNotBlank(config.getBatchifier())){
|
||||
params.put("batchifier", config.getBatchifier());
|
||||
}
|
||||
if(config.getModelEnum() == 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;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,536 @@
|
||||
package cn.smartjavaai.ocr.model.common.detect.translator;
|
||||
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.modality.cv.util.NDImageUtils;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDArrays;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.ndarray.NDManager;
|
||||
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 cn.smartjavaai.ocr.opencv.OcrNDArrayUtils;
|
||||
import org.opencv.core.*;
|
||||
import org.opencv.imgproc.Imgproc;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
/**
|
||||
* 文字检测前后处理
|
||||
*
|
||||
* @author Calvin
|
||||
* @mail 179209347@qq.com
|
||||
* @website www.aias.top
|
||||
*/
|
||||
public class PPOCRDetTranslator implements Translator<Image, NDList> {
|
||||
// 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;
|
||||
|
||||
private String batchifier;
|
||||
|
||||
public PPOCRDetTranslator(Map<String, ?> arguments) {
|
||||
limit_side_len =
|
||||
arguments.containsKey("limit_side_len")
|
||||
? Integer.parseInt(arguments.get("limit_side_len").toString())
|
||||
: 960;
|
||||
max_candidates =
|
||||
arguments.containsKey("max_candidates")
|
||||
? Integer.parseInt(arguments.get("max_candidates").toString())
|
||||
: 1000;
|
||||
min_size =
|
||||
arguments.containsKey("min_size")
|
||||
? Integer.parseInt(arguments.get("min_size").toString())
|
||||
: 3;
|
||||
box_thresh =
|
||||
arguments.containsKey("box_thresh")
|
||||
? Float.parseFloat(arguments.get("box_thresh").toString())
|
||||
: 0.6f; // 0.5f
|
||||
unclip_ratio =
|
||||
arguments.containsKey("unclip_ratio")
|
||||
? Float.parseFloat(arguments.get("unclip_ratio").toString())
|
||||
: 1.6f;
|
||||
|
||||
batchifier = arguments.containsKey("batchifier")
|
||||
? arguments.get("batchifier").toString()
|
||||
: "stack";
|
||||
}
|
||||
|
||||
@Override
|
||||
public NDList processOutput(TranslatorContext ctx, NDList list) {
|
||||
NDManager manager = ctx.getNDManager();
|
||||
NDArray pred = list.get(0);
|
||||
pred = pred.squeeze();
|
||||
NDArray segmentation = pred.gt(thresh); // thresh=0.3 .mul(255f)
|
||||
|
||||
segmentation = segmentation.toType(DataType.UINT8, true);
|
||||
Shape shape = segmentation.getShape();
|
||||
int rows = (int) shape.get(0);
|
||||
int cols = (int) shape.get(1);
|
||||
|
||||
Mat newMask = new Mat();
|
||||
if (this.use_dilation) {
|
||||
Mat mask = new Mat();
|
||||
//convert from NDArray to Mat
|
||||
Mat srcMat = OcrNDArrayUtils.uint8NDArrayToMat(segmentation);
|
||||
// size 越小,腐蚀的单位越小,图片越接近原图
|
||||
// Mat dilation_kernel = Imgproc.getStructuringElement(Imgproc.MORPH_RECT, new Size(2, 2));
|
||||
Mat dilation_kernel = OcrNDArrayUtils.uint8ArrayToMat(new byte[][]{{1, 1}, {1, 1}});
|
||||
/**
|
||||
* 膨胀说明: 图像的一部分区域与指定的核进行卷积, 求核的最`大`值并赋值给指定区域。 膨胀可以理解为图像中`高亮区域`的'领域扩大'。
|
||||
* 意思是高亮部分会侵蚀不是高亮的部分,使高亮部分越来越多。
|
||||
*/
|
||||
Imgproc.dilate(srcMat, mask, dilation_kernel);
|
||||
//destination Matrix
|
||||
Scalar scalar = new Scalar(255);
|
||||
Core.multiply(mask, scalar, newMask);
|
||||
// release Mat
|
||||
mask.release();
|
||||
srcMat.release();
|
||||
dilation_kernel.release();
|
||||
} else {
|
||||
Mat srcMat = OcrNDArrayUtils.uint8NDArrayToMat(segmentation);
|
||||
//destination Matrix
|
||||
Scalar scalar = new Scalar(255);
|
||||
Core.multiply(srcMat, scalar, newMask);
|
||||
// release Mat
|
||||
srcMat.release();
|
||||
}
|
||||
|
||||
NDArray boxes = boxes_from_bitmap(manager, pred, newMask);
|
||||
|
||||
//boxes[:, :, 0] = boxes[:, :, 0] / ratio_w
|
||||
NDArray boxes1 = boxes.get(":, :, 0").div(ratio_w);
|
||||
boxes.set(new NDIndex(":, :, 0"), boxes1);
|
||||
//boxes[:, :, 1] = boxes[:, :, 1] / ratio_h
|
||||
NDArray boxes2 = boxes.get(":, :, 1").div(ratio_h);
|
||||
boxes.set(new NDIndex(":, :, 1"), boxes2);
|
||||
|
||||
NDList dt_boxes = this.filter_tag_det_res(boxes);
|
||||
|
||||
dt_boxes.detach();
|
||||
|
||||
// release Mat
|
||||
newMask.release();
|
||||
|
||||
return dt_boxes;
|
||||
}
|
||||
|
||||
|
||||
private NDList filter_tag_det_res(NDArray dt_boxes) {
|
||||
NDList boxesList = new NDList();
|
||||
|
||||
int num = (int) dt_boxes.getShape().get(0);
|
||||
for (int i = 0; i < num; i++) {
|
||||
NDArray box = dt_boxes.get(i);
|
||||
box = order_points_clockwise(box);
|
||||
box = clip_det_res(box);
|
||||
float[] box0 = box.get(0).toFloatArray();
|
||||
float[] box1 = box.get(1).toFloatArray();
|
||||
float[] box3 = box.get(3).toFloatArray();
|
||||
int rect_width = (int) Math.sqrt(Math.pow(box1[0] - box0[0], 2) + Math.pow(box1[1] - box0[1], 2));
|
||||
int rect_height = (int) Math.sqrt(Math.pow(box3[0] - box0[0], 2) + Math.pow(box3[1] - box0[1], 2));
|
||||
if (rect_width <= 3 || rect_height <= 3)
|
||||
continue;
|
||||
boxesList.add(box);
|
||||
}
|
||||
|
||||
return boxesList;
|
||||
}
|
||||
|
||||
private NDArray clip_det_res(NDArray points) {
|
||||
for (int i = 0; i < points.getShape().get(0); i++) {
|
||||
int value = Math.max((int) points.get(i, 0).toFloatArray()[0], 0);
|
||||
value = Math.min(value, img_width - 1);
|
||||
points.set(new NDIndex(i + ",0"), value);
|
||||
value = Math.max((int) points.get(i, 1).toFloatArray()[0], 0);
|
||||
value = Math.min(value, img_height - 1);
|
||||
points.set(new NDIndex(i + ",1"), value);
|
||||
}
|
||||
|
||||
return points;
|
||||
}
|
||||
|
||||
/**
|
||||
* sort the points based on their x-coordinates
|
||||
* 顺时针
|
||||
*
|
||||
* @param pts
|
||||
* @return
|
||||
*/
|
||||
|
||||
private NDArray order_points_clockwise(NDArray pts) {
|
||||
NDList list = new NDList();
|
||||
long[] indexes = pts.get(":, 0").argSort().toLongArray();
|
||||
|
||||
// grab the left-most and right-most points from the sorted
|
||||
// x-roodinate points
|
||||
Shape s1 = pts.getShape();
|
||||
NDArray leftMost1 = pts.get(indexes[0] + ",:");
|
||||
NDArray leftMost2 = pts.get(indexes[1] + ",:");
|
||||
NDArray leftMost = leftMost1.concat(leftMost2).reshape(2, 2);
|
||||
NDArray rightMost1 = pts.get(indexes[2] + ",:");
|
||||
NDArray rightMost2 = pts.get(indexes[3] + ",:");
|
||||
NDArray rightMost = rightMost1.concat(rightMost2).reshape(2, 2);
|
||||
|
||||
// now, sort the left-most coordinates according to their
|
||||
// y-coordinates so we can grab the top-left and bottom-left
|
||||
// points, respectively
|
||||
indexes = leftMost.get(":, 1").argSort().toLongArray();
|
||||
NDArray lt = leftMost.get(indexes[0] + ",:");
|
||||
NDArray lb = leftMost.get(indexes[1] + ",:");
|
||||
indexes = rightMost.get(":, 1").argSort().toLongArray();
|
||||
NDArray rt = rightMost.get(indexes[0] + ",:");
|
||||
NDArray rb = rightMost.get(indexes[1] + ",:");
|
||||
|
||||
list.add(lt);
|
||||
list.add(rt);
|
||||
list.add(rb);
|
||||
list.add(lb);
|
||||
|
||||
NDArray rect = NDArrays.concat(list).reshape(4, 2);
|
||||
return rect;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get boxes from the binarized image predicted by DB
|
||||
*
|
||||
* @param manager
|
||||
* @param pred the binarized image predicted by DB.
|
||||
* @param bitmap new 'pred' after threshold filtering.
|
||||
*/
|
||||
private NDArray boxes_from_bitmap(NDManager manager, NDArray pred, Mat bitmap) {
|
||||
int dest_height = (int) pred.getShape().get(0);
|
||||
int dest_width = (int) pred.getShape().get(1);
|
||||
int height = bitmap.rows();
|
||||
int width = bitmap.cols();
|
||||
|
||||
List<MatOfPoint> contours = new ArrayList<>();
|
||||
Mat hierarchy = new Mat();
|
||||
// 寻找轮廓
|
||||
Imgproc.findContours(
|
||||
bitmap,
|
||||
contours,
|
||||
hierarchy,
|
||||
Imgproc.RETR_LIST,
|
||||
Imgproc.CHAIN_APPROX_SIMPLE);
|
||||
|
||||
int num_contours = Math.min(contours.size(), max_candidates);
|
||||
NDList boxList = new NDList();
|
||||
float[] scores = new float[num_contours];
|
||||
|
||||
for (int index = 0; index < num_contours; index++) {
|
||||
MatOfPoint contour = contours.get(index);
|
||||
MatOfPoint2f newContour = new MatOfPoint2f(contour.toArray());
|
||||
float[][] pointsArr = new float[4][2];
|
||||
int sside = get_mini_boxes(newContour, pointsArr);
|
||||
if (sside < this.min_size)
|
||||
continue;
|
||||
NDArray points = manager.create(pointsArr);
|
||||
float score = box_score_fast(manager, pred, points);
|
||||
if (score < this.box_thresh)
|
||||
continue;
|
||||
|
||||
NDArray box = unclip(manager, points); // TODO get_mini_boxes(box)
|
||||
|
||||
// box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
|
||||
NDArray boxes1 = box.get(":,0").div(width).mul(dest_width).round().clip(0, dest_width);
|
||||
box.set(new NDIndex(":, 0"), boxes1);
|
||||
// box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
|
||||
NDArray boxes2 = box.get(":,1").div(height).mul(dest_height).round().clip(0, dest_height);
|
||||
box.set(new NDIndex(":, 1"), boxes2);
|
||||
|
||||
boxList.add(box);
|
||||
scores[index] = score;
|
||||
|
||||
// release memory
|
||||
contour.release();
|
||||
newContour.release();
|
||||
}
|
||||
|
||||
NDArray boxes = NDArrays.stack(boxList);
|
||||
|
||||
// release
|
||||
hierarchy.release();
|
||||
|
||||
return boxes;
|
||||
}
|
||||
|
||||
/**
|
||||
* Shrink or expand the boxaccording to 'unclip_ratio'
|
||||
*
|
||||
* @param points The predicted box.
|
||||
* @return uncliped box
|
||||
*/
|
||||
private NDArray unclip(NDManager manager, NDArray points) {
|
||||
points = order_points_clockwise(points);
|
||||
float[] pointsArr = points.toFloatArray();
|
||||
float[] lt = java.util.Arrays.copyOfRange(pointsArr, 0, 2);
|
||||
float[] lb = java.util.Arrays.copyOfRange(pointsArr, 6, 8);
|
||||
|
||||
float[] rt = java.util.Arrays.copyOfRange(pointsArr, 2, 4);
|
||||
float[] rb = java.util.Arrays.copyOfRange(pointsArr, 4, 6);
|
||||
|
||||
float width = distance(lt, rt);
|
||||
float height = distance(lt, lb);
|
||||
|
||||
if (width > height) {
|
||||
float k = (lt[1] - rt[1]) / (lt[0] - rt[0]); // y = k * x + b
|
||||
|
||||
float delta_dis = height;
|
||||
float delta_x = (float) Math.sqrt((delta_dis * delta_dis) / (k * k + 1));
|
||||
float delta_y = Math.abs(k * delta_x);
|
||||
|
||||
if (k > 0) {
|
||||
pointsArr[0] = lt[0] - delta_x + delta_y;
|
||||
pointsArr[1] = lt[1] - delta_y - delta_x;
|
||||
pointsArr[2] = rt[0] + delta_x + delta_y;
|
||||
pointsArr[3] = rt[1] + delta_y - delta_x;
|
||||
|
||||
pointsArr[4] = rb[0] + delta_x - delta_y;
|
||||
pointsArr[5] = rb[1] + delta_y + delta_x;
|
||||
pointsArr[6] = lb[0] - delta_x - delta_y;
|
||||
pointsArr[7] = lb[1] - delta_y + delta_x;
|
||||
} else {
|
||||
pointsArr[0] = lt[0] - delta_x - delta_y;
|
||||
pointsArr[1] = lt[1] + delta_y - delta_x;
|
||||
pointsArr[2] = rt[0] + delta_x - delta_y;
|
||||
pointsArr[3] = rt[1] - delta_y - delta_x;
|
||||
|
||||
pointsArr[4] = rb[0] + delta_x + delta_y;
|
||||
pointsArr[5] = rb[1] - delta_y + delta_x;
|
||||
pointsArr[6] = lb[0] - delta_x + delta_y;
|
||||
pointsArr[7] = lb[1] + delta_y + delta_x;
|
||||
}
|
||||
} else {
|
||||
float k = (lt[1] - rt[1]) / (lt[0] - rt[0]); // y = k * x + b
|
||||
|
||||
float delta_dis = width;
|
||||
float delta_y = (float) Math.sqrt((delta_dis * delta_dis) / (k * k + 1));
|
||||
float delta_x = Math.abs(k * delta_y);
|
||||
|
||||
if (k > 0) {
|
||||
pointsArr[0] = lt[0] + delta_x - delta_y;
|
||||
pointsArr[1] = lt[1] - delta_y - delta_x;
|
||||
pointsArr[2] = rt[0] + delta_x + delta_y;
|
||||
pointsArr[3] = rt[1] - delta_y + delta_x;
|
||||
|
||||
pointsArr[4] = rb[0] - delta_x + delta_y;
|
||||
pointsArr[5] = rb[1] + delta_y + delta_x;
|
||||
pointsArr[6] = lb[0] - delta_x - delta_y;
|
||||
pointsArr[7] = lb[1] + delta_y - delta_x;
|
||||
} else {
|
||||
pointsArr[0] = lt[0] - delta_x - delta_y;
|
||||
pointsArr[1] = lt[1] - delta_y + delta_x;
|
||||
pointsArr[2] = rt[0] - delta_x + delta_y;
|
||||
pointsArr[3] = rt[1] - delta_y - delta_x;
|
||||
|
||||
pointsArr[4] = rb[0] + delta_x + delta_y;
|
||||
pointsArr[5] = rb[1] + delta_y - delta_x;
|
||||
pointsArr[6] = lb[0] + delta_x - delta_y;
|
||||
pointsArr[7] = lb[1] + delta_y + delta_x;
|
||||
}
|
||||
}
|
||||
points = manager.create(pointsArr).reshape(4, 2);
|
||||
|
||||
return points;
|
||||
}
|
||||
|
||||
private float distance(float[] point1, float[] point2) {
|
||||
float disX = point1[0] - point2[0];
|
||||
float disY = point1[1] - point2[1];
|
||||
float dis = (float) Math.sqrt(disX * disX + disY * disY);
|
||||
return dis;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get boxes from the contour or box.
|
||||
*
|
||||
* @param contour The predicted contour.
|
||||
* @param pointsArr The predicted box.
|
||||
* @return smaller side of box
|
||||
*/
|
||||
private int get_mini_boxes(MatOfPoint2f contour, float[][] pointsArr) {
|
||||
// https://blog.csdn.net/qq_37385726/article/details/82313558
|
||||
// bounding_box[1] - rect 返回矩形的长和宽
|
||||
RotatedRect rect = Imgproc.minAreaRect(contour);
|
||||
Mat points = new Mat();
|
||||
Imgproc.boxPoints(rect, points);
|
||||
|
||||
float[][] fourPoints = new float[4][2];
|
||||
for (int row = 0; row < 4; row++) {
|
||||
fourPoints[row][0] = (float) points.get(row, 0)[0];
|
||||
fourPoints[row][1] = (float) points.get(row, 1)[0];
|
||||
}
|
||||
|
||||
float[] tmpPoint = new float[2];
|
||||
for (int i = 0; i < 4; i++) {
|
||||
for (int j = i + 1; j < 4; j++) {
|
||||
if (fourPoints[j][0] < fourPoints[i][0]) {
|
||||
tmpPoint[0] = fourPoints[i][0];
|
||||
tmpPoint[1] = fourPoints[i][1];
|
||||
fourPoints[i][0] = fourPoints[j][0];
|
||||
fourPoints[i][1] = fourPoints[j][1];
|
||||
fourPoints[j][0] = tmpPoint[0];
|
||||
fourPoints[j][1] = tmpPoint[1];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int index_1 = 0;
|
||||
int index_2 = 1;
|
||||
int index_3 = 2;
|
||||
int index_4 = 3;
|
||||
|
||||
if (fourPoints[1][1] > fourPoints[0][1]) {
|
||||
index_1 = 0;
|
||||
index_4 = 1;
|
||||
} else {
|
||||
index_1 = 1;
|
||||
index_4 = 0;
|
||||
}
|
||||
|
||||
if (fourPoints[3][1] > fourPoints[2][1]) {
|
||||
index_2 = 2;
|
||||
index_3 = 3;
|
||||
} else {
|
||||
index_2 = 3;
|
||||
index_3 = 2;
|
||||
}
|
||||
|
||||
pointsArr[0] = fourPoints[index_1];
|
||||
pointsArr[1] = fourPoints[index_2];
|
||||
pointsArr[2] = fourPoints[index_3];
|
||||
pointsArr[3] = fourPoints[index_4];
|
||||
|
||||
int height = rect.boundingRect().height;
|
||||
int width = rect.boundingRect().width;
|
||||
int sside = Math.min(height, width);
|
||||
|
||||
// release
|
||||
points.release();
|
||||
|
||||
return sside;
|
||||
}
|
||||
|
||||
/**
|
||||
* Calculate the score of box.
|
||||
*
|
||||
* @param bitmap The binarized image predicted by DB.
|
||||
* @param points The predicted box
|
||||
* @return
|
||||
*/
|
||||
private float box_score_fast(NDManager manager, NDArray bitmap, NDArray points) {
|
||||
NDArray box = points.get(":");
|
||||
long h = bitmap.getShape().get(0);
|
||||
long w = bitmap.getShape().get(1);
|
||||
// xmin = np.clip(np.floor(box[:, 0].min()).astype(np.int), 0, w - 1)
|
||||
int xmin = box.get(":, 0").min().floor().clip(0, w - 1).toType(DataType.INT32, true).toIntArray()[0];
|
||||
int xmax = box.get(":, 0").max().ceil().clip(0, w - 1).toType(DataType.INT32, true).toIntArray()[0];
|
||||
int ymin = box.get(":, 1").min().floor().clip(0, h - 1).toType(DataType.INT32, true).toIntArray()[0];
|
||||
int ymax = box.get(":, 1").max().ceil().clip(0, h - 1).toType(DataType.INT32, true).toIntArray()[0];
|
||||
|
||||
NDArray mask = manager.zeros(new Shape(ymax - ymin + 1, xmax - xmin + 1), DataType.UINT8);
|
||||
|
||||
box.set(new NDIndex(":, 0"), box.get(":, 0").sub(xmin));
|
||||
box.set(new NDIndex(":, 1"), box.get(":, 1").sub(ymin));
|
||||
|
||||
//mask - convert from NDArray to Mat
|
||||
Mat maskMat = OcrNDArrayUtils.uint8NDArrayToMat(mask);
|
||||
|
||||
//mask - convert from NDArray to Mat - 4 rows, 2 cols
|
||||
Mat boxMat = OcrNDArrayUtils.floatNDArrayToMat(box, CvType.CV_32S);
|
||||
|
||||
// boxMat.reshape(1, new int[]{1, 4, 2});
|
||||
List<MatOfPoint> pts = new ArrayList<>();
|
||||
MatOfPoint matOfPoint = OcrNDArrayUtils.matToMatOfPoint(boxMat); // new MatOfPoint(boxMat);
|
||||
pts.add(matOfPoint);
|
||||
Imgproc.fillPoly(maskMat, pts, new Scalar(1));
|
||||
|
||||
|
||||
NDArray subBitMap = bitmap.get(ymin + ":" + (ymax + 1) + "," + xmin + ":" + (xmax + 1));
|
||||
Mat bitMapMat = OcrNDArrayUtils.floatNDArrayToMat(subBitMap);
|
||||
|
||||
Scalar score = Core.mean(bitMapMat, maskMat);
|
||||
float scoreValue = (float) score.val[0];
|
||||
// release
|
||||
maskMat.release();
|
||||
boxMat.release();
|
||||
bitMapMat.release();
|
||||
|
||||
return scoreValue;
|
||||
}
|
||||
|
||||
@Override
|
||||
public NDList processInput(TranslatorContext ctx, Image input) {
|
||||
NDArray img = input.toNDArray(ctx.getNDManager());
|
||||
int h = input.getHeight();
|
||||
int w = input.getWidth();
|
||||
img_height = h;
|
||||
img_width = w;
|
||||
|
||||
// limit the max side
|
||||
float ratio = 1.0f;
|
||||
if (Math.max(h, w) > limit_side_len) {
|
||||
if (h > w) {
|
||||
ratio = (float) limit_side_len / (float) h;
|
||||
} else {
|
||||
ratio = (float) limit_side_len / (float) w;
|
||||
}
|
||||
}
|
||||
|
||||
int resize_h = (int) (h * ratio);
|
||||
int resize_w = (int) (w * ratio);
|
||||
|
||||
resize_h = Math.round((float) resize_h / 32f) * 32;
|
||||
resize_w = Math.round((float) resize_w / 32f) * 32;
|
||||
|
||||
ratio_h = resize_h / (float) h;
|
||||
ratio_w = resize_w / (float) w;
|
||||
|
||||
img = NDImageUtils.resize(img, resize_w, resize_h);
|
||||
|
||||
img = NDImageUtils.toTensor(img);
|
||||
|
||||
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 Batchifier getBatchifier() {
|
||||
return Batchifier.fromString(batchifier);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package cn.smartjavaai.ocr.model.common.direction;
|
||||
|
||||
import ai.djl.inference.Predictor;
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.ndarray.NDManager;
|
||||
import cn.smartjavaai.ocr.config.DirectionModelConfig;
|
||||
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.apache.commons.pool2.impl.GenericObjectPool;
|
||||
import org.opencv.core.Mat;
|
||||
|
||||
import java.awt.image.BufferedImage;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* OCR 文本方向分类模型
|
||||
* @author dwj
|
||||
*/
|
||||
public interface OcrDirectionModel extends AutoCloseable{
|
||||
|
||||
default void setTextDetModel(OcrCommonDetModel detModel){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
default OcrCommonDetModel getTextDetModel(){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 加载模型
|
||||
* @param config
|
||||
*/
|
||||
void loadModel(DirectionModelConfig config); // 加载模型
|
||||
|
||||
/**
|
||||
* 文本方向检测
|
||||
* @param imagePath 图片路径
|
||||
* @return
|
||||
*/
|
||||
default List<OcrItem> detect(String imagePath) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 文本方向检测
|
||||
* @param image BufferedImage
|
||||
* @return
|
||||
*/
|
||||
default List<OcrItem> detect(BufferedImage image) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 文本方向检测
|
||||
* @param imageData 图片字节数组
|
||||
* @return
|
||||
*/
|
||||
default List<OcrItem> detect(byte[] imageData) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 文本方向检测
|
||||
* @param image
|
||||
* @return
|
||||
*/
|
||||
default List<OcrItem> detect(Image image) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 文本方向检测(基于检测结果)
|
||||
* @param boxList
|
||||
* @param srcMat
|
||||
* @param manager
|
||||
* @return
|
||||
*/
|
||||
default List<OcrItem> detect(List<OcrBox> boxList, Mat srcMat) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
default List<List<OcrItem>> batchDetect(List<List<OcrBox>> boxList, List<Mat> srcMatList) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 检测并绘制结果
|
||||
* @param imagePath 图片输入路径(包含文件名称)
|
||||
* @param outputPath 图片输出路径(包含文件名称)
|
||||
*/
|
||||
default void detectAndDraw(String imagePath, String outputPath) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 检测并绘制结果
|
||||
* @param sourceImage
|
||||
* @return
|
||||
*/
|
||||
default BufferedImage detectAndDraw(BufferedImage sourceImage){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
default GenericObjectPool<Predictor<Image, DirectionInfo>> getPool() {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,396 @@
|
||||
package cn.smartjavaai.ocr.model.common.direction;
|
||||
|
||||
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.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;
|
||||
import cn.smartjavaai.common.utils.OpenCVUtils;
|
||||
import cn.smartjavaai.ocr.config.DirectionModelConfig;
|
||||
import cn.smartjavaai.ocr.entity.*;
|
||||
import cn.smartjavaai.ocr.enums.AngleEnum;
|
||||
import cn.smartjavaai.ocr.exception.OcrException;
|
||||
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.utils.OcrUtils;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
import org.apache.commons.collections.CollectionUtils;
|
||||
import org.apache.commons.lang3.StringUtils;
|
||||
import org.apache.commons.pool2.ObjectPool;
|
||||
import org.apache.commons.pool2.impl.GenericObjectPool;
|
||||
import org.opencv.core.Mat;
|
||||
|
||||
import javax.imageio.ImageIO;
|
||||
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.ArrayList;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
import java.util.Objects;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
|
||||
/**
|
||||
* PPOCRMobileV2Model 方向分类模型
|
||||
* @author dwj
|
||||
* @date 2025/4/21
|
||||
*/
|
||||
@Slf4j
|
||||
public class PPOCRMobileV2ClsModel implements OcrDirectionModel {
|
||||
|
||||
|
||||
private GenericObjectPool<Predictor<Image, DirectionInfo>> predictorPool;
|
||||
|
||||
private DirectionModelConfig config;
|
||||
|
||||
private ZooModel<Image, DirectionInfo> 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();
|
||||
ConcurrentHashMap params = new ConcurrentHashMap<String, String>();
|
||||
if(StringUtils.isNotBlank(config.getBatchifier())){
|
||||
params.put("batchifier", config.getBatchifier());
|
||||
}
|
||||
Criteria<Image, DirectionInfo> criteria = DirectionCriteriaFactory.createCriteria(config);
|
||||
try{
|
||||
model = ModelZoo.loadModel(criteria);
|
||||
// 创建池子:每个线程独享 Predictor
|
||||
this.predictorPool = new GenericObjectPool<>(new PredictorFactory<>(model));
|
||||
int predictorPoolSize = config.getPredictorPoolSize();
|
||||
if(config.getPredictorPoolSize() <= 0){
|
||||
predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
|
||||
}
|
||||
predictorPool.setMaxTotal(predictorPoolSize);
|
||||
log.debug("当前设备: " + model.getNDManager().getDevice());
|
||||
log.debug("当前引擎: " + Engine.getInstance().getEngineName());
|
||||
log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
|
||||
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
|
||||
throw new OcrException("模型加载失败", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<OcrItem> detect(String imagePath){
|
||||
if(!FileUtils.isFileExists(imagePath)){
|
||||
throw new OcrException("图像文件不存在");
|
||||
}
|
||||
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();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@Override
|
||||
public List<OcrItem> detect(Image image){
|
||||
if(Objects.isNull(textDetModel)){
|
||||
throw new OcrException("textDetModel is null");
|
||||
}
|
||||
//检测文本
|
||||
List<OcrBox> boxeList = textDetModel.detect(image);
|
||||
if(Objects.isNull(boxeList) || boxeList.isEmpty()){
|
||||
throw new OcrException("未检测到文本");
|
||||
}
|
||||
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<Image, DirectionInfo> 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<OcrItem> detect(List<OcrBox> boxList,Mat srcMat){
|
||||
if(Objects.isNull(boxList) || boxList.isEmpty()){
|
||||
throw new OcrException("boxList为空");
|
||||
}
|
||||
List<List<OcrItem>> ocrItemList = batchDetect(Collections.singletonList(boxList), Collections.singletonList(srcMat));
|
||||
if(Objects.isNull(ocrItemList) || ocrItemList.isEmpty()){
|
||||
throw new OcrException("方向检测失败");
|
||||
}
|
||||
return ocrItemList.get(0);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void detectAndDraw(String imagePath, String outputPath) {
|
||||
if(!FileUtils.isFileExists(imagePath)){
|
||||
throw new OcrException("图像文件不存在");
|
||||
}
|
||||
Image img = null;
|
||||
try {
|
||||
img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
|
||||
List<OcrItem> itemList = detect(img);
|
||||
if(Objects.isNull(itemList) || itemList.isEmpty()){
|
||||
throw new OcrException("未检测到文字");
|
||||
}
|
||||
OcrUtils.drawRectWithText((Mat) img.getWrappedImage(), itemList);
|
||||
Path output = Paths.get(outputPath);
|
||||
log.debug("Saving to {}", output.toAbsolutePath().toString());
|
||||
img.save(Files.newOutputStream(output), "png");
|
||||
} catch (IOException e) {
|
||||
throw new OcrException(e);
|
||||
} finally {
|
||||
if (img != null){
|
||||
((Mat)img.getWrappedImage()).release();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<OcrItem> detect(BufferedImage image) {
|
||||
if(!ImageUtils.isImageValid(image)){
|
||||
throw new OcrException("图像无效");
|
||||
}
|
||||
Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image));
|
||||
List<OcrItem> ocrItemList = detect(img);
|
||||
((Mat)img.getWrappedImage()).release();
|
||||
return ocrItemList;
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<OcrItem> detect(byte[] imageData) {
|
||||
if(Objects.isNull(imageData)){
|
||||
throw new OcrException("图像无效");
|
||||
}
|
||||
try {
|
||||
BufferedImage image = ImageIO.read(new ByteArrayInputStream(imageData));
|
||||
return detect(image);
|
||||
} catch (IOException e) {
|
||||
throw new OcrException("错误的图像", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public BufferedImage detectAndDraw(BufferedImage sourceImage) {
|
||||
if(!ImageUtils.isImageValid(sourceImage)){
|
||||
throw new OcrException("图像无效");
|
||||
}
|
||||
Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(sourceImage));
|
||||
List<OcrItem> ocrItemList = detect(img);
|
||||
if(Objects.isNull(ocrItemList) || ocrItemList.isEmpty()){
|
||||
throw new OcrException("未检测到文字");
|
||||
}
|
||||
OcrUtils.drawRectWithText((Mat) img.getWrappedImage(), ocrItemList);
|
||||
try {
|
||||
ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
|
||||
// 调用 save 方法将 Image 写入字节流
|
||||
img.save(outputStream, "png");
|
||||
// 将字节流转换为 BufferedImage
|
||||
byte[] imageBytes = outputStream.toByteArray();
|
||||
return ImageIO.read(new ByteArrayInputStream(imageBytes));
|
||||
} catch (IOException e) {
|
||||
throw new OcrException("导出图片失败", e);
|
||||
} finally {
|
||||
if (img != null){
|
||||
((Mat) img.getWrappedImage()).release();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<List<OcrItem>> batchDetect(List<List<OcrBox>> boxList, List<Mat> 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<OcrBox> 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<Image> imageList = new ArrayList<Image>();
|
||||
List<Boolean> isRotatedList = new ArrayList<Boolean>();
|
||||
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<List<OcrItem>> result = new ArrayList<>();
|
||||
List<DirectionInfo> directionInfos = batchDetect(imageList);
|
||||
if(CollectionUtils.isEmpty(directionInfos)){
|
||||
throw new OcrException("方向检测失败");
|
||||
}
|
||||
index = 0;
|
||||
for(int i = 0; i < srcMatList.size(); i++){
|
||||
List<OcrItem> 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<DirectionInfo> batchDetect(List<Image> imageList) {
|
||||
Predictor<Image, DirectionInfo> 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 GenericObjectPool<Predictor<Image, DirectionInfo>> getPool() {
|
||||
return predictorPool;
|
||||
}
|
||||
|
||||
@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);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<Image, DirectionInfo> createCriteria(DirectionModelConfig config) {
|
||||
Device device = null;
|
||||
if(!Objects.isNull(config.getDevice())){
|
||||
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
|
||||
}
|
||||
Criteria<Image, DirectionInfo> criteria = null;
|
||||
ConcurrentHashMap params = new ConcurrentHashMap<String, String>();
|
||||
params.putAll(config.getCustomParams());
|
||||
if(StringUtils.isNotBlank(config.getBatchifier())){
|
||||
params.put("batchifier", config.getBatchifier());
|
||||
}
|
||||
if(config.getModelEnum() == 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;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
package cn.smartjavaai.ocr.model.common.direction.translator;
|
||||
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.modality.cv.util.NDImageUtils;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.ndarray.index.NDIndex;
|
||||
import ai.djl.ndarray.types.Shape;
|
||||
import ai.djl.translate.Batchifier;
|
||||
import ai.djl.translate.Translator;
|
||||
import ai.djl.translate.TranslatorContext;
|
||||
import cn.smartjavaai.ocr.entity.DirectionInfo;
|
||||
|
||||
import java.util.Arrays;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
/**
|
||||
* 方向检测
|
||||
*
|
||||
* @author Calvin
|
||||
* @mail 179209347@qq.com
|
||||
* @website www.aias.top
|
||||
*/
|
||||
public class PpWordRotateTranslator implements Translator<Image, DirectionInfo> {
|
||||
List<String> classes = Arrays.asList("No Rotate", "Rotate");
|
||||
|
||||
private String batchifier;
|
||||
|
||||
private int resizeHeight;
|
||||
|
||||
private int resizeWidth;
|
||||
|
||||
public PpWordRotateTranslator(Map<String, ?> 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
|
||||
public DirectionInfo processOutput(TranslatorContext ctx, NDList list) {
|
||||
NDArray prob = list.singletonOrThrow();
|
||||
float[] res = prob.toFloatArray();
|
||||
int maxIndex = 0;
|
||||
if (res[1] > res[0]) {
|
||||
maxIndex = 1;
|
||||
}
|
||||
|
||||
return new DirectionInfo(classes.get(maxIndex), Double.valueOf(res[maxIndex]));
|
||||
}
|
||||
|
||||
// public NDList processInput2(TranslatorContext ctx, Image input){
|
||||
// NDArray img = input.toNDArray(ctx.getNDManager());
|
||||
// img = NDImageUtils.resize(img, 192, 48);
|
||||
// img = NDImageUtils.toTensor(img).sub(0.5F).div(0.5F);
|
||||
// img = img.expandDims(0);
|
||||
// return new NDList(new NDArray[]{img});
|
||||
// }
|
||||
|
||||
@Override
|
||||
public NDList processInput(TranslatorContext ctx, Image input) {
|
||||
NDArray img = input.toNDArray(ctx.getNDManager());
|
||||
int imgC = 3;
|
||||
int imgH = resizeHeight;
|
||||
int imgW = resizeWidth;
|
||||
|
||||
NDArray array = ctx.getNDManager().zeros(new Shape(imgC, imgH, imgW));
|
||||
|
||||
int h = input.getHeight();
|
||||
int w = input.getWidth();
|
||||
int resized_w = 0;
|
||||
|
||||
float ratio = (float) w / (float) h;
|
||||
if (Math.ceil(imgH * ratio) > imgW) {
|
||||
resized_w = imgW;
|
||||
} else {
|
||||
resized_w = (int) (Math.ceil(imgH * ratio));
|
||||
}
|
||||
|
||||
img = NDImageUtils.resize(img, resized_w, imgH);
|
||||
|
||||
img = NDImageUtils.toTensor(img).sub(0.5F).div(0.5F);
|
||||
// img = img.transpose(2, 0, 1);
|
||||
|
||||
array.set(new NDIndex(":,:,0:" + resized_w), img);
|
||||
|
||||
// array = array.expandDims(0);
|
||||
|
||||
return new NDList(new NDArray[]{array});
|
||||
}
|
||||
|
||||
@Override
|
||||
public Batchifier getBatchifier() {
|
||||
return Batchifier.fromString(batchifier);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
package cn.smartjavaai.ocr.model.common.recognize;
|
||||
|
||||
import ai.djl.inference.Predictor;
|
||||
import ai.djl.modality.cv.Image;
|
||||
import cn.smartjavaai.ocr.config.OcrDetModelConfig;
|
||||
import cn.smartjavaai.ocr.config.OcrRecModelConfig;
|
||||
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 org.apache.commons.pool2.impl.GenericObjectPool;
|
||||
|
||||
import java.awt.image.BufferedImage;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* OCR 通用识别模型
|
||||
* @author dwj
|
||||
*/
|
||||
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
|
||||
*/
|
||||
void loadModel(OcrRecModelConfig config); // 加载模型
|
||||
|
||||
/**
|
||||
* 文本识别
|
||||
* @param imagePath 图片路径
|
||||
* @return
|
||||
*/
|
||||
default OcrInfo recognize(String imagePath, OcrRecOptions options) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 文本识别
|
||||
* @param image
|
||||
* @return
|
||||
*/
|
||||
default OcrInfo recognize(Image image, OcrRecOptions options) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 文本检测
|
||||
* @param image BufferedImage
|
||||
* @return
|
||||
*/
|
||||
default OcrInfo recognize(BufferedImage image, OcrRecOptions options) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 文本检测
|
||||
* @param imageData 图片字节数组
|
||||
* @return
|
||||
*/
|
||||
default OcrInfo recognize(byte[] imageData, OcrRecOptions options) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 识别并绘制结果
|
||||
* @param imagePath
|
||||
* @param outputPath
|
||||
*/
|
||||
default void recognizeAndDraw(String imagePath, String outputPath, int fontSize, OcrRecOptions options) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 识别并绘制结果
|
||||
* @param sourceImage
|
||||
* @return
|
||||
*/
|
||||
default BufferedImage recognizeAndDraw(BufferedImage sourceImage, int fontSize, OcrRecOptions options){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
/**
|
||||
* 识别并绘制Base64结果
|
||||
* @param imageData 图片字节数组
|
||||
* @return
|
||||
*/
|
||||
default String recognizeAndDrawToBase64(byte[] imageData, int fontSize, OcrRecOptions options){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 识别并绘制结果
|
||||
* @param imageData 图片字节数组
|
||||
* @return
|
||||
*/
|
||||
default OcrInfo recognizeAndDraw(byte[] imageData, int fontSize, OcrRecOptions options){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
default List<OcrInfo> batchRecognize(List<BufferedImage> imageList, OcrRecOptions options) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
default List<OcrInfo> batchRecognizeDJLImage(List<Image> imageList, OcrRecOptions options) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
default GenericObjectPool<Predictor<Image, String>> getPool() {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,500 @@
|
||||
package cn.smartjavaai.ocr.model.common.recognize;
|
||||
|
||||
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.NDArray;
|
||||
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.hutool.core.img.ImgUtil;
|
||||
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.OcrRecModelConfig;
|
||||
import cn.smartjavaai.ocr.config.OcrRecOptions;
|
||||
import cn.smartjavaai.ocr.entity.*;
|
||||
import cn.smartjavaai.ocr.exception.OcrException;
|
||||
import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel;
|
||||
import cn.smartjavaai.ocr.model.common.direction.OcrDirectionModel;
|
||||
import cn.smartjavaai.ocr.model.common.recognize.criteria.OcrCommonRecCriterialFactory;
|
||||
import cn.smartjavaai.ocr.utils.OcrUtils;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
import org.apache.commons.collections.CollectionUtils;
|
||||
import org.apache.commons.lang3.StringUtils;
|
||||
import org.apache.commons.pool2.impl.GenericObjectPool;
|
||||
import org.opencv.core.Mat;
|
||||
|
||||
import javax.imageio.ImageIO;
|
||||
import java.awt.image.BufferedImage;
|
||||
import java.io.ByteArrayInputStream;
|
||||
import java.io.ByteArrayOutputStream;
|
||||
import java.io.IOException;
|
||||
import java.nio.file.Paths;
|
||||
import java.util.*;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
/**
|
||||
* PPOCRV5 识别模型
|
||||
*
|
||||
* @author dwj
|
||||
*/
|
||||
@Slf4j
|
||||
public class OcrCommonRecModelImpl implements OcrCommonRecModel {
|
||||
|
||||
private GenericObjectPool<Predictor<Image, String>> recPredictorPool;
|
||||
|
||||
private OcrRecModelConfig config;
|
||||
|
||||
private ZooModel<Image, String> recognitionModel;
|
||||
|
||||
private OcrDirectionModel directionModel;
|
||||
|
||||
private OcrCommonDetModel textDetModel;
|
||||
|
||||
@Override
|
||||
public void loadModel(OcrRecModelConfig config) {
|
||||
if (StringUtils.isBlank(config.getRecModelPath())) {
|
||||
throw new OcrException("recModelPath is null");
|
||||
}
|
||||
this.config = config;
|
||||
this.directionModel = config.getDirectionModel();
|
||||
this.textDetModel = config.getTextDetModel();
|
||||
//初始化 识别Criteria
|
||||
Criteria<Image, String> recCriteria = OcrCommonRecCriterialFactory.createCriteria(config);
|
||||
try {
|
||||
recognitionModel = ModelZoo.loadModel(recCriteria);
|
||||
this.recPredictorPool = new GenericObjectPool<>(new PredictorFactory<>(recognitionModel));
|
||||
int predictorPoolSize = config.getPredictorPoolSize();
|
||||
if (config.getPredictorPoolSize() <= 0) {
|
||||
predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
|
||||
}
|
||||
recPredictorPool.setMaxTotal(predictorPoolSize);
|
||||
log.debug("当前设备: " + recognitionModel.getNDManager().getDevice());
|
||||
log.debug("当前引擎: " + Engine.getInstance().getEngineName());
|
||||
log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
|
||||
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
|
||||
throw new OcrException("识别模型加载失败", e);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
|
||||
@Override
|
||||
public OcrInfo recognize(String imagePath, OcrRecOptions options) {
|
||||
if (StringUtils.isBlank(config.getRecModelPath())) {
|
||||
throw new OcrException("recModelPath为空,无法识别");
|
||||
}
|
||||
if (!FileUtils.isFileExists(imagePath)) {
|
||||
throw new OcrException("图像文件不存在");
|
||||
}
|
||||
Image img = null;
|
||||
try {
|
||||
img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
|
||||
return recognize(img, options);
|
||||
} catch (IOException e) {
|
||||
throw new OcrException("无效的图片", e);
|
||||
} finally {
|
||||
if (img != null) {
|
||||
((Mat) img.getWrappedImage()).release();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @param image
|
||||
* @param options
|
||||
* @return
|
||||
*/
|
||||
@Override
|
||||
public OcrInfo recognize(Image image, OcrRecOptions options) {
|
||||
List<OcrInfo> result = batchRecognizeDJLImage(Collections.singletonList(image), options);
|
||||
if (CollectionUtils.isEmpty(result)) {
|
||||
throw new OcrException("OCR识别结果为空");
|
||||
}
|
||||
return result.get(0);
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 批量矫正文本框
|
||||
*
|
||||
* @param boxList
|
||||
* @param srcMat
|
||||
* @param manager
|
||||
* @return
|
||||
*/
|
||||
private List<Image> batchAlign(List<OcrBox> boxList, Mat srcMat, NDManager manager) {
|
||||
List<Image> imageList = new ArrayList<>(boxList.size());
|
||||
for (int i = 0; i < boxList.size(); i++) {
|
||||
//透视变换 + 裁剪
|
||||
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) {
|
||||
//旋转图片90度
|
||||
subImg = OcrUtils.rotateImg(manager, subImg);
|
||||
//ImageUtils.saveImage(subImg, i + "rotate.png", "build/output");
|
||||
}
|
||||
imageList.add(subImg);
|
||||
}
|
||||
return imageList;
|
||||
}
|
||||
|
||||
/**
|
||||
* 批量矫正文本框
|
||||
*
|
||||
* @param itemList
|
||||
* @param srcMat
|
||||
* @param manager
|
||||
* @return
|
||||
*/
|
||||
private List<Image> batchAlignWithDirection(List<OcrItem> itemList, Mat srcMat, NDManager manager) {
|
||||
List<Image> 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<Image, String> 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<RotatedBox> rotatedBoxes, OcrRecOptions ocrRecOptions) {
|
||||
//不分行
|
||||
if (!ocrRecOptions.isEnableLineSplit()) {
|
||||
return OcrUtils.convertRotatedBoxesToOcrItems(rotatedBoxes);
|
||||
}
|
||||
//Y坐标升序排序
|
||||
List<RotatedBox> initList = new ArrayList<>();
|
||||
for (RotatedBox result : rotatedBoxes) {
|
||||
initList.add(result);
|
||||
}
|
||||
Collections.sort(initList);
|
||||
//多行文本框的集合
|
||||
List<ArrayList<RotatedBoxCompX>> lines = new ArrayList<>();
|
||||
List<RotatedBoxCompX> line = new ArrayList<>();
|
||||
RotatedBoxCompX firstBox = new RotatedBoxCompX(initList.get(0).getBox(), initList.get(0).getText());
|
||||
line.add(firstBox);
|
||||
lines.add((ArrayList) line);
|
||||
//分行判断
|
||||
for (int i = 1; i < initList.size(); i++) {
|
||||
RotatedBoxCompX tmpBox = new RotatedBoxCompX(initList.get(i).getBox(), initList.get(i).getText());
|
||||
float y1 = firstBox.getBox().toFloatArray()[1];
|
||||
float y2 = tmpBox.getBox().toFloatArray()[1];
|
||||
float dis = Math.abs(y2 - y1);
|
||||
if (dis < 20) { // 认为是同 1 行 - Considered to be in the same line
|
||||
line.add(tmpBox);
|
||||
} else { // 换行 - Line break
|
||||
firstBox = tmpBox;
|
||||
Collections.sort(line);
|
||||
line = new ArrayList<>();
|
||||
line.add(firstBox);
|
||||
lines.add((ArrayList) line);
|
||||
}
|
||||
}
|
||||
return OcrUtils.convertToOcrInfo(lines);
|
||||
}
|
||||
|
||||
|
||||
@Override
|
||||
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, options);
|
||||
if (Objects.isNull(ocrInfo) || Objects.isNull(ocrInfo.getLineList()) || ocrInfo.getLineList().isEmpty()) {
|
||||
throw new OcrException("未检测到文字");
|
||||
}
|
||||
Mat wrappedImage = (Mat) img.getWrappedImage();
|
||||
BufferedImage bufferedImage = OpenCVUtils.mat2Image(wrappedImage);
|
||||
OcrUtils.drawRectWithText(bufferedImage, ocrInfo, fontSize);
|
||||
ImageUtils.saveImage(bufferedImage, outputPath);
|
||||
wrappedImage.release();
|
||||
} catch (IOException e) {
|
||||
throw new OcrException(e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
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, options);
|
||||
((Mat) img.getWrappedImage()).release();
|
||||
return ocrInfo;
|
||||
}
|
||||
|
||||
@Override
|
||||
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, options);
|
||||
} catch (IOException e) {
|
||||
throw new OcrException("错误的图像", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
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, options);
|
||||
if (Objects.isNull(ocrInfo) || Objects.isNull(ocrInfo.getLineList()) || ocrInfo.getLineList().isEmpty()) {
|
||||
throw new OcrException("未检测到文字");
|
||||
}
|
||||
OcrUtils.drawRectWithText(sourceImage, ocrInfo, fontSize);
|
||||
return sourceImage;
|
||||
}
|
||||
|
||||
@Override
|
||||
public String recognizeAndDrawToBase64(byte[] imageData, int fontSize, OcrRecOptions options) {
|
||||
if (Objects.isNull(imageData)) {
|
||||
throw new OcrException("图像无效");
|
||||
}
|
||||
OcrInfo ocrInfo = recognize(imageData, options);
|
||||
if (Objects.isNull(ocrInfo) || Objects.isNull(ocrInfo.getLineList()) || ocrInfo.getLineList().isEmpty()) {
|
||||
throw new OcrException("未检测到文字");
|
||||
}
|
||||
try {
|
||||
BufferedImage sourceImage = ImageIO.read(new ByteArrayInputStream(imageData));
|
||||
OcrUtils.drawRectWithText(sourceImage, ocrInfo, fontSize);
|
||||
return ImgUtil.toBase64(sourceImage, "png");
|
||||
} catch (IOException e) {
|
||||
throw new OcrException("导出图片失败", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public OcrInfo recognizeAndDraw(byte[] imageData, int fontSize, OcrRecOptions options) {
|
||||
if (Objects.isNull(imageData)) {
|
||||
throw new OcrException("图像无效");
|
||||
}
|
||||
OcrInfo ocrInfo = recognize(imageData, options);
|
||||
if (Objects.isNull(ocrInfo) || Objects.isNull(ocrInfo.getLineList()) || ocrInfo.getLineList().isEmpty()) {
|
||||
throw new OcrException("未检测到文字");
|
||||
}
|
||||
try {
|
||||
BufferedImage sourceImage = ImageIO.read(new ByteArrayInputStream(imageData));
|
||||
OcrUtils.drawRectWithText(sourceImage, ocrInfo, fontSize);
|
||||
ocrInfo.setBase64Img(ImgUtil.toBase64(sourceImage, "png"));
|
||||
return ocrInfo;
|
||||
} catch (IOException e) {
|
||||
throw new OcrException("导出图片失败", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<OcrInfo> batchRecognize(List<BufferedImage> imageList, OcrRecOptions options) {
|
||||
List<Image> 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<OcrInfo> batchRecognizeDJLImage(List<Image> 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<List<OcrBox>> boxeList = textDetModel.batchDetectDJLImage(imageList);
|
||||
if (CollectionUtils.isEmpty(boxeList) || boxeList.size() != imageList.size()) {
|
||||
throw new OcrException("未检测到文本");
|
||||
}
|
||||
Predictor<Image, String> predictor = null;
|
||||
List<OcrInfo> ocrInfoList = new ArrayList<OcrInfo>();
|
||||
try (NDManager manager = NDManager.newBaseManager()) {
|
||||
predictor = recPredictorPool.borrowObject();
|
||||
List<Image> allImageAlignList = new ArrayList<Image>();
|
||||
//检测方向
|
||||
if (ocrRecOptions.isEnableDirectionCorrect()) {
|
||||
if (Objects.isNull(directionModel)) {
|
||||
throw new OcrException("请配置方向模型");
|
||||
}
|
||||
List<Mat> matList = imageList.stream()
|
||||
.map(image -> (Mat) image.getWrappedImage())
|
||||
.collect(Collectors.toList());
|
||||
List<List<OcrItem>> ocrItemList = directionModel.batchDetect(boxeList, matList);
|
||||
if (CollectionUtils.isEmpty(ocrItemList) || ocrItemList.size() != imageList.size()) {
|
||||
throw new OcrException("方向检测失败");
|
||||
}
|
||||
allImageAlignList = new ArrayList<Image>();
|
||||
for (int i = 0; i < ocrItemList.size(); i++) {
|
||||
Mat srcMat = (Mat) imageList.get(i).getWrappedImage();
|
||||
List<Image> 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<Image> 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<String> textList = batchRecognize(allImageAlignList);
|
||||
int textIndex = 0;
|
||||
for (int i = 0; i < boxeList.size(); i++) {
|
||||
List<RotatedBox> 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<String> batchRecognize(List<Image> imageAlignList) {
|
||||
Predictor<Image, String> predictor = null;
|
||||
try {
|
||||
predictor = recPredictorPool.borrowObject();
|
||||
List<String> 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;
|
||||
}
|
||||
|
||||
|
||||
public GenericObjectPool<Predictor<Image, String>> getRecPredictorPool() {
|
||||
return recPredictorPool;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() throws Exception {
|
||||
try {
|
||||
if (recPredictorPool != null) {
|
||||
recPredictorPool.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 predictorPool 失败", e);
|
||||
}
|
||||
try {
|
||||
if (recognitionModel != null) {
|
||||
recognitionModel.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 model 失败", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<Image, String> createCriteria(OcrRecModelConfig config) {
|
||||
Device device = null;
|
||||
if(!Objects.isNull(config.getDevice())){
|
||||
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
|
||||
}
|
||||
Criteria<Image, String> criteria = null;
|
||||
ConcurrentHashMap params = new ConcurrentHashMap<String, String>();
|
||||
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;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
package cn.smartjavaai.ocr.model.common.recognize.translator;
|
||||
|
||||
import ai.djl.Model;
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.modality.cv.util.NDImageUtils;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.ndarray.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 java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.util.Arrays;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
/**
|
||||
* 文字识别前后处理
|
||||
*
|
||||
*/
|
||||
public class PPOCRRecTranslator implements Translator<Image, String> {
|
||||
private List<String> table;
|
||||
private final boolean use_space_char;
|
||||
|
||||
private String batchifier;
|
||||
|
||||
public PPOCRRecTranslator(Map<String, ?> 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
|
||||
public void prepare(TranslatorContext ctx) throws IOException {
|
||||
Model model = ctx.getModel();
|
||||
try (InputStream is = model.getArtifact("dict.txt").openStream()) {
|
||||
table = Utils.readLines(is, true);
|
||||
table.add(0, "blank");
|
||||
if(use_space_char){
|
||||
table.add(" ");
|
||||
table.add(" ");
|
||||
}
|
||||
else{
|
||||
table.add("");
|
||||
table.add("");
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public String processOutput(TranslatorContext ctx, NDList list) throws IOException {
|
||||
StringBuilder sb = new StringBuilder();
|
||||
NDArray tokens = list.singletonOrThrow();
|
||||
|
||||
// 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++) {
|
||||
if (indices[i] == indices[i - 1]) {
|
||||
selection[i] = false;
|
||||
}
|
||||
}
|
||||
|
||||
// 字符置信度
|
||||
// float[] probs = new float[indices.length];
|
||||
// for (int row = 0; row < indices.length; row++) {
|
||||
// NDArray value = tokens.get(0).get(new NDIndex(""+ row +":" + (row + 1) +"," + indices[row] +":" + ( indices[row] + 1)));
|
||||
// probs[row] = value.toFloatArray()[0];
|
||||
// }
|
||||
|
||||
int lastIdx = 0;
|
||||
for (int i = 0; i < indices.length; i++) {
|
||||
if (selection[i] == true && indices[i] > 0 && !(i > 0 && indices[i] == lastIdx)) {
|
||||
sb.append(table.get((int) indices[i]));
|
||||
}
|
||||
}
|
||||
return sb.toString();
|
||||
}
|
||||
|
||||
@Override
|
||||
public NDList processInput(TranslatorContext ctx, Image input) {
|
||||
NDArray img = input.toNDArray(ctx.getNDManager(), Image.Flag.COLOR);
|
||||
int imgC = 3;
|
||||
int imgH = 48;
|
||||
int imgW = 320;
|
||||
|
||||
float max_wh_ratio = (float) imgW / (float) imgH;
|
||||
|
||||
int h = input.getHeight();
|
||||
int w = input.getWidth();
|
||||
float wh_ratio = (float) w / (float) h;
|
||||
|
||||
max_wh_ratio = Math.max(max_wh_ratio,wh_ratio);
|
||||
imgW = (int)(imgH * max_wh_ratio);
|
||||
|
||||
int resized_w;
|
||||
if (Math.ceil(imgH * wh_ratio) > imgW) {
|
||||
resized_w = imgW;
|
||||
} else {
|
||||
resized_w = (int) (Math.ceil(imgH * wh_ratio));
|
||||
}
|
||||
NDArray resized_image = NDImageUtils.resize(img, resized_w, imgH);
|
||||
resized_image = resized_image.transpose(2, 0, 1).toType(DataType.FLOAT32,false);
|
||||
resized_image.divi(255f).subi(0.5f).divi(0.5f);
|
||||
NDArray padding_im = ctx.getNDManager().zeros(new Shape(imgC, imgH, imgW), DataType.FLOAT32);
|
||||
padding_im.set(new NDIndex(":,:,0:" + resized_w), resized_image);
|
||||
|
||||
padding_im = padding_im.flip(0);
|
||||
// padding_im = padding_im.expandDims(0);
|
||||
return new NDList(padding_im);
|
||||
}
|
||||
|
||||
@Override
|
||||
public Batchifier getBatchifier() {
|
||||
return Batchifier.fromString(batchifier);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,332 @@
|
||||
package cn.smartjavaai.ocr.model.plate;
|
||||
|
||||
import ai.djl.MalformedModelException;
|
||||
import ai.djl.engine.Engine;
|
||||
import ai.djl.inference.Predictor;
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.modality.cv.ImageFactory;
|
||||
import ai.djl.modality.cv.output.BoundingBox;
|
||||
import ai.djl.modality.cv.output.DetectedObjects;
|
||||
import ai.djl.repository.zoo.Criteria;
|
||||
import ai.djl.repository.zoo.ModelNotFoundException;
|
||||
import ai.djl.repository.zoo.ModelZoo;
|
||||
import ai.djl.repository.zoo.ZooModel;
|
||||
import cn.hutool.core.lang.UUID;
|
||||
import cn.hutool.core.lang.generator.UUIDGenerator;
|
||||
import cn.smartjavaai.common.entity.DetectionRectangle;
|
||||
import cn.smartjavaai.common.entity.R;
|
||||
import cn.smartjavaai.common.pool.PredictorFactory;
|
||||
import cn.smartjavaai.common.utils.Base64ImageUtils;
|
||||
import cn.smartjavaai.common.utils.FileUtils;
|
||||
import cn.smartjavaai.common.utils.ImageUtils;
|
||||
import cn.smartjavaai.common.utils.OpenCVUtils;
|
||||
import cn.smartjavaai.ocr.config.PlateDetModelConfig;
|
||||
import cn.smartjavaai.ocr.config.PlateRecModelConfig;
|
||||
import cn.smartjavaai.ocr.entity.PlateInfo;
|
||||
import cn.smartjavaai.ocr.entity.PlateResult;
|
||||
import cn.smartjavaai.ocr.enums.PlateType;
|
||||
import cn.smartjavaai.ocr.exception.OcrException;
|
||||
import cn.smartjavaai.ocr.model.plate.criteria.PlateDetCriterialFactory;
|
||||
import cn.smartjavaai.ocr.model.plate.criteria.PlateRecCriterialFactory;
|
||||
import cn.smartjavaai.ocr.utils.OcrUtils;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
import org.apache.commons.collections.CollectionUtils;
|
||||
import org.apache.commons.lang3.StringUtils;
|
||||
import org.apache.commons.pool2.ObjectPool;
|
||||
import org.apache.commons.pool2.impl.GenericObjectPool;
|
||||
import org.opencv.core.Core;
|
||||
import org.opencv.core.Mat;
|
||||
import org.opencv.core.Rect;
|
||||
import org.opencv.core.Size;
|
||||
import org.opencv.imgproc.Imgproc;
|
||||
|
||||
import javax.imageio.ImageIO;
|
||||
import java.awt.*;
|
||||
import java.awt.image.BufferedImage;
|
||||
import java.io.*;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.nio.file.Paths;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Iterator;
|
||||
import java.util.List;
|
||||
import java.util.Objects;
|
||||
|
||||
/**
|
||||
* @author dwj
|
||||
*/
|
||||
@Slf4j
|
||||
public class CRNNPlateRecModel implements PlateRecModel{
|
||||
|
||||
|
||||
private GenericObjectPool<Predictor<Image, PlateResult>> recPredictorPool;
|
||||
|
||||
private ZooModel<Image, PlateResult> recModel;
|
||||
|
||||
private PlateRecModelConfig config;
|
||||
|
||||
@Override
|
||||
public void loadModel(PlateRecModelConfig config) {
|
||||
if(StringUtils.isBlank(config.getModelPath())){
|
||||
throw new OcrException("modelPath is null");
|
||||
}
|
||||
this.config = config;
|
||||
//初始化 检测Criteria
|
||||
Criteria<Image, PlateResult> detCriteria = PlateRecCriterialFactory.createCriteria(config);
|
||||
try{
|
||||
recModel = ModelZoo.loadModel(detCriteria);
|
||||
// 创建池子:每个线程独享 Predictor
|
||||
this.recPredictorPool = new GenericObjectPool<>(new PredictorFactory<>(recModel));
|
||||
int predictorPoolSize = config.getPredictorPoolSize();
|
||||
if(config.getPredictorPoolSize() <= 0){
|
||||
predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
|
||||
}
|
||||
recPredictorPool.setMaxTotal(predictorPoolSize);
|
||||
log.debug("当前设备: " + recModel.getNDManager().getDevice());
|
||||
log.debug("当前引擎: " + Engine.getInstance().getEngineName());
|
||||
log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
|
||||
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
|
||||
throw new OcrException("检测模型加载失败", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<List<PlateInfo>> recognize(String imagePath) {
|
||||
if(!FileUtils.isFileExists(imagePath)){
|
||||
return R.fail(R.Status.FILE_NOT_FOUND);
|
||||
}
|
||||
Image img = null;
|
||||
try {
|
||||
img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
|
||||
R<List<PlateInfo>> plateResult = recognize(img);
|
||||
return plateResult;
|
||||
} catch (IOException e) {
|
||||
throw new OcrException("无效的图片", e);
|
||||
} finally {
|
||||
((Mat)img.getWrappedImage()).release();
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<List<PlateInfo>> recognizeBase64(String base64Image) {
|
||||
if(StringUtils.isBlank(base64Image)){
|
||||
return R.fail(R.Status.INVALID_IMAGE);
|
||||
}
|
||||
byte[] imageData = Base64ImageUtils.base64ToImage(base64Image);
|
||||
return recognize(imageData);
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<List<PlateInfo>> recognize(BufferedImage image) {
|
||||
if(!ImageUtils.isImageValid(image)){
|
||||
return R.fail(R.Status.INVALID_IMAGE);
|
||||
}
|
||||
Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image));
|
||||
R<List<PlateInfo>> plateResult = recognize(img);
|
||||
((Mat)img.getWrappedImage()).release();
|
||||
return plateResult;
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<List<PlateInfo>> recognize(byte[] imageData) {
|
||||
if(Objects.isNull(imageData)){
|
||||
return R.fail(R.Status.INVALID_IMAGE);
|
||||
}
|
||||
return recognize(new ByteArrayInputStream(imageData));
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<List<PlateInfo>> recognize(Image image) {
|
||||
if(Objects.isNull(config.getPlateDetModel())){
|
||||
return R.fail(R.Status.PARAM_ERROR.getCode(), "未指定车牌检测模型");
|
||||
}
|
||||
DetectedObjects detectedObjects = config.getPlateDetModel().detect(image);
|
||||
if(Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
|
||||
return R.fail(R.Status.NO_OBJECT_DETECTED);
|
||||
}
|
||||
List<PlateInfo> plateInfoList = OcrUtils.convertToPlateInfo(detectedObjects, image);
|
||||
Predictor<Image, PlateResult> predictor = null;
|
||||
try {
|
||||
predictor = recPredictorPool.borrowObject();
|
||||
for (PlateInfo plateInfo : plateInfoList){
|
||||
DetectionRectangle detectionRectangle = plateInfo.getDetectionRectangle();
|
||||
// Image subImage = image.getSubImage(detectionRectangle.getX(), detectionRectangle.getY(), detectionRectangle.getWidth(), detectionRectangle.getHeight());
|
||||
//透视变换
|
||||
Image subImage = OcrUtils.transformAndCrop((Mat)image.getWrappedImage(), plateInfo.getBox());
|
||||
//双层车牌
|
||||
if(plateInfo.getPlateType() == PlateType.DOUBLE){
|
||||
Mat mergeImage = getSplitMerge((Mat)subImage.getWrappedImage());
|
||||
subImage = ImageFactory.getInstance().fromImage(mergeImage);
|
||||
}
|
||||
PlateResult plateResult = predictor.predict(subImage);
|
||||
if(Objects.nonNull(plateResult)){
|
||||
plateInfo.setPlateNumber(plateResult.getPlateNo());
|
||||
plateInfo.setPlateColor(plateResult.getPlateColor());
|
||||
}
|
||||
}
|
||||
return R.ok(plateInfoList);
|
||||
} catch (Exception e) {
|
||||
throw new OcrException("车牌识别错误", e);
|
||||
}finally {
|
||||
if (predictor != null) {
|
||||
try {
|
||||
recPredictorPool.returnObject(predictor); //归还
|
||||
} catch (Exception e) {
|
||||
log.warn("归还Predictor失败", e);
|
||||
try {
|
||||
predictor.close(); // 归还失败才销毁
|
||||
} catch (Exception ex) {
|
||||
log.error("关闭Predictor失败", ex);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 双层车牌进行分割后识别
|
||||
* @param img
|
||||
* @return
|
||||
*/
|
||||
private Mat getSplitMerge(Mat img) {
|
||||
int h = img.rows();
|
||||
int w = img.cols();
|
||||
|
||||
// 上半部分:高度的前 5/12
|
||||
Rect upperRect = new Rect(0, 0, w, (int)(5.0 / 12 * h));
|
||||
Mat imgUpper = new Mat(img, upperRect);
|
||||
|
||||
// 下半部分:高度从 1/3 开始
|
||||
Rect lowerRect = new Rect(0, (int)(1.0 / 3 * h), w, h - (int)(1.0 / 3 * h));
|
||||
Mat imgLower = new Mat(img, lowerRect);
|
||||
|
||||
// 将上半部分 resize 到与下半部分相同大小
|
||||
Mat resizedUpper = new Mat();
|
||||
Size lowerSize = imgLower.size();
|
||||
Imgproc.resize(imgUpper, resizedUpper, lowerSize);
|
||||
|
||||
// 水平拼接(将上下拼成左右)
|
||||
List<Mat> mergeList = new ArrayList<>();
|
||||
mergeList.add(resizedUpper);
|
||||
mergeList.add(imgLower);
|
||||
|
||||
Mat merged = new Mat();
|
||||
Core.hconcat(mergeList, merged);
|
||||
return merged;
|
||||
}
|
||||
|
||||
@Override
|
||||
public PlateResult recognizeCropped(Image image) {
|
||||
Predictor<Image, PlateResult> predictor = null;
|
||||
try {
|
||||
predictor = recPredictorPool.borrowObject();
|
||||
return predictor.predict(image);
|
||||
} catch (Exception e) {
|
||||
throw new OcrException("车牌检测错误", e);
|
||||
}finally {
|
||||
if (predictor != null) {
|
||||
try {
|
||||
recPredictorPool.returnObject(predictor); //归还
|
||||
} catch (Exception e) {
|
||||
log.warn("归还Predictor失败", e);
|
||||
try {
|
||||
predictor.close(); // 归还失败才销毁
|
||||
} catch (Exception ex) {
|
||||
log.error("关闭Predictor失败", ex);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<List<PlateInfo>> recognize(InputStream inputStream) {
|
||||
if(Objects.isNull(inputStream)){
|
||||
return R.fail(R.Status.INVALID_IMAGE);
|
||||
}
|
||||
Image img = null;
|
||||
try {
|
||||
img = ImageFactory.getInstance().fromInputStream(inputStream);
|
||||
return recognize(img);
|
||||
} catch (IOException e) {
|
||||
throw new OcrException("无效图片输入流", e);
|
||||
} finally {
|
||||
if (img != null){
|
||||
((Mat)img.getWrappedImage()).release();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<Void> recognizeAndDraw(String imagePath, String outputPath) {
|
||||
if(!FileUtils.isFileExists(imagePath)){
|
||||
return R.fail(R.Status.FILE_NOT_FOUND);
|
||||
}
|
||||
Image img = null;
|
||||
try {
|
||||
img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
|
||||
R<List<PlateInfo>> plateResult = recognize(img);
|
||||
if(!plateResult.isSuccess()){
|
||||
return R.fail(plateResult.getCode(), plateResult.getMessage());
|
||||
}
|
||||
if(CollectionUtils.isEmpty(plateResult.getData())){
|
||||
return R.fail(R.Status.NO_OBJECT_DETECTED);
|
||||
}
|
||||
BufferedImage bufferedImage = OpenCVUtils.mat2Image((Mat)img.getWrappedImage());
|
||||
OcrUtils.drawPlateInfo(bufferedImage, plateResult.getData());
|
||||
ImageIO.write(bufferedImage, "png", new File(outputPath));
|
||||
return R.ok();
|
||||
} catch (IOException e) {
|
||||
throw new OcrException(e);
|
||||
} finally {
|
||||
if (img != null){
|
||||
((Mat)img.getWrappedImage()).release();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<BufferedImage> recognizeAndDraw(BufferedImage sourceImage) {
|
||||
if(!ImageUtils.isImageValid(sourceImage)){
|
||||
return R.fail(R.Status.INVALID_IMAGE);
|
||||
}
|
||||
try {
|
||||
R<List<PlateInfo>> plateResult = recognize(sourceImage);
|
||||
if(!plateResult.isSuccess()){
|
||||
return R.fail(plateResult.getCode(), plateResult.getMessage());
|
||||
}
|
||||
if(CollectionUtils.isEmpty(plateResult.getData())){
|
||||
return R.fail(R.Status.NO_OBJECT_DETECTED);
|
||||
}
|
||||
OcrUtils.drawPlateInfo(sourceImage, plateResult.getData());
|
||||
return R.ok(sourceImage);
|
||||
} catch (Exception e) {
|
||||
throw new OcrException("导出图片失败", e);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@Override
|
||||
public GenericObjectPool<Predictor<Image, PlateResult>> getPool() {
|
||||
return recPredictorPool;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() throws Exception {
|
||||
try {
|
||||
if (recPredictorPool != null) {
|
||||
recPredictorPool.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 predictorPool 失败", e);
|
||||
}
|
||||
try {
|
||||
if (recModel != null) {
|
||||
recModel.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 model 失败", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
package cn.smartjavaai.ocr.model.plate;
|
||||
|
||||
import ai.djl.inference.Predictor;
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.modality.cv.output.DetectedObjects;
|
||||
import cn.smartjavaai.common.entity.DetectionResponse;
|
||||
import cn.smartjavaai.ocr.config.OcrDetModelConfig;
|
||||
import cn.smartjavaai.ocr.config.PlateDetModelConfig;
|
||||
import cn.smartjavaai.ocr.entity.OcrBox;
|
||||
import cn.smartjavaai.ocr.entity.PlateInfo;
|
||||
import cn.smartjavaai.common.entity.R;
|
||||
import org.apache.commons.pool2.impl.GenericObjectPool;
|
||||
|
||||
import java.awt.image.BufferedImage;
|
||||
import java.io.InputStream;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* 车牌检测模型
|
||||
* @author dwj
|
||||
*/
|
||||
public interface PlateDetModel extends AutoCloseable{
|
||||
|
||||
/**
|
||||
* 加载模型
|
||||
* @param config
|
||||
*/
|
||||
void loadModel(PlateDetModelConfig config); // 加载模型
|
||||
|
||||
/**
|
||||
* 车牌检测
|
||||
* @param imagePath 图片路径
|
||||
* @return
|
||||
*/
|
||||
default R<List<PlateInfo>> detect(String imagePath) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 车牌检测
|
||||
* @param inputStream
|
||||
* @return
|
||||
*/
|
||||
default R<List<PlateInfo>> detect(InputStream inputStream) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 车牌检测
|
||||
* @param base64Image
|
||||
* @return
|
||||
*/
|
||||
default R<List<PlateInfo>> detectBase64(String base64Image){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 车牌检测
|
||||
* @param image BufferedImage
|
||||
* @return
|
||||
*/
|
||||
default R<List<PlateInfo>> detect(BufferedImage image) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 车牌检测
|
||||
* @param imageData 图片字节数组
|
||||
* @return
|
||||
*/
|
||||
default R<List<PlateInfo>> detect(byte[] imageData) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 车牌检测
|
||||
* @param image DJL Image
|
||||
* @return
|
||||
*/
|
||||
default DetectedObjects detect(Image image){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 检测并绘制结果
|
||||
* @param imagePath 图片输入路径(包含文件名称)
|
||||
* @param outputPath 图片输出路径(包含文件名称)
|
||||
*/
|
||||
default R<Void> detectAndDraw(String imagePath, String outputPath) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 检测并绘制结果
|
||||
* @param sourceImage
|
||||
* @return
|
||||
*/
|
||||
default R<BufferedImage> detectAndDraw(BufferedImage sourceImage){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
default GenericObjectPool<Predictor<Image, DetectedObjects>> getPool(){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
package cn.smartjavaai.ocr.model.plate;
|
||||
|
||||
import ai.djl.inference.Predictor;
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.modality.cv.output.DetectedObjects;
|
||||
import cn.smartjavaai.common.entity.R;
|
||||
import cn.smartjavaai.ocr.config.PlateDetModelConfig;
|
||||
import cn.smartjavaai.ocr.config.PlateRecModelConfig;
|
||||
import cn.smartjavaai.ocr.entity.PlateInfo;
|
||||
import cn.smartjavaai.ocr.entity.PlateResult;
|
||||
import org.apache.commons.pool2.impl.GenericObjectPool;
|
||||
|
||||
import java.awt.image.BufferedImage;
|
||||
import java.io.InputStream;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* 车牌识别模型
|
||||
* @author dwj
|
||||
*/
|
||||
public interface PlateRecModel extends AutoCloseable{
|
||||
|
||||
/**
|
||||
* 加载模型
|
||||
* @param config
|
||||
*/
|
||||
void loadModel(PlateRecModelConfig config); // 加载模型
|
||||
|
||||
/**
|
||||
* 车牌识别
|
||||
* @param imagePath 图片路径
|
||||
* @return
|
||||
*/
|
||||
default R<List<PlateInfo>> recognize(String imagePath) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 车牌识别
|
||||
* @param inputStream
|
||||
* @return
|
||||
*/
|
||||
default R<List<PlateInfo>> recognize(InputStream inputStream) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 车牌识别
|
||||
* @param base64Image
|
||||
* @return
|
||||
*/
|
||||
default R<List<PlateInfo>> recognizeBase64(String base64Image){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 车牌识别
|
||||
* @param image BufferedImage
|
||||
* @return
|
||||
*/
|
||||
default R<List<PlateInfo>> recognize(BufferedImage image) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 车牌识别
|
||||
* @param imageData 图片字节数组
|
||||
* @return
|
||||
*/
|
||||
default R<List<PlateInfo>> recognize(byte[] imageData) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 车牌识别
|
||||
* @param image DJL Image
|
||||
* @return
|
||||
*/
|
||||
default R<List<PlateInfo>> recognize(Image image){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 识别裁剪后的图片
|
||||
* @return
|
||||
*/
|
||||
default PlateResult recognizeCropped(Image image){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 检测并绘制结果
|
||||
* @param imagePath 图片输入路径(包含文件名称)
|
||||
* @param outputPath 图片输出路径(包含文件名称)
|
||||
*/
|
||||
default R<Void> recognizeAndDraw(String imagePath, String outputPath) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 检测并绘制结果
|
||||
* @param sourceImage
|
||||
* @return
|
||||
*/
|
||||
default R<BufferedImage> recognizeAndDraw(BufferedImage sourceImage){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
default GenericObjectPool<Predictor<Image, PlateResult>> getPool() {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,244 @@
|
||||
package cn.smartjavaai.ocr.model.plate;
|
||||
|
||||
import ai.djl.MalformedModelException;
|
||||
import ai.djl.engine.Engine;
|
||||
import ai.djl.inference.Predictor;
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.modality.cv.ImageFactory;
|
||||
import ai.djl.modality.cv.output.DetectedObjects;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.ndarray.NDManager;
|
||||
import ai.djl.repository.zoo.Criteria;
|
||||
import ai.djl.repository.zoo.ModelNotFoundException;
|
||||
import ai.djl.repository.zoo.ModelZoo;
|
||||
import ai.djl.repository.zoo.ZooModel;
|
||||
import cn.smartjavaai.common.entity.R;
|
||||
import cn.smartjavaai.common.pool.PredictorFactory;
|
||||
import cn.smartjavaai.common.utils.Base64ImageUtils;
|
||||
import cn.smartjavaai.common.utils.FileUtils;
|
||||
import cn.smartjavaai.common.utils.ImageUtils;
|
||||
import cn.smartjavaai.common.utils.OpenCVUtils;
|
||||
import cn.smartjavaai.ocr.config.OcrDetModelConfig;
|
||||
import cn.smartjavaai.ocr.config.PlateDetModelConfig;
|
||||
import cn.smartjavaai.ocr.entity.OcrBox;
|
||||
import cn.smartjavaai.ocr.entity.PlateInfo;
|
||||
import cn.smartjavaai.ocr.exception.OcrException;
|
||||
import cn.smartjavaai.ocr.model.common.detect.criteria.OcrCommonDetCriterialFactory;
|
||||
import cn.smartjavaai.ocr.model.plate.criteria.PlateDetCriterialFactory;
|
||||
import cn.smartjavaai.ocr.utils.OcrUtils;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
import org.apache.commons.lang3.StringUtils;
|
||||
import org.apache.commons.pool2.ObjectPool;
|
||||
import org.apache.commons.pool2.impl.GenericObjectPool;
|
||||
import org.opencv.core.Mat;
|
||||
|
||||
import javax.imageio.ImageIO;
|
||||
import java.awt.image.BufferedImage;
|
||||
import java.io.ByteArrayInputStream;
|
||||
import java.io.ByteArrayOutputStream;
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.nio.file.Paths;
|
||||
import java.util.List;
|
||||
import java.util.Objects;
|
||||
|
||||
/**
|
||||
* Yolov5 车牌检测模型
|
||||
* @author dwj
|
||||
* @date 2025/7/23
|
||||
*/
|
||||
@Slf4j
|
||||
public class Yolov5PlateDetModel implements PlateDetModel{
|
||||
|
||||
private GenericObjectPool<Predictor<Image, DetectedObjects>> detPredictorPool;
|
||||
|
||||
private ZooModel<Image, DetectedObjects> detectionModel;
|
||||
|
||||
private PlateDetModelConfig config;
|
||||
|
||||
@Override
|
||||
public void loadModel(PlateDetModelConfig config) {
|
||||
if(StringUtils.isBlank(config.getModelPath())){
|
||||
throw new OcrException("modelPath is null");
|
||||
}
|
||||
this.config = config;
|
||||
//初始化 检测Criteria
|
||||
Criteria<Image, DetectedObjects> detCriteria = PlateDetCriterialFactory.createCriteria(config);
|
||||
try{
|
||||
detectionModel = ModelZoo.loadModel(detCriteria);
|
||||
// 创建池子:每个线程独享 Predictor
|
||||
this.detPredictorPool = new GenericObjectPool<>(new PredictorFactory<>(detectionModel));
|
||||
int predictorPoolSize = config.getPredictorPoolSize();
|
||||
if(config.getPredictorPoolSize() <= 0){
|
||||
predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
|
||||
}
|
||||
detPredictorPool.setMaxTotal(predictorPoolSize);
|
||||
log.debug("当前设备: " + detectionModel.getNDManager().getDevice());
|
||||
log.debug("当前引擎: " + Engine.getInstance().getEngineName());
|
||||
log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
|
||||
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
|
||||
throw new OcrException("检测模型加载失败", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<List<PlateInfo>> detect(String imagePath) {
|
||||
if(!FileUtils.isFileExists(imagePath)){
|
||||
return R.fail(R.Status.FILE_NOT_FOUND);
|
||||
}
|
||||
Image img = null;
|
||||
try {
|
||||
img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
|
||||
} catch (IOException e) {
|
||||
throw new OcrException("无效的图片", e);
|
||||
}
|
||||
DetectedObjects detectedObjects = detect(img);
|
||||
if (Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
|
||||
return R.fail(R.Status.NO_OBJECT_DETECTED);
|
||||
}
|
||||
List<PlateInfo> plateInfoList = OcrUtils.convertToPlateInfo(detectedObjects, img);
|
||||
((Mat)img.getWrappedImage()).release();
|
||||
return R.ok(plateInfoList);
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<List<PlateInfo>> detectBase64(String base64Image) {
|
||||
if(StringUtils.isBlank(base64Image)){
|
||||
return R.fail(R.Status.INVALID_IMAGE);
|
||||
}
|
||||
byte[] imageData = Base64ImageUtils.base64ToImage(base64Image);
|
||||
return detect(imageData);
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<List<PlateInfo>> detect(BufferedImage image) {
|
||||
if(!ImageUtils.isImageValid(image)){
|
||||
return R.fail(R.Status.INVALID_IMAGE);
|
||||
}
|
||||
Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image));
|
||||
DetectedObjects detectedObjects = detect(img);
|
||||
if (Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
|
||||
return R.fail(R.Status.NO_OBJECT_DETECTED);
|
||||
}
|
||||
List<PlateInfo> plateInfoList = OcrUtils.convertToPlateInfo(detectedObjects, img);
|
||||
((Mat)img.getWrappedImage()).release();
|
||||
return R.ok(plateInfoList);
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<List<PlateInfo>> detect(byte[] imageData) {
|
||||
if(Objects.isNull(imageData)){
|
||||
return R.fail(R.Status.INVALID_IMAGE);
|
||||
}
|
||||
return detect(new ByteArrayInputStream(imageData));
|
||||
}
|
||||
|
||||
@Override
|
||||
public DetectedObjects detect(Image image) {
|
||||
Predictor<Image, DetectedObjects> predictor = null;
|
||||
try {
|
||||
predictor = detPredictorPool.borrowObject();
|
||||
return predictor.predict(image);
|
||||
} catch (Exception e) {
|
||||
throw new OcrException("车牌检测错误", e);
|
||||
}finally {
|
||||
if (predictor != null) {
|
||||
try {
|
||||
detPredictorPool.returnObject(predictor); //归还
|
||||
} catch (Exception e) {
|
||||
log.warn("归还Predictor失败", e);
|
||||
try {
|
||||
predictor.close(); // 归还失败才销毁
|
||||
} catch (Exception ex) {
|
||||
log.error("关闭Predictor失败", ex);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<List<PlateInfo>> detect(InputStream inputStream) {
|
||||
if(Objects.isNull(inputStream)){
|
||||
return R.fail(R.Status.INVALID_IMAGE);
|
||||
}
|
||||
try {
|
||||
Image img = ImageFactory.getInstance().fromInputStream(inputStream);
|
||||
DetectedObjects detection = detect(img);
|
||||
List<PlateInfo> plateInfoList = OcrUtils.convertToPlateInfo(detection, img);
|
||||
((Mat)img.getWrappedImage()).release();
|
||||
return R.ok(plateInfoList);
|
||||
} catch (IOException e) {
|
||||
throw new OcrException("无效图片输入流", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<Void> detectAndDraw(String imagePath, String outputPath) {
|
||||
if(!FileUtils.isFileExists(imagePath)){
|
||||
return R.fail(R.Status.FILE_NOT_FOUND);
|
||||
}
|
||||
try {
|
||||
Image img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
|
||||
DetectedObjects detectedObjects = detect(img);
|
||||
if(Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
|
||||
return R.fail(R.Status.NO_FACE_DETECTED);
|
||||
}
|
||||
img.drawBoundingBoxes(detectedObjects);
|
||||
Path output = Paths.get(outputPath);
|
||||
log.debug("Saving to {}", output.toAbsolutePath().toString());
|
||||
img.save(Files.newOutputStream(output), "png");
|
||||
return R.ok();
|
||||
} catch (IOException e) {
|
||||
throw new OcrException(e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<BufferedImage> detectAndDraw(BufferedImage sourceImage) {
|
||||
if(!ImageUtils.isImageValid(sourceImage)){
|
||||
return R.fail(R.Status.INVALID_IMAGE);
|
||||
}
|
||||
Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(sourceImage));
|
||||
DetectedObjects detectedObjects = detect(img);
|
||||
if(Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
|
||||
return R.fail(R.Status.NO_FACE_DETECTED);
|
||||
}
|
||||
img.drawBoundingBoxes(detectedObjects);
|
||||
try {
|
||||
ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
|
||||
// 调用 save 方法将 Image 写入字节流
|
||||
img.save(outputStream, "png");
|
||||
// 将字节流转换为 BufferedImage
|
||||
byte[] imageBytes = outputStream.toByteArray();
|
||||
return R.ok(ImageIO.read(new ByteArrayInputStream(imageBytes)));
|
||||
} catch (IOException e) {
|
||||
throw new OcrException("导出图片失败", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public GenericObjectPool<Predictor<Image, DetectedObjects>> getPool() {
|
||||
return detPredictorPool;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() throws Exception {
|
||||
try {
|
||||
if (detPredictorPool != null) {
|
||||
detPredictorPool.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 predictorPool 失败", e);
|
||||
}
|
||||
try {
|
||||
if (detectionModel != null) {
|
||||
detectionModel.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 model 失败", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
package cn.smartjavaai.ocr.model.plate.criteria;
|
||||
|
||||
import ai.djl.Device;
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.modality.cv.output.DetectedObjects;
|
||||
import ai.djl.repository.zoo.Criteria;
|
||||
import ai.djl.training.util.ProgressBar;
|
||||
import cn.smartjavaai.common.enums.DeviceEnum;
|
||||
import cn.smartjavaai.ocr.config.PlateDetModelConfig;
|
||||
import cn.smartjavaai.ocr.enums.PlateDetModelEnum;
|
||||
import cn.smartjavaai.ocr.model.plate.translator.Yolo5PlateDetectTranslator;
|
||||
import cn.smartjavaai.ocr.model.plate.translator.Yolov7PlateDetectTranslator;
|
||||
import cn.smartjavaai.ocr.model.plate.translator.Yolov8PlateDetectTranslator;
|
||||
import org.apache.commons.lang3.StringUtils;
|
||||
|
||||
import java.nio.file.Paths;
|
||||
import java.util.Objects;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
|
||||
/**
|
||||
* @author dwj
|
||||
* @date 2025/7/8
|
||||
*/
|
||||
public class PlateDetCriterialFactory {
|
||||
|
||||
|
||||
public static Criteria<Image, DetectedObjects> createCriteria(PlateDetModelConfig config) {
|
||||
Device device = null;
|
||||
if(!Objects.isNull(config.getDevice())){
|
||||
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
|
||||
}
|
||||
Criteria<Image, DetectedObjects> criteria = null;
|
||||
ConcurrentHashMap params = new ConcurrentHashMap<String, String>();
|
||||
params.putAll(config.getCustomParams());
|
||||
if(StringUtils.isNotBlank(config.getBatchifier())){
|
||||
params.put("batchifier", config.getBatchifier());
|
||||
}
|
||||
if(config.getModelEnum() == PlateDetModelEnum.YOLOV5){
|
||||
criteria =
|
||||
Criteria.builder()
|
||||
.optEngine("OnnxRuntime")
|
||||
.setTypes(Image.class, DetectedObjects.class)
|
||||
.optModelPath(Paths.get(config.getModelPath()))
|
||||
.optTranslator(new Yolo5PlateDetectTranslator(params))
|
||||
.optDevice(device)
|
||||
.optProgress(new ProgressBar())
|
||||
.build();
|
||||
}else if (config.getModelEnum() == PlateDetModelEnum.YOLOV7){
|
||||
criteria =
|
||||
Criteria.builder()
|
||||
.optEngine("OnnxRuntime")
|
||||
.setTypes(Image.class, DetectedObjects.class)
|
||||
.optModelPath(Paths.get(config.getModelPath()))
|
||||
.optTranslator(new Yolov7PlateDetectTranslator(params))
|
||||
.optDevice(device)
|
||||
.optProgress(new ProgressBar())
|
||||
.build();
|
||||
}
|
||||
// else if (config.getModelEnum() == PlateDetModelEnum.YOLOV8){
|
||||
// criteria =
|
||||
// Criteria.builder()
|
||||
// .optEngine("OnnxRuntime")
|
||||
// .setTypes(Image.class, DetectedObjects.class)
|
||||
// .optModelPath(Paths.get(config.getModelPath()))
|
||||
// .optTranslator(new Yolov8PlateDetectTranslator(params))
|
||||
// .optDevice(device)
|
||||
// .optProgress(new ProgressBar())
|
||||
// .build();
|
||||
// }
|
||||
return criteria;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package cn.smartjavaai.ocr.model.plate.criteria;
|
||||
|
||||
import ai.djl.Device;
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.repository.zoo.Criteria;
|
||||
import ai.djl.training.util.ProgressBar;
|
||||
import cn.smartjavaai.common.enums.DeviceEnum;
|
||||
import cn.smartjavaai.ocr.config.PlateRecModelConfig;
|
||||
import cn.smartjavaai.ocr.entity.PlateResult;
|
||||
import cn.smartjavaai.ocr.enums.PlateRecModelEnum;
|
||||
import cn.smartjavaai.ocr.model.plate.translator.CRNNPlateRecTranslator;
|
||||
|
||||
import java.nio.file.Paths;
|
||||
import java.util.Objects;
|
||||
|
||||
/**
|
||||
* @author dwj
|
||||
* @date 2025/7/8
|
||||
*/
|
||||
public class PlateRecCriterialFactory {
|
||||
|
||||
|
||||
public static Criteria<Image, PlateResult> createCriteria(PlateRecModelConfig config) {
|
||||
Device device = null;
|
||||
if(!Objects.isNull(config.getDevice())){
|
||||
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
|
||||
}
|
||||
Criteria<Image, PlateResult> criteria = null;
|
||||
if(config.getModelEnum() == PlateRecModelEnum.PLATE_REC_CRNN){
|
||||
criteria =
|
||||
Criteria.builder()
|
||||
.optEngine("OnnxRuntime")
|
||||
.setTypes(Image.class, PlateResult.class)
|
||||
.optModelPath(Paths.get(config.getModelPath()))
|
||||
.optTranslator(new CRNNPlateRecTranslator())
|
||||
.optDevice(device)
|
||||
.optProgress(new ProgressBar())
|
||||
.build();
|
||||
}
|
||||
return criteria;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
package cn.smartjavaai.ocr.model.plate.translator;
|
||||
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.modality.cv.util.NDImageUtils;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.ndarray.NDManager;
|
||||
import ai.djl.ndarray.types.DataType;
|
||||
import ai.djl.translate.Batchifier;
|
||||
import ai.djl.translate.Translator;
|
||||
import ai.djl.translate.TranslatorContext;
|
||||
import cn.smartjavaai.ocr.entity.PlateResult;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* @author dwj
|
||||
*/
|
||||
public class CRNNPlateRecTranslator implements Translator<Image, PlateResult> {
|
||||
|
||||
private static final String plateName = "#京沪津渝冀晋蒙辽吉黑苏浙皖闽赣鲁豫鄂湘粤桂琼川贵云藏陕甘青宁新学警港澳挂使领民航危0123456789ABCDEFGHJKLMNPQRSTUVWXYZ险品";
|
||||
private static final String[] plateColors = {"黑色", "蓝色", "绿色", "白色", "黄色"};
|
||||
private static final float MEAN = 0.588f;
|
||||
private static final float STD = 0.193f;
|
||||
|
||||
@Override
|
||||
public NDList processInput(TranslatorContext ctx, Image input) {
|
||||
NDManager manager = ctx.getNDManager();
|
||||
|
||||
// Resize to (168, 48)
|
||||
NDArray array = input.toNDArray(manager, Image.Flag.COLOR);
|
||||
array = NDImageUtils.resize(array, 168, 48);
|
||||
|
||||
// Normalize
|
||||
array = array.toType(DataType.FLOAT32, false)
|
||||
.div(255f)
|
||||
.sub(MEAN)
|
||||
.div(STD);
|
||||
|
||||
// HWC to CHW
|
||||
array = array.transpose(2, 0, 1);
|
||||
array = array.expandDims(0); // batch dimension
|
||||
|
||||
return new NDList(array);
|
||||
}
|
||||
|
||||
@Override
|
||||
public PlateResult processOutput(TranslatorContext ctx, NDList list) {
|
||||
NDArray plateOutput = list.get(0); // shape: [1, T, num_classes]
|
||||
NDArray colorOutput = list.get(1); // shape: [1, num_colors]
|
||||
|
||||
int[] plateIdx = plateOutput.argMax(-1)
|
||||
.toType(DataType.INT32, false)
|
||||
.toIntArray();
|
||||
int colorIdx = colorOutput.argMax(1).toType(DataType.INT32, false).toIntArray()[0];
|
||||
|
||||
String plateNo = decodePlate(plateIdx);
|
||||
String plateColor = plateColors[colorIdx];
|
||||
|
||||
return new PlateResult(plateNo, plateColor);
|
||||
}
|
||||
|
||||
private String decodePlate(int[] preds) {
|
||||
int pre = 0;
|
||||
List<Integer> newPreds = new ArrayList<>();
|
||||
for (int idx : preds) {
|
||||
if (idx != 0 && idx != pre) {
|
||||
newPreds.add(idx);
|
||||
}
|
||||
pre = idx;
|
||||
}
|
||||
|
||||
StringBuilder sb = new StringBuilder();
|
||||
for (int i : newPreds) {
|
||||
if (i >= 0 && i < plateName.length()) {
|
||||
sb.append(plateName.charAt(i));
|
||||
}
|
||||
}
|
||||
return sb.toString();
|
||||
}
|
||||
|
||||
@Override
|
||||
public Batchifier getBatchifier() {
|
||||
return null; // 非批量任务
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,194 @@
|
||||
package cn.smartjavaai.ocr.model.plate.translator;
|
||||
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.modality.cv.output.BoundingBox;
|
||||
import ai.djl.modality.cv.output.DetectedObjects;
|
||||
import ai.djl.modality.cv.output.Landmark;
|
||||
import ai.djl.modality.cv.output.Point;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDArrays;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.ndarray.NDManager;
|
||||
import ai.djl.ndarray.types.DataType;
|
||||
import ai.djl.translate.Batchifier;
|
||||
import ai.djl.translate.Translator;
|
||||
import ai.djl.translate.TranslatorContext;
|
||||
import cn.smartjavaai.common.utils.LetterBoxUtils;
|
||||
import cn.smartjavaai.common.utils.NMSUtils;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
/**
|
||||
* @author dwj
|
||||
*/
|
||||
public class Yolo5PlateDetectTranslator implements Translator<Image, DetectedObjects> {
|
||||
|
||||
private int inputSize = 640;
|
||||
private float minConfThreshold = 0.3f;
|
||||
private float iouThreshold = 0.5f;
|
||||
|
||||
private float confThreshold = 0;
|
||||
|
||||
private int imageWidth;
|
||||
private int imageHeight;
|
||||
|
||||
private int topK;
|
||||
|
||||
private LetterBoxUtils.ResizeResult letterBoxResult;
|
||||
|
||||
public Yolo5PlateDetectTranslator(Map<String, ?> arguments) {
|
||||
confThreshold =
|
||||
arguments.containsKey("confThreshold")
|
||||
? Integer.parseInt(arguments.get("confThreshold").toString())
|
||||
: 0.3f;
|
||||
|
||||
iouThreshold =
|
||||
arguments.containsKey("iouThreshold")
|
||||
? Integer.parseInt(arguments.get("iouThreshold").toString())
|
||||
: 0.5f;
|
||||
|
||||
topK = arguments.containsKey("topk")
|
||||
? Integer.parseInt(arguments.get("topk").toString())
|
||||
: 100;
|
||||
}
|
||||
|
||||
@Override
|
||||
public NDList processInput(TranslatorContext ctx, Image input) {
|
||||
NDManager manager = ctx.getNDManager();
|
||||
NDArray array = input.toNDArray(manager, Image.Flag.COLOR);
|
||||
imageWidth = (int) array.getShape().get(1);
|
||||
imageHeight = (int) array.getShape().get(0);
|
||||
//Letter box resize 640x640 with padding (保持比例,补边缘)
|
||||
letterBoxResult = LetterBoxUtils.letterbox(manager, array, inputSize, inputSize, 114f, LetterBoxUtils.PaddingPosition.CENTER);
|
||||
array = letterBoxResult.image;
|
||||
// 转为 float32 且归一化到 0~1
|
||||
array = array.toType(DataType.FLOAT32, false).div(255f); // HWC
|
||||
// HWC -> CHW
|
||||
array = array.transpose(2, 0, 1); // CHW
|
||||
return new NDList(array.expandDims(0));
|
||||
}
|
||||
|
||||
@Override
|
||||
public DetectedObjects processOutput(TranslatorContext ctx, NDList list) {
|
||||
NDManager manager = ctx.getNDManager();
|
||||
//[x_center, y_center, w, h, obj_conf, 8个关键点, class1_conf, class2_conf]
|
||||
//目标置信度 obj_conf 5:13 关键点 [13:15]分类得分:单层车牌 / 双层车牌
|
||||
NDArray dets = list.singletonOrThrow();
|
||||
//置信度过滤 (1,25200, 15)
|
||||
NDArray dets0 = dets.get(0);
|
||||
NDArray conf = dets0.get(":, 4"); // shape [N]
|
||||
NDArray mask = conf.gt(minConfThreshold);
|
||||
//筛选出符合条件的框(17,15)
|
||||
NDArray detsFiltered = dets0.get(mask); // 筛掉低置信度
|
||||
|
||||
//把分类得分 [13:15] * 置信度 [4:5] 做联合概率
|
||||
NDArray clsLogits = detsFiltered.get(":, 13:15"); // (N, 2)
|
||||
NDArray confFiltered = detsFiltered.get(":, 4").reshape(-1, 1); // (N, 1)
|
||||
clsLogits = clsLogits.mul(confFiltered); // (N, 2),变成 obj_conf * class_conf
|
||||
|
||||
NDArray jointScore = clsLogits.max(new int[]{1}); // shape (N,)
|
||||
// 联合过滤
|
||||
NDArray jointMask = jointScore.gt(confThreshold);
|
||||
detsFiltered = detsFiltered.get(jointMask);
|
||||
|
||||
if (detsFiltered.isEmpty()) {
|
||||
return new DetectedObjects(Collections.emptyList(), Collections.emptyList(), Collections.emptyList());
|
||||
}
|
||||
|
||||
clsLogits = clsLogits.get(jointMask);
|
||||
|
||||
|
||||
//中心点框 [x,y,w,h] ➔ 左上右下 [x1,y1,x2,y2]
|
||||
NDArray xywh = detsFiltered.get(":, 0:4"); // (N, 4)
|
||||
NDArray halfWH = xywh.get(":, 2:4").div(2); // (N, 2)
|
||||
NDArray xy1 = xywh.get(":, 0:2").sub(halfWH); // (N, 2)
|
||||
NDArray xy2 = xywh.get(":, 0:2").add(halfWH); // (N, 2)
|
||||
NDArray boxes = NDArrays.concat(new NDList(xy1, xy2), 1); // (N, 4)
|
||||
|
||||
// 分类得分最大值:score (N, 1),对应类别 index (N, 1)
|
||||
NDArray scores = clsLogits.max(new int[]{1}, true); // (N, 1)
|
||||
NDArray indices = clsLogits.argMax(1).reshape(-1, 1).toType(DataType.FLOAT32, false); // (N, 1)
|
||||
|
||||
// 关键点坐标 [5:13]
|
||||
NDArray keyPoints = detsFiltered.get(":, 5:13"); // (N, 8)
|
||||
|
||||
// 拼成最终结果:(x1, y1, x2, y2, score, 8关键点, index)
|
||||
NDArray output = NDArrays.concat(new NDList(boxes, scores, keyPoints, indices), 1); // (N, 14)
|
||||
|
||||
// NMS 过滤掉重叠框
|
||||
int[] keepIndices = NMSUtils.nms(boxes, scores.squeeze(), iouThreshold); // scores.squeeze() ➝ (N,)】
|
||||
if (keepIndices.length == 0) {
|
||||
return new DetectedObjects(Collections.emptyList(), Collections.emptyList(), Collections.emptyList());
|
||||
}
|
||||
NDArray kept = output.get(manager.create(keepIndices));
|
||||
// 如果超过 topK,则截断
|
||||
if (keepIndices.length > topK) {
|
||||
int[] topkIndices = new int[topK];
|
||||
System.arraycopy(keepIndices, 0, topkIndices, 0, topK);
|
||||
keepIndices = topkIndices;
|
||||
}
|
||||
//恢复原图坐标(除回比例,减掉 padding)
|
||||
NDArray restored = LetterBoxUtils.restoreBox(kept, letterBoxResult.r, letterBoxResult.left, letterBoxResult.top, 5,8);
|
||||
|
||||
List<String> classNames = new ArrayList<>();
|
||||
List<Double> probabilities = new ArrayList<>();
|
||||
List<BoundingBox> boundingBoxes = new ArrayList<>();
|
||||
|
||||
float[] flatData = restored.toFloatArray();
|
||||
long[] shape = restored.getShape().getShape(); // 比如 (N, 14)
|
||||
int rows = (int) shape[0];
|
||||
int cols = (int) shape[1];
|
||||
|
||||
// 把一维数组重组为二维数组
|
||||
float[][] data = new float[rows][cols];
|
||||
for (int i = 0; i < rows; i++) {
|
||||
System.arraycopy(flatData, i * cols, data[i], 0, cols);
|
||||
}
|
||||
|
||||
for (float[] row : data) {
|
||||
// row结构:(x1, y1, x2, y2, score, kp1,..., kp8, classIndex)
|
||||
float x1 = row[0];
|
||||
float y1 = row[1];
|
||||
float x2 = row[2];
|
||||
float y2 = row[3];
|
||||
float score = row[4];
|
||||
int classIndex = (int) row[13];
|
||||
|
||||
double prob = score;
|
||||
String className = classIndex == 0 ? "single" : "double";
|
||||
|
||||
// 转相对坐标,DJL的Rectangle用比例坐标(0~1)
|
||||
double rectX = x1 / imageWidth;
|
||||
double rectY = y1 / imageHeight;
|
||||
double rectW = (x2 - x1) / imageWidth;
|
||||
double rectH = (y2 - y1) / imageHeight;
|
||||
|
||||
// 构建 Polygon 四个角点
|
||||
List<Point> pointsSrc = new ArrayList<>();
|
||||
pointsSrc.add(new Point(row[5], row[6]));
|
||||
pointsSrc.add(new Point(row[7], row[8]));
|
||||
pointsSrc.add(new Point(row[9], row[10]));
|
||||
pointsSrc.add(new Point(row[11], row[12]));
|
||||
|
||||
Landmark box = new Landmark(rectX, rectY, rectW, rectH, pointsSrc);
|
||||
classNames.add(className);
|
||||
probabilities.add(prob);
|
||||
boundingBoxes.add(box);
|
||||
}
|
||||
DetectedObjects detectedObjects = new DetectedObjects(classNames, probabilities, boundingBoxes);
|
||||
return detectedObjects;
|
||||
|
||||
}
|
||||
|
||||
@Override
|
||||
public Batchifier getBatchifier() {
|
||||
return null;
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
package cn.smartjavaai.ocr.model.plate.translator;
|
||||
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.modality.cv.output.*;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDArrays;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.ndarray.NDManager;
|
||||
import ai.djl.ndarray.types.DataType;
|
||||
import ai.djl.translate.Batchifier;
|
||||
import ai.djl.translate.Translator;
|
||||
import ai.djl.translate.TranslatorContext;
|
||||
import cn.smartjavaai.common.utils.LetterBoxUtils;
|
||||
import cn.smartjavaai.common.utils.NMSUtils;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
/**
|
||||
* @author dwj
|
||||
*/
|
||||
public class Yolov7PlateDetectTranslator implements Translator<Image, DetectedObjects> {
|
||||
|
||||
private int inputSize = 640;
|
||||
private float minConfThreshold = 0.3f;
|
||||
private float iouThreshold = 0.5f;
|
||||
|
||||
private float confThreshold = 0;
|
||||
|
||||
private int imageWidth;
|
||||
private int imageHeight;
|
||||
|
||||
private int topK;
|
||||
|
||||
private LetterBoxUtils.ResizeResult letterBoxResult;
|
||||
|
||||
public Yolov7PlateDetectTranslator(Map<String, ?> arguments) {
|
||||
confThreshold =
|
||||
arguments.containsKey("confThreshold")
|
||||
? Integer.parseInt(arguments.get("confThreshold").toString())
|
||||
: 0.3f;
|
||||
|
||||
iouThreshold =
|
||||
arguments.containsKey("iouThreshold")
|
||||
? Integer.parseInt(arguments.get("iouThreshold").toString())
|
||||
: 0.5f;
|
||||
|
||||
topK = arguments.containsKey("topk")
|
||||
? Integer.parseInt(arguments.get("topk").toString())
|
||||
: 100;
|
||||
}
|
||||
|
||||
@Override
|
||||
public NDList processInput(TranslatorContext ctx, Image input) {
|
||||
NDManager manager = ctx.getNDManager();
|
||||
NDArray array = input.toNDArray(manager, Image.Flag.COLOR);
|
||||
imageWidth = (int) array.getShape().get(1);
|
||||
imageHeight = (int) array.getShape().get(0);
|
||||
//Letter box resize 640x640 with padding (保持比例,补边缘)
|
||||
letterBoxResult = LetterBoxUtils.letterbox(manager, array, inputSize, inputSize, 114f, LetterBoxUtils.PaddingPosition.CENTER);
|
||||
array = letterBoxResult.image;
|
||||
// 转为 float32 且归一化到 0~1
|
||||
array = array.toType(DataType.FLOAT32, false).div(255f); // HWC
|
||||
// HWC -> CHW
|
||||
array = array.transpose(2, 0, 1); // CHW
|
||||
return new NDList(array.expandDims(0));
|
||||
}
|
||||
|
||||
@Override
|
||||
public DetectedObjects processOutput(TranslatorContext ctx, NDList list) {
|
||||
NDManager manager = ctx.getNDManager();
|
||||
int num_cls = 2;
|
||||
//[x_center, y_center, w, h, obj_conf, class1_conf, class2_conf,8个关键点]
|
||||
//目标置信度 obj_conf 5:13 关键点 [13:15]分类得分:单层车牌 / 双层车牌
|
||||
NDArray dets = list.singletonOrThrow();
|
||||
//置信度过滤 (1,25200, 15)
|
||||
NDArray dets0 = dets.get(0);
|
||||
NDArray conf = dets0.get(":, 4"); // shape [N]
|
||||
NDArray mask = conf.gt(minConfThreshold);
|
||||
//筛选出符合条件的框(17,15)
|
||||
NDArray detsFiltered = dets0.get(mask); // 筛掉低置信度
|
||||
|
||||
//把分类得分 [5:7] * 置信度 [4:5] 做联合概率
|
||||
NDArray clsLogits = detsFiltered.get(":, 5:7"); // (N, 2)
|
||||
NDArray confFiltered = detsFiltered.get(":, 4").reshape(-1, 1); // (N, 1)
|
||||
clsLogits = clsLogits.mul(confFiltered); // (N, 2),变成 obj_conf * class_conf
|
||||
|
||||
NDArray jointScore = clsLogits.max(new int[]{1}); // shape (N,)
|
||||
// 联合过滤
|
||||
NDArray jointMask = jointScore.gt(confThreshold);
|
||||
detsFiltered = detsFiltered.get(jointMask);
|
||||
clsLogits = clsLogits.get(jointMask);
|
||||
|
||||
|
||||
//中心点框 [x,y,w,h] ➔ 左上右下 [x1,y1,x2,y2]
|
||||
NDArray xywh = detsFiltered.get(":, 0:4"); // (N, 4)
|
||||
NDArray halfWH = xywh.get(":, 2:4").div(2); // (N, 2)
|
||||
NDArray xy1 = xywh.get(":, 0:2").sub(halfWH); // (N, 2)
|
||||
NDArray xy2 = xywh.get(":, 0:2").add(halfWH); // (N, 2)
|
||||
NDArray boxes = NDArrays.concat(new NDList(xy1, xy2), 1); // (N, 4)
|
||||
|
||||
// 分类得分最大值:score (N, 1),对应类别 index (N, 1)
|
||||
NDArray scores = clsLogits.max(new int[]{1}, true); // (N, 1)
|
||||
NDArray indices = clsLogits.argMax(1).reshape(-1, 1).toType(DataType.FLOAT32, false); // (N, 1)
|
||||
|
||||
// 关键点坐标 [7,8,10,11,13,14,16,17]
|
||||
NDArray keyPoints = NDArrays.concat(new NDList(
|
||||
detsFiltered.get(":, 7:8"),
|
||||
detsFiltered.get(":, 8:9"),
|
||||
detsFiltered.get(":, 10:11"),
|
||||
detsFiltered.get(":, 11:12"),
|
||||
detsFiltered.get(":, 13:14"),
|
||||
detsFiltered.get(":, 14:15"),
|
||||
detsFiltered.get(":, 16:17"),
|
||||
detsFiltered.get(":, 17:18")
|
||||
), 1); // 拼成 (N, 8)
|
||||
|
||||
// 拼成最终结果:(x1, y1, x2, y2, score, 8关键点, index)
|
||||
NDArray output = NDArrays.concat(new NDList(boxes, scores, keyPoints, indices), 1); // (N, 14)
|
||||
|
||||
// NMS 过滤掉重叠框
|
||||
int[] keepIndices = NMSUtils.nms(boxes, scores.squeeze(), iouThreshold); // scores.squeeze() ➝ (N,)
|
||||
NDArray kept = output.get(manager.create(keepIndices));
|
||||
// 如果超过 topK,则截断
|
||||
if (keepIndices.length > topK) {
|
||||
int[] topkIndices = new int[topK];
|
||||
System.arraycopy(keepIndices, 0, topkIndices, 0, topK);
|
||||
keepIndices = topkIndices;
|
||||
}
|
||||
//恢复原图坐标(除回比例,减掉 padding)
|
||||
NDArray restored = LetterBoxUtils.restoreBox(kept, letterBoxResult.r, letterBoxResult.left, letterBoxResult.top, 5,8);
|
||||
|
||||
List<String> classNames = new ArrayList<>();
|
||||
List<Double> probabilities = new ArrayList<>();
|
||||
List<BoundingBox> boundingBoxes = new ArrayList<>();
|
||||
|
||||
float[] flatData = restored.toFloatArray();
|
||||
long[] shape = restored.getShape().getShape(); // 比如 (N, 14)
|
||||
int rows = (int) shape[0];
|
||||
int cols = (int) shape[1];
|
||||
|
||||
// 把一维数组重组为二维数组
|
||||
float[][] data = new float[rows][cols];
|
||||
for (int i = 0; i < rows; i++) {
|
||||
System.arraycopy(flatData, i * cols, data[i], 0, cols);
|
||||
}
|
||||
|
||||
for (float[] row : data) {
|
||||
// row结构:(x1, y1, x2, y2, score, kp1,..., kp8, classIndex)
|
||||
float x1 = row[0];
|
||||
float y1 = row[1];
|
||||
float x2 = row[2];
|
||||
float y2 = row[3];
|
||||
float score = row[4];
|
||||
int classIndex = (int) row[13];
|
||||
|
||||
double prob = score;
|
||||
String className = classIndex == 0 ? "single" : "double";
|
||||
|
||||
// 转相对坐标,DJL的Rectangle用比例坐标(0~1)
|
||||
double rectX = x1 / imageWidth;
|
||||
double rectY = y1 / imageHeight;
|
||||
double rectW = (x2 - x1) / imageWidth;
|
||||
double rectH = (y2 - y1) / imageHeight;
|
||||
|
||||
// 构建 Polygon 四个角点
|
||||
List<Point> pointsSrc = new ArrayList<>();
|
||||
pointsSrc.add(new Point(row[5], row[6]));
|
||||
pointsSrc.add(new Point(row[7], row[8]));
|
||||
pointsSrc.add(new Point(row[9], row[10]));
|
||||
pointsSrc.add(new Point(row[11], row[12]));
|
||||
|
||||
Landmark box = new Landmark(rectX, rectY, rectW, rectH, pointsSrc);
|
||||
classNames.add(className);
|
||||
probabilities.add(prob);
|
||||
boundingBoxes.add(box);
|
||||
}
|
||||
DetectedObjects detectedObjects = new DetectedObjects(classNames, probabilities, boundingBoxes);
|
||||
return detectedObjects;
|
||||
|
||||
}
|
||||
|
||||
@Override
|
||||
public Batchifier getBatchifier() {
|
||||
return null;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,186 @@
|
||||
package cn.smartjavaai.ocr.model.plate.translator;
|
||||
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.modality.cv.output.*;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDArrays;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.ndarray.NDManager;
|
||||
import ai.djl.ndarray.types.DataType;
|
||||
import ai.djl.translate.Batchifier;
|
||||
import ai.djl.translate.Translator;
|
||||
import ai.djl.translate.TranslatorContext;
|
||||
import cn.smartjavaai.common.utils.LetterBoxUtils;
|
||||
import cn.smartjavaai.common.utils.NMSUtils;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
/**
|
||||
* @author dwj
|
||||
*/
|
||||
public class Yolov8PlateDetectTranslator implements Translator<Image, DetectedObjects> {
|
||||
|
||||
private int inputSize = 640;
|
||||
private float minConfThreshold = 0.3f;
|
||||
private float iouThreshold = 0.5f;
|
||||
|
||||
private float confThreshold = 0;
|
||||
|
||||
private int imageWidth;
|
||||
private int imageHeight;
|
||||
|
||||
private int topK;
|
||||
|
||||
private LetterBoxUtils.ResizeResult letterBoxResult;
|
||||
|
||||
public Yolov8PlateDetectTranslator(Map<String, ?> arguments) {
|
||||
confThreshold =
|
||||
arguments.containsKey("confThreshold")
|
||||
? Integer.parseInt(arguments.get("confThreshold").toString())
|
||||
: 0.3f;
|
||||
|
||||
iouThreshold =
|
||||
arguments.containsKey("iouThreshold")
|
||||
? Integer.parseInt(arguments.get("iouThreshold").toString())
|
||||
: 0.5f;
|
||||
|
||||
topK = arguments.containsKey("topk")
|
||||
? Integer.parseInt(arguments.get("topk").toString())
|
||||
: 100;
|
||||
}
|
||||
|
||||
@Override
|
||||
public NDList processInput(TranslatorContext ctx, Image input) {
|
||||
NDManager manager = ctx.getNDManager();
|
||||
NDArray array = input.toNDArray(manager, Image.Flag.COLOR);
|
||||
imageWidth = (int) array.getShape().get(1);
|
||||
imageHeight = (int) array.getShape().get(0);
|
||||
//Letter box resize 640x640 with padding (保持比例,补边缘)
|
||||
letterBoxResult = LetterBoxUtils.letterbox(manager, array, inputSize, inputSize, 114f, LetterBoxUtils.PaddingPosition.CENTER);
|
||||
array = letterBoxResult.image;
|
||||
// 转为 float32 且归一化到 0~1
|
||||
array = array.toType(DataType.FLOAT32, false).div(255f); // HWC
|
||||
// HWC -> CHW
|
||||
array = array.transpose(2, 0, 1); // CHW
|
||||
return new NDList(array.expandDims(0));
|
||||
}
|
||||
|
||||
@Override
|
||||
public DetectedObjects processOutput(TranslatorContext ctx, NDList list) {
|
||||
NDManager manager = ctx.getNDManager();
|
||||
|
||||
NDArray preds = list.get(0); // shape: (1, 6, 8400)
|
||||
preds = preds.squeeze(0).transpose(1, 0); // shape: (8400, 6)
|
||||
|
||||
// preds shape: (8400, 6)
|
||||
NDArray classScores = preds.get(":, 4:6"); // shape: (8400, 2)
|
||||
|
||||
// 获取每行最大值(对应 Python 的 .amax(1))
|
||||
NDArray maxScores = classScores.max(new int[]{1}); // shape: (8400,)
|
||||
|
||||
// 构造 mask:score > conf
|
||||
NDArray confMask = maxScores.gt(minConfThreshold); // shape: (8400,)
|
||||
|
||||
// 应用 mask 筛选
|
||||
preds = preds.get(confMask); // shape: (N_filtered, 6)
|
||||
|
||||
if (preds.isEmpty()) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// 提取 box (xywh),转换为 xyxy
|
||||
NDArray boxes = preds.get(":, 0:4"); // shape: (N, 4)
|
||||
boxes = xywh2xyxy(boxes); // 自定义函数:center xywh -> xyxy
|
||||
|
||||
// 1. 得分和类别索引
|
||||
NDArray scoresAndClasses = preds.get(":, 4:6"); // shape (num, 2)
|
||||
NDArray scores = scoresAndClasses.max(new int[]{1}, true); // keepDim = true
|
||||
NDArray index = scoresAndClasses.argMax(1).expandDims(1); // 最大值索引,类别,shape (num, 1)
|
||||
|
||||
// 4. 拼接
|
||||
NDArray result = NDArrays.concat(new NDList(boxes, scores, index), 1); // 在列方向拼接
|
||||
|
||||
// NMS 过滤掉重叠框
|
||||
int[] keepIndices = NMSUtils.nms(boxes, scores.squeeze(), iouThreshold); // scores.squeeze() ➝ (N,)
|
||||
NDArray kept = result.get(manager.create(keepIndices));
|
||||
// 如果超过 topK,则截断
|
||||
if (keepIndices.length > topK) {
|
||||
int[] topkIndices = new int[topK];
|
||||
System.arraycopy(keepIndices, 0, topkIndices, 0, topK);
|
||||
keepIndices = topkIndices;
|
||||
}
|
||||
//恢复原图坐标(除回比例,减掉 padding)
|
||||
NDArray restored = LetterBoxUtils.restoreBox(kept, letterBoxResult.r, letterBoxResult.left, letterBoxResult.top, 5,0);
|
||||
|
||||
List<String> classNames = new ArrayList<>();
|
||||
List<Double> probabilities = new ArrayList<>();
|
||||
List<BoundingBox> boundingBoxes = new ArrayList<>();
|
||||
|
||||
float[] flatData = restored.toFloatArray();
|
||||
long[] shape = restored.getShape().getShape(); // 比如 (N, 14)
|
||||
int rows = (int) shape[0];
|
||||
int cols = (int) shape[1];
|
||||
|
||||
// 把一维数组重组为二维数组
|
||||
float[][] data = new float[rows][cols];
|
||||
for (int i = 0; i < rows; i++) {
|
||||
System.arraycopy(flatData, i * cols, data[i], 0, cols);
|
||||
}
|
||||
|
||||
for (float[] row : data) {
|
||||
// row结构:(x1, y1, x2, y2, score, classIndex)
|
||||
float x1 = row[0];
|
||||
float y1 = row[1];
|
||||
float x2 = row[2];
|
||||
float y2 = row[3];
|
||||
float score = row[4];
|
||||
int classIndex = (int) row[5];
|
||||
|
||||
double prob = score;
|
||||
String className = classIndex == 0 ? "single" : "double";
|
||||
|
||||
// 转相对坐标,DJL的Rectangle用比例坐标(0~1)
|
||||
double rectX = x1 / imageWidth;
|
||||
double rectY = y1 / imageHeight;
|
||||
double rectW = (x2 - x1) / imageWidth;
|
||||
double rectH = (y2 - y1) / imageHeight;
|
||||
|
||||
// 构建 Polygon 四个角点
|
||||
// List<Point> pointsSrc = new ArrayList<>();
|
||||
// pointsSrc.add(new Point(row[5], row[6]));
|
||||
// pointsSrc.add(new Point(row[7], row[8]));
|
||||
// pointsSrc.add(new Point(row[9], row[10]));
|
||||
// pointsSrc.add(new Point(row[11], row[12]));
|
||||
|
||||
Rectangle rectangle = new Rectangle(rectX, rectY, rectW, rectH);
|
||||
classNames.add(className);
|
||||
probabilities.add(prob);
|
||||
boundingBoxes.add(rectangle);
|
||||
}
|
||||
DetectedObjects detectedObjects = new DetectedObjects(classNames, probabilities, boundingBoxes);
|
||||
return detectedObjects;
|
||||
|
||||
}
|
||||
|
||||
@Override
|
||||
public Batchifier getBatchifier() {
|
||||
return null;
|
||||
}
|
||||
|
||||
|
||||
|
||||
public static NDArray xywh2xyxy(NDArray xywh) {
|
||||
NDArray x = xywh.get(":, 0");
|
||||
NDArray y = xywh.get(":, 1");
|
||||
NDArray w = xywh.get(":, 2").div(2);
|
||||
NDArray h = xywh.get(":, 3").div(2);
|
||||
NDArray x1 = x.sub(w);
|
||||
NDArray y1 = y.sub(h);
|
||||
NDArray x2 = x.add(w);
|
||||
NDArray y2 = y.add(h);
|
||||
return NDArrays.stack(new NDList(x1, y1, x2, y2), 1);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,171 @@
|
||||
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<Image, TableStructureResult> model;
|
||||
|
||||
private GenericObjectPool<Predictor<Image, TableStructureResult>> predictorPool;
|
||||
|
||||
@Override
|
||||
public void loadModel(TableStructureConfig config) {
|
||||
if(StringUtils.isBlank(config.getModelPath())){
|
||||
throw new OcrException("modelPath is null");
|
||||
}
|
||||
Criteria<Image, TableStructureResult> criteria = StructureCriteriaFactory.createCriteria(config);
|
||||
try{
|
||||
model = ModelZoo.loadModel(criteria);
|
||||
// 创建池子:每个线程独享 Predictor
|
||||
this.predictorPool = new GenericObjectPool<>(new PredictorFactory<>(model));
|
||||
int predictorPoolSize = config.getPredictorPoolSize();
|
||||
if(config.getPredictorPoolSize() <= 0){
|
||||
predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
|
||||
}
|
||||
predictorPool.setMaxTotal(predictorPoolSize);
|
||||
log.debug("当前设备: " + model.getNDManager().getDevice());
|
||||
log.debug("当前引擎: " + Engine.getInstance().getEngineName());
|
||||
log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
|
||||
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
|
||||
throw new OcrException("表格结构识别模型加载失败", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<TableStructureResult> 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<TableStructureResult> 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<TableStructureResult> 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<TableStructureResult> detect(Image image) {
|
||||
Predictor<Image, TableStructureResult> 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 GenericObjectPool<Predictor<Image, TableStructureResult>> getPool() {
|
||||
return predictorPool;
|
||||
}
|
||||
|
||||
@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);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,491 @@
|
||||
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.io.OutputStream;
|
||||
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<TableStructureResult> recognize(Image image) {
|
||||
//表格结构识别
|
||||
R<TableStructureResult> 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<String> tableContentList = buildTable(result.getData(), ocrInfo);
|
||||
String html = convertHtml(result.getData().getTableTagList(), tableContentList);
|
||||
result.getData().setHtml(html);
|
||||
return result;
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 表格识别
|
||||
* @param image
|
||||
* @return
|
||||
*/
|
||||
public R<TableStructureResult> 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<TableStructureResult> 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<TableStructureResult> 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);
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 绘制表格
|
||||
* @param tableStructureResult
|
||||
* @param image
|
||||
* @return
|
||||
*/
|
||||
public BufferedImage drawTable(TableStructureResult tableStructureResult, BufferedImage image){
|
||||
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);
|
||||
}
|
||||
return image;
|
||||
}
|
||||
|
||||
/**
|
||||
* 删除 HTML 中第一个 <style> ... </style> 段落
|
||||
* @param html 原始 HTML
|
||||
* @return 去掉 <style> 的 HTML
|
||||
*/
|
||||
public static String removeStyleBlock(String html) {
|
||||
String lowerHtml = html.toLowerCase();
|
||||
int styleStart = lowerHtml.indexOf("<style");
|
||||
if (styleStart == -1) {
|
||||
return html; // 没有 style,返回原文
|
||||
}
|
||||
int styleEnd = lowerHtml.indexOf("</style>", styleStart);
|
||||
if (styleEnd == -1) {
|
||||
return html; // 没闭合标签,不处理
|
||||
}
|
||||
styleEnd += "</style>".length();
|
||||
// 去掉 style 块
|
||||
return html.substring(0, styleStart) + html.substring(styleEnd);
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 导出 Excel
|
||||
* @param html
|
||||
* @param out
|
||||
*/
|
||||
public void exportExcel(String html, OutputStream out){
|
||||
String content = removeStyleBlock(html);
|
||||
content = content.replace("<html><body>", "");
|
||||
content = content.replace("</body></html>", "");
|
||||
try (HSSFWorkbook workbook = ConvertHtml2Excel.table2Excel(content)){
|
||||
workbook.write(out);
|
||||
out.flush();
|
||||
} catch (Exception e) {
|
||||
throw new OcrException("导出excel失败,请检查表结构是否识别正确");
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 导出 Excel
|
||||
* @param html
|
||||
* @param savePath
|
||||
*/
|
||||
public void exportExcel(String html, String savePath){
|
||||
String content = removeStyleBlock(html);
|
||||
content = content.replace("<html><body>", "");
|
||||
content = content.replace("</body></html>", "");
|
||||
try (HSSFWorkbook workbook = ConvertHtml2Excel.table2Excel(content)){
|
||||
workbook.write(new File(savePath));
|
||||
} catch (Exception e) {
|
||||
throw new OcrException("导出excel失败,请检查表结构是否识别正确");
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 构建表格
|
||||
* @param tableStructureResult
|
||||
* @param ocrInfo
|
||||
* @return
|
||||
*/
|
||||
public List<String> buildTable(TableStructureResult tableStructureResult, OcrInfo ocrInfo) {
|
||||
// 获取 Cell 与 文本检测框 的对应关系(1:N)。
|
||||
Map<Integer, List<Integer>> matched = new ConcurrentHashMap<>();
|
||||
List<OcrItem> 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<Pair<Float, Float>> 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<Float, Float> nearest = sorted(distances);
|
||||
|
||||
// 获取最小距离对应的下标id,也等价于cell的下标id (distances列表是根据遍历cells生成的)
|
||||
int id = 0;
|
||||
for (int idx = 0; idx < distances.size(); idx++) {
|
||||
Pair<Float, Float> 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<Integer> 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<String> cell_contents = new ArrayList<>();
|
||||
List<Double> probs = new ArrayList<>();
|
||||
for (int i = 0; i < tableStructureResult.getOcrItemList().size(); i++) {
|
||||
List<Integer> textIds = matched.get(i);
|
||||
List<String> 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<Float, Float> sorted(List<Pair<Float, Float>> distances) {
|
||||
Comparator<Pair<Float, Float>> comparator =
|
||||
new Comparator<Pair<Float, Float>>() {
|
||||
@Override
|
||||
public int compare(Pair<Float, Float> a1, Pair<Float, Float> 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<Pair<Float, Float>> 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<String> pred_structures, List<String> cell_contents) {
|
||||
StringBuffer html = new StringBuffer();
|
||||
// 添加统一的样式(可选放到<head>中)
|
||||
html.append("<style>\n");
|
||||
html.append("table { border-collapse: collapse; }\n");
|
||||
html.append("td, th, table { border: 1px solid black; padding: 5px; }\n");
|
||||
html.append("</style>\n");
|
||||
int td_index = 0;
|
||||
for (String tag : pred_structures) {
|
||||
if (tag.contains("<td></td>")) {
|
||||
String content = cell_contents.get(td_index);
|
||||
html.append("<td>");
|
||||
html.append(content);
|
||||
html.append("</td>");
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
package cn.smartjavaai.ocr.model.table;
|
||||
|
||||
import ai.djl.inference.Predictor;
|
||||
import ai.djl.modality.cv.Image;
|
||||
import cn.smartjavaai.common.entity.R;
|
||||
import cn.smartjavaai.ocr.config.OcrDetModelConfig;
|
||||
import cn.smartjavaai.ocr.config.TableStructureConfig;
|
||||
import cn.smartjavaai.ocr.entity.OcrBox;
|
||||
import cn.smartjavaai.ocr.entity.OcrItem;
|
||||
import cn.smartjavaai.ocr.entity.TableStructureResult;
|
||||
import org.apache.commons.pool2.impl.GenericObjectPool;
|
||||
|
||||
import java.awt.image.BufferedImage;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* 表格结构识别模型
|
||||
* @author dwj
|
||||
*/
|
||||
public interface TableStructureModel extends AutoCloseable{
|
||||
|
||||
/**
|
||||
* 加载模型
|
||||
* @param config
|
||||
*/
|
||||
void loadModel(TableStructureConfig config);
|
||||
|
||||
|
||||
/**
|
||||
* 表格结构检测
|
||||
* @param image
|
||||
* @return
|
||||
*/
|
||||
default R<TableStructureResult> detect(BufferedImage image){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 表格结构检测
|
||||
* @param imagePath 图片路径
|
||||
* @return
|
||||
*/
|
||||
default R<TableStructureResult> detect(String imagePath) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 表格结构检测
|
||||
* @param imageData 图片字节数组
|
||||
* @return
|
||||
*/
|
||||
default R<TableStructureResult> detect(byte[] imageData) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 表格结构检测
|
||||
* @param image DJL Image
|
||||
* @return
|
||||
*/
|
||||
default R<TableStructureResult> detect(Image image){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
default GenericObjectPool<Predictor<Image, TableStructureResult>> getPool() {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
}
|
||||
@@ -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<Image, TableStructureResult> createCriteria(TableStructureConfig config) {
|
||||
Device device = null;
|
||||
if(!Objects.isNull(config.getDevice())){
|
||||
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
|
||||
}
|
||||
Criteria<Image, TableStructureResult> 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;
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
@@ -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<Image, TableStructureResult> {
|
||||
|
||||
private final int maxLength = 488;
|
||||
private int height;
|
||||
private int width;
|
||||
private float scale = 1.0f;
|
||||
private float xScale;
|
||||
private float yScale;
|
||||
private List<String> dict;
|
||||
private String beg_str = "sos";
|
||||
private String end_str = "eos";
|
||||
private List<String> 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("<td>"))
|
||||
dict.remove("<td>");
|
||||
if(!dict.contains("<td></td>"))
|
||||
dict.add("<td></td>");
|
||||
dict.add(end_str);
|
||||
}
|
||||
|
||||
td_token.add("<td>");
|
||||
td_token.add("<td");
|
||||
td_token.add("<td></td>");
|
||||
}
|
||||
|
||||
@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<List<String>> structure_batch_list = new ArrayList<>();
|
||||
List<List<NDArray>> bbox_batch_list = new ArrayList<>();
|
||||
List<List<NDArray>> 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<String> structure_list = new ArrayList<>();
|
||||
List<NDArray> bbox_list = new ArrayList<>();
|
||||
List<NDArray> 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<String> structure_str_list =structure_batch_list.get(0);
|
||||
List<NDArray> bbox_list = bbox_batch_list.get(0);
|
||||
List<NDArray> score_list = result_score_list.get(0);
|
||||
|
||||
structure_str_list.add(0,"<html>");
|
||||
structure_str_list.add(1,"<body>");
|
||||
structure_str_list.add(2,"<table>");
|
||||
structure_str_list.add("</table>");
|
||||
structure_str_list.add("</body>");
|
||||
structure_str_list.add("</html>");
|
||||
|
||||
List<OcrItem> 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); // 前面加了<html><body><table> 所以偏移+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;
|
||||
}
|
||||
|
||||
}
|
||||
228
ocr/src/main/java/cn/smartjavaai/ocr/opencv/OcrNDArrayUtils.java
Normal file
228
ocr/src/main/java/cn/smartjavaai/ocr/opencv/OcrNDArrayUtils.java
Normal file
@@ -0,0 +1,228 @@
|
||||
package cn.smartjavaai.ocr.opencv;
|
||||
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import org.opencv.core.CvType;
|
||||
import org.opencv.core.Mat;
|
||||
import org.opencv.core.MatOfPoint;
|
||||
import org.opencv.core.Point;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
/**
|
||||
* NDArray Utils
|
||||
*
|
||||
*/
|
||||
public class OcrNDArrayUtils {
|
||||
/**
|
||||
* Mat To MatOfPoint
|
||||
* @param mat
|
||||
* @return
|
||||
*/
|
||||
public static MatOfPoint matToMatOfPoint(Mat mat) {
|
||||
int rows = mat.rows();
|
||||
MatOfPoint matOfPoint = new MatOfPoint();
|
||||
|
||||
List<Point> list = new ArrayList<>();
|
||||
for (int i = 0; i < rows; i++) {
|
||||
Point point = new Point((float) mat.get(i, 0)[0], (float) mat.get(i, 1)[0]);
|
||||
list.add(point);
|
||||
}
|
||||
matOfPoint.fromList(list);
|
||||
|
||||
return matOfPoint;
|
||||
}
|
||||
|
||||
/**
|
||||
* float NDArray To float[][] Array
|
||||
* @param ndArray
|
||||
* @return
|
||||
*/
|
||||
public static float[][] floatNDArrayToArray(NDArray ndArray) {
|
||||
int rows = (int) (ndArray.getShape().get(0));
|
||||
int cols = (int) (ndArray.getShape().get(1));
|
||||
float[][] arr = new float[rows][cols];
|
||||
|
||||
float[] arrs = ndArray.toFloatArray();
|
||||
for (int i = 0; i < rows; i++) {
|
||||
for (int j = 0; j < cols; j++) {
|
||||
arr[i][j] = arrs[i * cols + j];
|
||||
}
|
||||
}
|
||||
return arr;
|
||||
}
|
||||
|
||||
/**
|
||||
* Mat To double[][] Array
|
||||
* @param mat
|
||||
* @return
|
||||
*/
|
||||
public static double[][] matToDoubleArray(Mat mat) {
|
||||
int rows = mat.rows();
|
||||
int cols = mat.cols();
|
||||
|
||||
double[][] doubles = new double[rows][cols];
|
||||
|
||||
for (int i = 0; i < rows; i++) {
|
||||
for (int j = 0; j < cols; j++) {
|
||||
doubles[i][j] = mat.get(i, j)[0];
|
||||
}
|
||||
}
|
||||
|
||||
return doubles;
|
||||
}
|
||||
|
||||
/**
|
||||
* Mat To float[][] Array
|
||||
* @param mat
|
||||
* @return
|
||||
*/
|
||||
public static float[][] matToFloatArray(Mat mat) {
|
||||
int rows = mat.rows();
|
||||
int cols = mat.cols();
|
||||
|
||||
float[][] floats = new float[rows][cols];
|
||||
|
||||
for (int i = 0; i < rows; i++) {
|
||||
for (int j = 0; j < cols; j++) {
|
||||
floats[i][j] = (float) mat.get(i, j)[0];
|
||||
}
|
||||
}
|
||||
|
||||
return floats;
|
||||
}
|
||||
|
||||
/**
|
||||
* Mat To byte[][] Array
|
||||
* @param mat
|
||||
* @return
|
||||
*/
|
||||
public static byte[][] matToUint8Array(Mat mat) {
|
||||
int rows = mat.rows();
|
||||
int cols = mat.cols();
|
||||
|
||||
byte[][] bytes = new byte[rows][cols];
|
||||
|
||||
for (int i = 0; i < rows; i++) {
|
||||
for (int j = 0; j < cols; j++) {
|
||||
bytes[i][j] = (byte) mat.get(i, j)[0];
|
||||
}
|
||||
}
|
||||
|
||||
return bytes;
|
||||
}
|
||||
|
||||
/**
|
||||
* float NDArray To float[][] Array
|
||||
* @param ndArray
|
||||
* @param cvType
|
||||
* @return
|
||||
*/
|
||||
public static Mat floatNDArrayToMat(NDArray ndArray, int cvType) {
|
||||
int rows = (int) (ndArray.getShape().get(0));
|
||||
int cols = (int) (ndArray.getShape().get(1));
|
||||
Mat mat = new Mat(rows, cols, cvType);
|
||||
|
||||
float[] arrs = ndArray.toFloatArray();
|
||||
for (int i = 0; i < rows; i++) {
|
||||
for (int j = 0; j < cols; j++) {
|
||||
mat.put(i, j, arrs[i * cols + j]);
|
||||
}
|
||||
}
|
||||
return mat;
|
||||
}
|
||||
|
||||
/**
|
||||
* float NDArray To Mat
|
||||
* @param ndArray
|
||||
* @return
|
||||
*/
|
||||
public static Mat floatNDArrayToMat(NDArray ndArray) {
|
||||
int rows = (int) (ndArray.getShape().get(0));
|
||||
int cols = (int) (ndArray.getShape().get(1));
|
||||
Mat mat = new Mat(rows, cols, CvType.CV_32F);
|
||||
|
||||
float[] arrs = ndArray.toFloatArray();
|
||||
for (int i = 0; i < rows; i++) {
|
||||
for (int j = 0; j < cols; j++) {
|
||||
mat.put(i, j, arrs[i * cols + j]);
|
||||
}
|
||||
}
|
||||
|
||||
return mat;
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* uint8 NDArray To Mat
|
||||
* @param ndArray
|
||||
* @return
|
||||
*/
|
||||
public static Mat uint8NDArrayToMat(NDArray ndArray) {
|
||||
int rows = (int) (ndArray.getShape().get(0));
|
||||
int cols = (int) (ndArray.getShape().get(1));
|
||||
Mat mat = new Mat(rows, cols, CvType.CV_8U);
|
||||
|
||||
byte[] arrs = ndArray.toByteArray();
|
||||
|
||||
for (int i = 0; i < rows; i++) {
|
||||
for (int j = 0; j < cols; j++) {
|
||||
mat.put(i, j, arrs[i * cols + j]);
|
||||
}
|
||||
}
|
||||
return mat;
|
||||
}
|
||||
|
||||
/**
|
||||
* float[][] Array To Mat
|
||||
* @param arr
|
||||
* @return
|
||||
*/
|
||||
public static Mat floatArrayToMat(float[][] arr) {
|
||||
int rows = arr.length;
|
||||
int cols = arr[0].length;
|
||||
Mat mat = new Mat(rows, cols, CvType.CV_32F);
|
||||
|
||||
for (int i = 0; i < rows; i++) {
|
||||
for (int j = 0; j < cols; j++) {
|
||||
mat.put(i, j, arr[i][j]);
|
||||
}
|
||||
}
|
||||
|
||||
return mat;
|
||||
}
|
||||
|
||||
/**
|
||||
* byte[][] Array To Mat
|
||||
* @param arr
|
||||
* @return
|
||||
*/
|
||||
public static Mat uint8ArrayToMat(byte[][] arr) {
|
||||
int rows = arr.length;
|
||||
int cols = arr[0].length;
|
||||
Mat mat = new Mat(rows, cols, CvType.CV_8U);
|
||||
|
||||
for (int i = 0; i < rows; i++) {
|
||||
for (int j = 0; j < cols; j++) {
|
||||
mat.put(i, j, arr[i][j]);
|
||||
}
|
||||
}
|
||||
|
||||
return mat;
|
||||
}
|
||||
|
||||
/**
|
||||
* List To Mat
|
||||
* @param points
|
||||
* @return
|
||||
*/
|
||||
public static Mat toMat(List<ai.djl.modality.cv.output.Point> points) {
|
||||
Mat mat = new Mat(points.size(), 2, CvType.CV_32F);
|
||||
for (int i = 0; i < points.size(); i++) {
|
||||
ai.djl.modality.cv.output.Point point = points.get(i);
|
||||
mat.put(i, 0, (float) point.getX());
|
||||
mat.put(i, 1, (float) point.getY());
|
||||
}
|
||||
|
||||
return mat;
|
||||
}
|
||||
}
|
||||
@@ -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 如
|
||||
* <table>
|
||||
* ..
|
||||
* </table>
|
||||
* @return
|
||||
*/
|
||||
public static HSSFWorkbook table2Excel(String tableHtml) {
|
||||
HSSFWorkbook wb = new HSSFWorkbook();
|
||||
HSSFSheet sheet = wb.createSheet();
|
||||
List<CrossRangeCellMeta> 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<Element> trLs = thead.elements("tr");
|
||||
for (Element trEle : trLs) {
|
||||
HSSFRow row = sheet.createRow(rowIndex);
|
||||
List<Element> 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<Element> trLs = tbody.elements("tr");
|
||||
for (Element trEle : trLs) {
|
||||
HSSFRow row = sheet.createRow(rowIndex);
|
||||
List<Element> thLs = trEle.elements("th");
|
||||
int cellIndex = makeRowCell(thLs, rowIndex, row, 0, titleStyle, crossRowEleMetaLs);
|
||||
List<Element> 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<sheet.getRow(0).getPhysicalNumberOfCells();i++){
|
||||
sheet.autoSizeColumn(i, true);//设置列宽
|
||||
if(sheet.getColumnWidth(i)<255*256){
|
||||
sheet.setColumnWidth(i, sheet.getColumnWidth(i) < 9000 ? 9000 : sheet.getColumnWidth(i));
|
||||
}else{
|
||||
sheet.setColumnWidth(i, 15000);
|
||||
}
|
||||
}
|
||||
} catch (DocumentException e) {
|
||||
e.printStackTrace();
|
||||
}
|
||||
|
||||
return wb;
|
||||
}
|
||||
|
||||
/**
|
||||
* 生产行内容
|
||||
*
|
||||
* @return 最后一列的cell index
|
||||
*/
|
||||
/**
|
||||
* @param tdLs th或者td集合
|
||||
* @param rowIndex 行号
|
||||
* @param row POI行对象
|
||||
* @param startCellIndex
|
||||
* @param cellStyle 样式
|
||||
* @param crossRowEleMetaLs 跨行元数据集合
|
||||
* @return
|
||||
*/
|
||||
private static int makeRowCell(List<Element> tdLs, int rowIndex, HSSFRow row, int startCellIndex, HSSFCellStyle cellStyle,
|
||||
List<CrossRangeCellMeta> 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<CrossRangeCellMeta> 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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
457
ocr/src/main/java/cn/smartjavaai/ocr/utils/OcrUtils.java
Normal file
457
ocr/src/main/java/cn/smartjavaai/ocr/utils/OcrUtils.java
Normal file
@@ -0,0 +1,457 @@
|
||||
package cn.smartjavaai.ocr.utils;
|
||||
|
||||
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.Landmark;
|
||||
import ai.djl.modality.cv.util.NDImageUtils;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.ndarray.NDManager;
|
||||
import ai.djl.opencv.OpenCVImageFactory;
|
||||
import cn.smartjavaai.common.entity.*;
|
||||
import cn.smartjavaai.common.entity.Point;
|
||||
import cn.smartjavaai.common.utils.ImageUtils;
|
||||
import cn.smartjavaai.common.utils.OpenCVUtils;
|
||||
import cn.smartjavaai.common.utils.PointUtils;
|
||||
import cn.smartjavaai.ocr.entity.*;
|
||||
import cn.smartjavaai.ocr.entity.RotatedBox;
|
||||
import cn.smartjavaai.ocr.enums.AngleEnum;
|
||||
import cn.smartjavaai.ocr.enums.PlateType;
|
||||
import cn.smartjavaai.ocr.opencv.OcrNDArrayUtils;
|
||||
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;
|
||||
|
||||
import java.awt.*;
|
||||
import java.awt.image.BufferedImage;
|
||||
import java.math.BigDecimal;
|
||||
import java.util.*;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* @author dwj
|
||||
* @date 2025/4/22
|
||||
*/
|
||||
@Slf4j
|
||||
public class OcrUtils {
|
||||
|
||||
|
||||
/**
|
||||
* 转换为OcrBox
|
||||
* @param dt_boxes
|
||||
* @return
|
||||
*/
|
||||
public static List<OcrBox> convertToOcrBox(NDList dt_boxes) {
|
||||
List<OcrBox> boxList = new ArrayList<>();
|
||||
for (NDArray box : dt_boxes) {
|
||||
float[] pointsArr = box.toFloatArray();
|
||||
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 ndLists
|
||||
* @return
|
||||
*/
|
||||
public static List<List<OcrBox>> convertToOcrBox(List<NDList> ndLists) {
|
||||
if (ndLists == null || ndLists.isEmpty()) {
|
||||
return Collections.emptyList();
|
||||
}
|
||||
List<List<OcrBox>> boxLists = new ArrayList<>();
|
||||
for (NDList dt_boxes : ndLists) {
|
||||
boxLists.add(convertToOcrBox(dt_boxes));
|
||||
}
|
||||
return boxLists;
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* 图片旋转
|
||||
*
|
||||
* @param manager
|
||||
* @param image
|
||||
* @return
|
||||
*/
|
||||
public static Image rotateImg(NDManager manager, Image image) {
|
||||
NDArray rotated = NDImageUtils.rotate90(image.toNDArray(manager), 1);
|
||||
return ImageFactory.getInstance().fromNDArray(rotated);
|
||||
}
|
||||
|
||||
/**
|
||||
* 逆时针旋转图片
|
||||
*
|
||||
* @param image
|
||||
* @param times
|
||||
* @return
|
||||
*/
|
||||
public static Image rotateImg(Image image, int times) {
|
||||
try (NDManager manager = NDManager.newBaseManager()) {
|
||||
NDArray rotated = NDImageUtils.rotate90(image.toNDArray(manager), times);
|
||||
return OpenCVImageFactory.getInstance().fromNDArray(rotated);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 逆时针旋转图片
|
||||
*
|
||||
* @param image
|
||||
* @param angleEnum
|
||||
* @return
|
||||
*/
|
||||
public static Image rotateImg(Image image, AngleEnum angleEnum) {
|
||||
try (NDManager manager = NDManager.newBaseManager()) {
|
||||
int times = 0;
|
||||
switch (angleEnum) {
|
||||
case ANGLE_90:
|
||||
times = 1;
|
||||
break;
|
||||
case ANGLE_180:
|
||||
times = 2;
|
||||
break;
|
||||
case ANGLE_270:
|
||||
times = 3;
|
||||
break;
|
||||
}
|
||||
NDArray rotated = NDImageUtils.rotate90(image.toNDArray(manager), times);
|
||||
return OpenCVImageFactory.getInstance().fromNDArray(rotated);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 转换为OcrInfo
|
||||
* @param lines
|
||||
* @return
|
||||
*/
|
||||
public static OcrInfo convertToOcrInfo(List<ArrayList<RotatedBoxCompX>> lines){
|
||||
if(Objects.isNull(lines) || lines.size() == 0){
|
||||
return null;
|
||||
}
|
||||
List<List<OcrItem>> lineList = new ArrayList<List<OcrItem>>();
|
||||
String fullText = "";
|
||||
for(ArrayList<RotatedBoxCompX> boxList : lines){
|
||||
List<OcrItem> line = new ArrayList<OcrItem>();
|
||||
for(RotatedBoxCompX box : boxList){
|
||||
float[] pointsArr = box.getBox().toFloatArray();
|
||||
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]));
|
||||
OcrItem ocrItem = new OcrItem(ocrBox, box.getText());
|
||||
line.add(ocrItem);
|
||||
String text = box.getText();
|
||||
if(text.trim().equals(""))
|
||||
continue;
|
||||
fullText += text + " ";
|
||||
}
|
||||
lineList.add(line);
|
||||
fullText += '\n';
|
||||
}
|
||||
return new OcrInfo(lineList, fullText);
|
||||
}
|
||||
|
||||
public static OcrInfo convertRotatedBoxesToOcrItems(List<RotatedBox> rotatedBoxes) {
|
||||
OcrInfo ocrInfo = new OcrInfo();
|
||||
List<OcrItem> 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;
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 透视变换 + 裁剪
|
||||
* @param srcMat
|
||||
* @param landMarks
|
||||
* @return
|
||||
*/
|
||||
public static Image transformAndCrop(Mat srcMat, List<ai.djl.modality.cv.output.Point> landMarks){
|
||||
if (landMarks == null || landMarks.size() != 4) {
|
||||
throw new IllegalArgumentException("必须提供4个关键点");
|
||||
}
|
||||
|
||||
// 步骤 1:排序为 左上、右上、右下、左下
|
||||
List<ai.djl.modality.cv.output.Point> ordered = PointUtils.orderPoints(landMarks);
|
||||
|
||||
ai.djl.modality.cv.output.Point lt = ordered.get(0);
|
||||
ai.djl.modality.cv.output.Point rt = ordered.get(1);
|
||||
ai.djl.modality.cv.output.Point rb = ordered.get(2);
|
||||
ai.djl.modality.cv.output.Point lb = ordered.get(3);
|
||||
|
||||
// 步骤 2:计算目标图像尺寸(宽、高)
|
||||
int img_crop_width = (int) Math.max(
|
||||
PointUtils.distance(lt, rt),
|
||||
PointUtils.distance(rb, lb)
|
||||
);
|
||||
int img_crop_height = (int) Math.max(
|
||||
PointUtils.distance(lt, lb),
|
||||
PointUtils.distance(rt, rb)
|
||||
);
|
||||
|
||||
// 步骤 3:构造目标坐标点
|
||||
List<ai.djl.modality.cv.output.Point> dstPoints = Arrays.asList(
|
||||
new ai.djl.modality.cv.output.Point(0, 0),
|
||||
new ai.djl.modality.cv.output.Point(img_crop_width, 0),
|
||||
new ai.djl.modality.cv.output.Point(img_crop_width, img_crop_height),
|
||||
new ai.djl.modality.cv.output.Point(0, img_crop_height)
|
||||
);
|
||||
|
||||
// 步骤 4:透视变换
|
||||
Mat srcPoint2f = OcrNDArrayUtils.toMat(ordered);
|
||||
Mat dstPoint2f = OcrNDArrayUtils.toMat(dstPoints);
|
||||
Mat cvMat = OpenCVUtils.perspectiveTransform(srcMat, srcPoint2f, dstPoint2f);
|
||||
|
||||
// 步骤 5:转为 DJL Image + 裁剪
|
||||
Image subImg = OpenCVImageFactory.getInstance().fromImage(cvMat);
|
||||
subImg = subImg.getSubImage(0, 0, img_crop_width, img_crop_height);
|
||||
|
||||
// 释放资源
|
||||
cvMat.release();
|
||||
srcPoint2f.release();
|
||||
dstPoint2f.release();
|
||||
|
||||
return subImg;
|
||||
}
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* 透视变换+裁剪
|
||||
* @param srcMat
|
||||
* @param box
|
||||
* @return
|
||||
*/
|
||||
public static Image transformAndCrop(Mat srcMat, OcrBox box){
|
||||
float[] pointsArr = box.toFloatArray();
|
||||
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);
|
||||
int img_crop_width = (int) Math.max(PointUtils.distance(lt, rt), PointUtils.distance(rb, lb));
|
||||
int img_crop_height = (int) Math.max(PointUtils.distance(lt, lb), PointUtils.distance(rt, rb));
|
||||
List<ai.djl.modality.cv.output.Point> srcPoints = new ArrayList<>();
|
||||
srcPoints.add(new ai.djl.modality.cv.output.Point(lt[0], lt[1]));
|
||||
srcPoints.add(new ai.djl.modality.cv.output.Point(rt[0], rt[1]));
|
||||
srcPoints.add(new ai.djl.modality.cv.output.Point(rb[0], rb[1]));
|
||||
srcPoints.add(new ai.djl.modality.cv.output.Point(lb[0], lb[1]));
|
||||
List<ai.djl.modality.cv.output.Point> dstPoints = new ArrayList<>();
|
||||
dstPoints.add(new ai.djl.modality.cv.output.Point(0, 0));
|
||||
dstPoints.add(new ai.djl.modality.cv.output.Point(img_crop_width, 0));
|
||||
dstPoints.add(new ai.djl.modality.cv.output.Point(img_crop_width, img_crop_height));
|
||||
dstPoints.add(new ai.djl.modality.cv.output.Point(0, img_crop_height));
|
||||
Mat srcPoint2f = OcrNDArrayUtils.toMat(srcPoints);
|
||||
Mat dstPoint2f = OcrNDArrayUtils.toMat(dstPoints);
|
||||
//透视变换
|
||||
Mat cvMat = OpenCVUtils.perspectiveTransform(srcMat, srcPoint2f, dstPoint2f);
|
||||
Image subImg = OpenCVImageFactory.getInstance().fromImage(cvMat);
|
||||
//ImageUtils.saveImage(subImg, i + ".png", "build/output");
|
||||
//变换后裁剪
|
||||
subImg = subImg.getSubImage(0, 0, img_crop_width, img_crop_height);
|
||||
cvMat.release();
|
||||
srcPoint2f.release();
|
||||
dstPoint2f.release();
|
||||
return subImg;
|
||||
}
|
||||
|
||||
/**
|
||||
* 绘制文本框
|
||||
*
|
||||
* @param mat
|
||||
* @param boxList
|
||||
*/
|
||||
public static void drawRect(Mat mat, List<OcrBox> boxList) {
|
||||
for(OcrBox ocrBox : boxList){
|
||||
Imgproc.line(mat, ocrBox.getTopLeft().toCvPoint(), ocrBox.getTopRight().toCvPoint(), new Scalar(0, 255, 0), 1);
|
||||
Imgproc.line(mat, ocrBox.getTopRight().toCvPoint(), ocrBox.getBottomRight().toCvPoint(), new Scalar(0, 255, 0),1);
|
||||
Imgproc.line(mat, ocrBox.getBottomRight().toCvPoint(), ocrBox.getBottomLeft().toCvPoint(), new Scalar(0, 255, 0),1);
|
||||
Imgproc.line(mat, ocrBox.getBottomLeft().toCvPoint(), ocrBox.getTopLeft().toCvPoint(), new Scalar(0, 255, 0), 1);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 绘制文本框及文本
|
||||
* @param image
|
||||
* @param ocrInfo
|
||||
*/
|
||||
public static void drawRectWithText(BufferedImage image, OcrInfo ocrInfo, int fontSize) {
|
||||
// 将绘制图像转换为Graphics2D
|
||||
Graphics2D g = (Graphics2D) image.getGraphics();
|
||||
try {
|
||||
Font font = new Font("楷体", Font.PLAIN, fontSize);
|
||||
g.setFont(font);
|
||||
g.setColor(new Color(0, 0, 255));
|
||||
// 声明画笔属性 :粗 细(单位像素)末端无修饰 折线处呈尖角
|
||||
BasicStroke bStroke = new BasicStroke(2, BasicStroke.CAP_BUTT, BasicStroke.JOIN_MITER);
|
||||
g.setStroke(bStroke);
|
||||
List<OcrItem> 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();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 绘制文本框及文本
|
||||
* @param srcMat
|
||||
* @param itemList
|
||||
*/
|
||||
public static void drawRectWithText(Mat srcMat, List<OcrItem> itemList) {
|
||||
for(OcrItem item : itemList){
|
||||
OcrBox ocrBox = item.getOcrBox();
|
||||
Imgproc.line(srcMat, ocrBox.getTopLeft().toCvPoint(), ocrBox.getTopRight().toCvPoint(), new Scalar(0, 255, 0), 1);
|
||||
Imgproc.line(srcMat, ocrBox.getTopRight().toCvPoint(), ocrBox.getBottomRight().toCvPoint(), new Scalar(0, 255, 0),1);
|
||||
Imgproc.line(srcMat, ocrBox.getBottomRight().toCvPoint(), ocrBox.getBottomLeft().toCvPoint(), new Scalar(0, 255, 0),1);
|
||||
Imgproc.line(srcMat, ocrBox.getBottomLeft().toCvPoint(), ocrBox.getTopLeft().toCvPoint(), new Scalar(0, 255, 0), 1);
|
||||
// 中文乱码
|
||||
Imgproc.putText(srcMat, item.getAngle().getValue(), ocrBox.getTopLeft().toCvPoint(), Imgproc.FONT_HERSHEY_SCRIPT_SIMPLEX, 1.0, new Scalar(0, 255, 0), 1);
|
||||
}
|
||||
}
|
||||
|
||||
public static List<PlateInfo> convertToPlateInfo(DetectedObjects detectedObjects, Image image) {
|
||||
List<PlateInfo> plateInfoList = new ArrayList<>();
|
||||
Iterator iterator = detectedObjects.items().iterator();
|
||||
int index = 0;
|
||||
while(iterator.hasNext()) {
|
||||
DetectedObjects.DetectedObject result = (DetectedObjects.DetectedObject)iterator.next();
|
||||
BoundingBox box = result.getBoundingBox();
|
||||
List<Point> keyPoints = new ArrayList<Point>();
|
||||
if(box instanceof Landmark){
|
||||
box.getBounds().getPath().forEach(point -> {
|
||||
keyPoints.add(new Point(point.getX(), point.getY()));
|
||||
});
|
||||
}
|
||||
int x = (int)(box.getBounds().getX() * image.getWidth());
|
||||
int y = (int)(box.getBounds().getY() * image.getHeight());
|
||||
int width = (int)(box.getBounds().getWidth() * image.getWidth());
|
||||
int height = (int)(box.getBounds().getHeight() * image.getHeight());
|
||||
// 修正边界,防止越界
|
||||
if (x < 0) x = 0;
|
||||
if (y < 0) y = 0;
|
||||
if (x + width > image.getWidth()) width = image.getWidth() - x;
|
||||
if (y + height > image.getHeight()) height = image.getHeight() - y;
|
||||
PlateInfo plateInfo = new PlateInfo();
|
||||
plateInfo.setPlateType(PlateType.fromClassName(detectedObjects.getClassNames().get(index)));
|
||||
plateInfo.setScore(detectedObjects.getProbabilities().get(index).floatValue());
|
||||
plateInfo.setDetectionRectangle(new DetectionRectangle(x, y, width, height));
|
||||
OcrBox ocrBox = new OcrBox(keyPoints.get(0), keyPoints.get(1), keyPoints.get(2), keyPoints.get(3));
|
||||
plateInfo.setBox(ocrBox);
|
||||
plateInfoList.add(plateInfo);
|
||||
index++;
|
||||
}
|
||||
return plateInfoList;
|
||||
}
|
||||
|
||||
/**
|
||||
* 绘制车牌信息
|
||||
* @param srcMat
|
||||
* @param plateInfoList
|
||||
*/
|
||||
public static void drawPlateInfo(Mat srcMat, List<PlateInfo> plateInfoList) {
|
||||
for(PlateInfo plateInfo : plateInfoList){
|
||||
OcrBox ocrBox = plateInfo.getBox();
|
||||
Imgproc.line(srcMat, ocrBox.getTopLeft().toCvPoint(), ocrBox.getTopRight().toCvPoint(), new Scalar(0, 0, 255), 1);
|
||||
Imgproc.line(srcMat, ocrBox.getTopRight().toCvPoint(), ocrBox.getBottomRight().toCvPoint(), new Scalar(0, 0, 255),1);
|
||||
Imgproc.line(srcMat, ocrBox.getBottomRight().toCvPoint(), ocrBox.getBottomLeft().toCvPoint(), new Scalar(0, 0, 255),1);
|
||||
Imgproc.line(srcMat, ocrBox.getBottomLeft().toCvPoint(), ocrBox.getTopLeft().toCvPoint(), new Scalar(0, 0, 255), 1);
|
||||
// 中文乱码
|
||||
ImageUtils.putTextWithBackground(srcMat, plateInfo.getPlateNumber() + " " + plateInfo.getPlateColor(), ocrBox.getTopLeft().toCvPoint(), new Scalar(255, 255, 255), new Scalar(0, 0, 0), 1);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 在图像上绘制带白色背景、黑色文字的文本
|
||||
*/
|
||||
public static void drawPlateInfo(BufferedImage image, List<PlateInfo> plateInfoList) {
|
||||
// 将绘制图像转换为Graphics2D
|
||||
Graphics2D graphics = (Graphics2D) image.getGraphics();
|
||||
try {
|
||||
graphics.setColor(Color.RED);// 边框颜色
|
||||
graphics.setStroke(new BasicStroke(2)); // 线宽2像素
|
||||
graphics.setRenderingHint(RenderingHints.KEY_ANTIALIASING,
|
||||
RenderingHints.VALUE_ANTIALIAS_ON); // 抗锯齿
|
||||
int stroke = 2;
|
||||
for(PlateInfo plateInfo : plateInfoList){
|
||||
DetectionRectangle rectangle = plateInfo.getDetectionRectangle();
|
||||
graphics.setColor(Color.RED);// 边框颜色
|
||||
//绘制车牌框
|
||||
graphics.drawRect(rectangle.getX(), rectangle.getY(), rectangle.getWidth(), rectangle.getHeight());
|
||||
graphics.setColor(Color.BLACK);// 字体颜色
|
||||
ImageUtils.drawText(graphics, plateInfo.getPlateNumber() + " " + plateInfo.getPlateColor(), rectangle.getX(), rectangle.getY(), stroke, 4);
|
||||
OcrBox ocrBox = plateInfo.getBox();
|
||||
//绘制关键点
|
||||
graphics.setColor(Color.BLUE);
|
||||
graphics.drawRect((int)ocrBox.getTopLeft().getX(), (int)ocrBox.getTopLeft().getY(), 2, 2);
|
||||
graphics.setColor(Color.GREEN);
|
||||
graphics.drawRect((int)ocrBox.getTopRight().getX(), (int)ocrBox.getTopRight().getY(), 2, 2);
|
||||
graphics.setColor(Color.RED);
|
||||
graphics.drawRect((int)ocrBox.getBottomLeft().getX(), (int)ocrBox.getBottomLeft().getY(), 2, 2);
|
||||
graphics.setColor(Color.CYAN);
|
||||
graphics.drawRect((int)ocrBox.getBottomRight().getX(), (int)ocrBox.getBottomRight().getY(), 2, 2);
|
||||
}
|
||||
} finally {
|
||||
graphics.dispose();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
Reference in New Issue
Block a user