package com.xindao.ocr.swingui.service;
|
|
import cn.smartjavaai.common.enums.DeviceEnum;
|
import com.alibaba.fastjson.JSONObject;
|
import com.xindao.ocr.smartjavaai.config.DirectionModelConfig;
|
import com.xindao.ocr.smartjavaai.config.OcrDetModelConfig;
|
import com.xindao.ocr.smartjavaai.config.OcrRecModelConfig;
|
import com.xindao.ocr.smartjavaai.config.OcrRecOptions;
|
import com.xindao.ocr.smartjavaai.entity.OcrInfo;
|
import com.xindao.ocr.smartjavaai.enums.CommonDetModelEnum;
|
import com.xindao.ocr.smartjavaai.enums.CommonRecModelEnum;
|
import com.xindao.ocr.smartjavaai.enums.DirectionModelEnum;
|
import com.xindao.ocr.smartjavaai.factory.OcrModelFactory;
|
import com.xindao.ocr.smartjavaai.model.common.detect.OcrCommonDetModel;
|
import com.xindao.ocr.smartjavaai.model.common.direction.OcrDirectionModel;
|
import com.xindao.ocr.smartjavaai.model.common.recognize.OcrCommonRecModel;
|
import com.xindao.ocr.swingui.constant.OcrSwingConstants;
|
import lombok.RequiredArgsConstructor;
|
import lombok.extern.slf4j.Slf4j;
|
import org.springframework.core.io.ClassPathResource;
|
import org.springframework.core.io.Resource;
|
import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
|
import org.springframework.stereotype.Service;
|
|
import javax.annotation.PostConstruct;
|
import java.io.File;
|
import java.io.IOException;
|
import java.nio.file.Files;
|
import java.nio.file.StandardCopyOption;
|
|
@Slf4j
|
@Service
|
@RequiredArgsConstructor
|
public class OcrService {
|
|
public static DeviceEnum device = DeviceEnum.CPU;
|
|
private final static String MAIN_PP_DIR = "PP_OCRv5";
|
|
private OcrCommonRecModel recModel;
|
|
/**
|
* 初始化加载ocr模型
|
*/
|
@PostConstruct
|
private void init() {
|
log.info("复制ocr文件到本地用户缓存...");
|
// 复制PP_OCRv5目录到本地缓存
|
copyPaddleCppToCache();
|
try {
|
recModel = getRecModel();
|
} catch (IOException e) {
|
log.error("加载OCR模型失败: {}", e.getMessage());
|
e.printStackTrace();
|
}
|
}
|
|
/**
|
* 复制resources中的Paddle_CPP目录到本地缓存
|
*/
|
private void copyPaddleCppToCache() {
|
try {
|
// 创建缓存目录 - 使用用户主目录避免权限问题
|
File cacheDir = OcrSwingConstants.ocrDir;
|
if (!cacheDir.exists()) {
|
cacheDir.mkdirs();
|
}
|
|
// 获取resources中的Paddle_CPP目录资源
|
PathMatchingResourcePatternResolver resolver = new PathMatchingResourcePatternResolver();
|
Resource[] resources = resolver.getResources("classpath:"+MAIN_PP_DIR+"/**");
|
|
// 复制所有资源到缓存目录，保持目录结构
|
for (Resource resource : resources) {
|
String resourcePath = resource.getURL().getPath();
|
// 获取相对于Paddle_CPP的路径
|
int startIndex = resourcePath.indexOf(MAIN_PP_DIR) + MAIN_PP_DIR.length();
|
String relativePath = resourcePath.substring(startIndex);
|
// 处理Windows路径分隔符
|
relativePath = relativePath.replace('/', File.separatorChar);
|
// 创建目标文件
|
File destFile = new File(cacheDir, relativePath);
|
// 如果是目录，创建目录
|
if (resource.isReadable() && resource.contentLength() == 0 && relativePath.endsWith(File.separator)) {
|
if (!destFile.exists()) {
|
destFile.mkdirs();
|
log.info("创建目录: {}", destFile.getAbsolutePath());
|
}
|
} else if (resource.isReadable()) {
|
// 确保父目录存在
|
File parentDir = destFile.getParentFile();
|
if (parentDir != null && !parentDir.exists()) {
|
parentDir.mkdirs();
|
}
|
// 复制文件
|
Files.copy(resource.getInputStream(), destFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
|
log.info("复制资源: {} 到 {}", relativePath, destFile.getAbsolutePath());
|
}
|
}
|
|
} catch (IOException e) {
|
log.error("复制PP_OCRv5目录到缓存失败: {}", e.getMessage());
|
e.printStackTrace();
|
}
|
}
|
|
|
/**
|
* 获取 resources 下模型文件的绝对路径
|
*
|
* @param relativePath 相对于 resources 的路径
|
*/
|
private String getModelPath(String relativePath) throws IOException {
|
String localPath = null;
|
try{
|
new ClassPathResource(relativePath).getFile().getAbsolutePath();
|
}catch (IOException e){
|
// 如果找不到文件，则尝试从本地缓存目录获取
|
localPath = getLocalPath(relativePath);
|
File modelFile = new File(localPath);
|
if (!modelFile.exists()) {
|
throw new IOException("模型文件不存在: " + localPath);
|
}
|
} finally {
|
log.info("OCR模型文件路径: {}", localPath);
|
}
|
return localPath;
|
}
|
|
/**
|
* 获取 resources 下模型文件的用户本地路径
|
*
|
* @param relativePath 相对于 resources 的路径
|
*/
|
private String getLocalPath(String relativePath) {
|
return new File(OcrSwingConstants.ocrDir,File.separator + relativePath).getAbsolutePath();
|
}
|
|
public OcrCommonRecModel getRecModel() throws IOException {
|
OcrRecModelConfig recModelConfig = new OcrRecModelConfig();
|
recModelConfig.setRecModelEnum(CommonRecModelEnum.PP_OCR_V5_MOBILE_REC_MODEL);
|
recModelConfig.setRecModelPath(
|
getModelPath("PP-OCRv5_server_rec_infer/PP-OCRv5_server_rec.onnx")
|
);
|
recModelConfig.setDevice(device);
|
recModelConfig.setTextDetModel(getDetectionModel());
|
return OcrModelFactory.getInstance().getRecModel(recModelConfig);
|
}
|
|
public OcrCommonDetModel getDetectionModel() throws IOException {
|
OcrDetModelConfig config = new OcrDetModelConfig();
|
config.setModelEnum(CommonDetModelEnum.PP_OCR_V5_MOBILE_DET_MODEL);
|
config.setDetModelPath(
|
getModelPath("PP-OCRv5_server_det_infer/PP-OCRv5_server_det.onnx")
|
);
|
config.setDevice(device);
|
return OcrModelFactory.getInstance().getDetModel(config);
|
}
|
|
public OcrDirectionModel getDirectionModel() throws IOException {
|
DirectionModelConfig directionModelConfig = new DirectionModelConfig();
|
directionModelConfig.setModelEnum(DirectionModelEnum.PP_LCNET_X0_25);
|
directionModelConfig.setModelPath(
|
getModelPath("PP-LCNet_x0_25_textline_ori_infer/PP-LCNet_x0_25_textline_ori_infer.onnx")
|
);
|
directionModelConfig.setDevice(device);
|
return OcrModelFactory.getInstance().getDirectionModel(directionModelConfig);
|
}
|
|
public OcrCommonRecModel getRecModelWithDirection() throws IOException {
|
OcrRecModelConfig recModelConfig = new OcrRecModelConfig();
|
recModelConfig.setRecModelEnum(CommonRecModelEnum.PP_OCR_V5_MOBILE_REC_MODEL);
|
recModelConfig.setRecModelPath(
|
getModelPath("PP-OCRv5_mobile_rec_infer/PP-OCRv5_mobile_rec_infer.onnx")
|
);
|
recModelConfig.setDevice(device);
|
recModelConfig.setTextDetModel(getDetectionModel());
|
recModelConfig.setDirectionModel(getDirectionModel());
|
return OcrModelFactory.getInstance().getRecModel(recModelConfig);
|
}
|
|
public String ocr(String url) {
|
String fullText = null;
|
try {
|
OcrRecOptions options = new OcrRecOptions(false, true);
|
OcrInfo ocrInfo = recModel.recognize(url, options);
|
log.info("OCR识别结果：{}", JSONObject.toJSONString(ocrInfo));
|
fullText = ocrInfo.getFullText();
|
} catch (Exception e) {
|
e.printStackTrace();
|
}
|
return fullText;
|
}
|
|
}
|
