package com.xindao.ocr.smartjavaai.model.common.direction;
|
|
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 cn.smartjavaai.common.pool.PredictorFactory;
|
import cn.smartjavaai.common.utils.FileUtils;
|
import cn.smartjavaai.common.utils.ImageUtils;
|
import cn.smartjavaai.common.utils.OpenCVUtils;
|
import com.xindao.ocr.smartjavaai.config.DirectionModelConfig;
|
import com.xindao.ocr.smartjavaai.entity.DirectionInfo;
|
import com.xindao.ocr.smartjavaai.entity.OcrBox;
|
import com.xindao.ocr.smartjavaai.entity.OcrItem;
|
import com.xindao.ocr.smartjavaai.enums.AngleEnum;
|
import com.xindao.ocr.smartjavaai.exception.OcrException;
|
import com.xindao.ocr.smartjavaai.model.common.detect.OcrCommonDetModel;
|
import com.xindao.ocr.smartjavaai.model.common.direction.criteria.DirectionCriteriaFactory;
|
import com.xindao.ocr.smartjavaai.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.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);
|
}
|
}
|
}
|
