Administrator / xlyAi

package com.xly.ocr.service;

import com.xly.ocr.util.OcrUtil;
import lombok.extern.slf4j.Slf4j;
import net.sourceforge.tess4j.Tesseract;
import net.sourceforge.tess4j.TesseractException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Arrays;
import java.util.List;

@Slf4j
@Service("ocrService")
public class OcrService {

    private static final Logger logger = LoggerFactory.getLogger(OcrService.class);

    @Value("${ocr.tmpPath}")
    private String tmpPath;


    private final Tesseract tesseract;

    // 配置参数
    private static final List<String> ALLOWED_EXTENSIONS = Arrays.asList(".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".gif");
    private static final long MAX_FILE_SIZE = 10 * 1024 * 1024; // 10MB
    private static final int BINARIZE_THRESHOLD = 127;
    private static final int MIN_WIDTH = 800;
    private static final int MIN_HEIGHT = 200;

    // 性能统计
    private static class OcrStats {
        long preprocessTime = 0;
        long ocrTime = 0;
        String imageSize = "";

        @Override
        public String toString() {
            return String.format("预处理耗时: %dms, OCR耗时: %dms, 图片尺寸: %s",
                    preprocessTime, ocrTime, imageSize);
        }
    }

    public OcrService(@Value("${tesseract.datapath}") String dataPath) {
        this.tesseract = new Tesseract();

        // 基础配置
        this.tesseract.setDatapath(dataPath);
        this.tesseract.setLanguage("chi_sim+eng");

        // 优化识别参数
        configureTesseract();

        logger.info("Tesseract 初始化完成，语言包路径: {}, 语言: chi_sim+eng", dataPath);
    }

    /**
     * 配置 Tesseract 参数
     */
    private void configureTesseract() {
        // 页面分割模式：3 = 自动页面分割，但没有方向检测
        this.tesseract.setPageSegMode(3);

        // OCR 引擎模式：3 = 默认，基于 LSTM 和传统引擎
        this.tesseract.setOcrEngineMode(3);

        // 提高中文识别率
        this.tesseract.setVariable("preserve_interword_spaces", "1");
        this.tesseract.setVariable("textord_force_make_prop_words", "true");

        // 可选：设置字符白名单（根据需要启用）
        // this.tesseract.setVariable("tessedit_char_whitelist",
        //     "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ，。！？；：\"‘’“”【】（）《》");

        // 可选：设置黑名单（排除干扰字符）
        // this.tesseract.setVariable("tessedit_char_blacklist", "|\\/`~@#$%^&*()_+={}[]");
    }

    /**
     * 图片预处理 - 优化的处理流程
     */
    private BufferedImage preprocessImage(BufferedImage originalImage) {
        if (originalImage == null) {
            return null;
        }

        try {
            long startTime = System.currentTimeMillis();

            // 1. 自动调整亮度和对比度
            BufferedImage adjusted = autoAdjustBrightnessContrast(originalImage);

            // 2. 灰度化
            BufferedImage grayImage = toGray(adjusted);

            // 3. 自适应二值化（比固定阈值更好）
            BufferedImage binaryImage = adaptiveBinarize(grayImage);

            // 4. 降噪处理
            BufferedImage denoisedImage = denoise(binaryImage);

            // 5. 放大图片（如果太小）
            BufferedImage scaledImage = scaleImageIfNeeded(denoisedImage);

            // 6. 可选：边缘增强（提高清晰度）
            BufferedImage enhancedImage = sharpen(scaledImage);

            long endTime = System.currentTimeMillis();
            logger.debug("图片预处理耗时: {}ms", endTime - startTime);

            return enhancedImage;

        } catch (Exception e) {
            logger.error("图片预处理失败: {}", e.getMessage(), e);
            return originalImage;
        }
    }

    /**
     * 自动调整亮度和对比度
     */
    private BufferedImage autoAdjustBrightnessContrast(BufferedImage image) {
        BufferedImage result = new BufferedImage(
                image.getWidth(), image.getHeight(), image.getType());

        // 计算亮度直方图
        int[] histogram = new int[256];
        for (int y = 0; y < image.getHeight(); y++) {
            for (int x = 0; x < image.getWidth(); x++) {
                int rgb = image.getRGB(x, y);
                int gray = (int)((rgb >> 16 & 0xFF) * 0.299 +
                        (rgb >> 8 & 0xFF) * 0.587 +
                        (rgb & 0xFF) * 0.114);
                histogram[gray]++;
            }
        }

        // 找到黑色和白色的阈值
        int total = image.getWidth() * image.getHeight();
        int blackThreshold = 0;
        int whiteThreshold = 255;

        int sum = 0;
        for (int i = 0; i < 256; i++) {
            sum += histogram[i];
            if (sum > total * 0.05) {
                blackThreshold = i;
                break;
            }
        }

        sum = 0;
        for (int i = 255; i >= 0; i--) {
            sum += histogram[i];
            if (sum > total * 0.05) {
                whiteThreshold = i;
                break;
            }
        }

        // 应用对比度拉伸
        for (int y = 0; y < image.getHeight(); y++) {
            for (int x = 0; x < image.getWidth(); x++) {
                int rgb = image.getRGB(x, y);
                int r = (rgb >> 16) & 0xFF;
                int g = (rgb >> 8) & 0xFF;
                int b = rgb & 0xFF;

                // 拉伸到 0-255 范围
                r = stretchValue(r, blackThreshold, whiteThreshold);
                g = stretchValue(g, blackThreshold, whiteThreshold);
                b = stretchValue(b, blackThreshold, whiteThreshold);

                result.setRGB(x, y, (r << 16) | (g << 8) | b);
            }
        }

        return result;
    }

    private int stretchValue(int value, int black, int white) {
        if (value <= black) return 0;
        if (value >= white) return 255;
        return (value - black) * 255 / (white - black);
    }

    /**
     * 灰度化
     */
    private BufferedImage toGray(BufferedImage image) {
        BufferedImage result = new BufferedImage(
                image.getWidth(), image.getHeight(), BufferedImage.TYPE_BYTE_GRAY);
        Graphics g = result.getGraphics();
        g.drawImage(image, 0, 0, null);
        g.dispose();
        return result;
    }

    /**
     * 自适应二值化 - 根据局部区域动态调整阈值
     */
    private BufferedImage adaptiveBinarize(BufferedImage image) {
        BufferedImage result = new BufferedImage(
                image.getWidth(), image.getHeight(), BufferedImage.TYPE_BYTE_BINARY);

        int blockSize = 15;
        int constant = 5;

        for (int y = 0; y < image.getHeight(); y++) {
            for (int x = 0; x < image.getWidth(); x++) {
                // 计算局部区域的平均值
                int sum = 0;
                int count = 0;
                for (int ky = -blockSize/2; ky <= blockSize/2; ky++) {
                    for (int kx = -blockSize/2; kx <= blockSize/2; kx++) {
                        int px = Math.min(Math.max(x + kx, 0), image.getWidth() - 1);
                        int py = Math.min(Math.max(y + ky, 0), image.getHeight() - 1);
                        sum += new Color(image.getRGB(px, py)).getRed();
                        count++;
                    }
                }
                int threshold = sum / count - constant;

                // 应用阈值
                int gray = new Color(image.getRGB(x, y)).getRed();
                int binary = gray > threshold ? 255 : 0;
                result.setRGB(x, y, new Color(binary, binary, binary).getRGB());
            }
        }
        return result;
    }

    /**
     * 降噪 - 优化的中值滤波
     */
    private BufferedImage denoise(BufferedImage image) {
        BufferedImage result = new BufferedImage(
                image.getWidth(), image.getHeight(), image.getType());

        for (int y = 1; y < image.getHeight() - 1; y++) {
            for (int x = 1; x < image.getWidth() - 1; x++) {
                int[] neighbors = new int[9];
                int index = 0;
                for (int ky = -1; ky <= 1; ky++) {
                    for (int kx = -1; kx <= 1; kx++) {
                        neighbors[index++] = new Color(image.getRGB(x + kx, y + ky)).getRed();
                    }
                }
                Arrays.sort(neighbors);
                int median = neighbors[4];
                result.setRGB(x, y, new Color(median, median, median).getRGB());
            }
        }

        // 处理边缘
        for (int x = 0; x < image.getWidth(); x++) {
            result.setRGB(x, 0, image.getRGB(x, 0));
            result.setRGB(x, image.getHeight() - 1, image.getRGB(x, image.getHeight() - 1));
        }
        for (int y = 0; y < image.getHeight(); y++) {
            result.setRGB(0, y, image.getRGB(0, y));
            result.setRGB(image.getWidth() - 1, y, image.getRGB(image.getWidth() - 1, y));
        }

        return result;
    }

    /**
     * 锐化处理 - 增强文字边缘
     */
    private BufferedImage sharpen(BufferedImage image) {
        BufferedImage result = new BufferedImage(
                image.getWidth(), image.getHeight(), image.getType());

        // 拉普拉斯锐化核
        float[] sharpenKernel = {
                0, -1, 0,
                -1, 5, -1,
                0, -1, 0
        };

        for (int y = 1; y < image.getHeight() - 1; y++) {
            for (int x = 1; x < image.getWidth() - 1; x++) {
                int sum = 0;
                int index = 0;
                for (int ky = -1; ky <= 1; ky++) {
                    for (int kx = -1; kx <= 1; kx++) {
                        int gray = new Color(image.getRGB(x + kx, y + ky)).getRed();
                        sum += gray * sharpenKernel[index++];
                    }
                }
                sum = Math.min(255, Math.max(0, sum));
                result.setRGB(x, y, new Color(sum, sum, sum).getRGB());
            }
        }

        return result;
    }

    /**
     * 放大图片（如果图片太小）
     */
    private BufferedImage scaleImageIfNeeded(BufferedImage image) {
        int width = image.getWidth();
        int height = image.getHeight();

        if (width >= MIN_WIDTH && height >= MIN_HEIGHT) {
            return image;
        }

        double scaleX = (double) MIN_WIDTH / width;
        double scaleY = (double) MIN_HEIGHT / height;
        double scale = Math.max(scaleX, scaleY);

        int newWidth = (int) (width * scale);
        int newHeight = (int) (height * scale);

        // 使用更好的插值算法
        BufferedImage result = new BufferedImage(newWidth, newHeight, image.getType());
        Graphics2D g2d = result.createGraphics();
        g2d.setRenderingHint(RenderingHints.KEY_INTERPOLATION,
                RenderingHints.VALUE_INTERPOLATION_BICUBIC);
        g2d.setRenderingHint(RenderingHints.KEY_RENDERING,
                RenderingHints.VALUE_RENDER_QUALITY);
        g2d.setRenderingHint(RenderingHints.KEY_ANTIALIASING,
                RenderingHints.VALUE_ANTIALIAS_ON);
        g2d.drawImage(image, 0, 0, newWidth, newHeight, null);
        g2d.dispose();

        logger.debug("图片已放大: {}x{} -> {}x{}", width, height, newWidth, newHeight);
        return result;
    }

    /**
     * 识别图片中的文字（增强版）
     */
    public String extractText(File imageFile) {
        if (imageFile == null || !imageFile.exists()) {
            logger.error("图片文件不存在或为空");
            return "图片文件不存在";
        }

        OcrStats stats = new OcrStats();

        try {
            logger.info("开始识别图片: {}, 大小: {} bytes",
                    imageFile.getAbsolutePath(), imageFile.length());

            // 读取原始图片
            long readStart = System.currentTimeMillis();
            BufferedImage originalImage = ImageIO.read(imageFile);
            if (originalImage == null) {
                return "无法读取图片文件，请确保图片格式正确";
            }
            stats.imageSize = originalImage.getWidth() + "x" + originalImage.getHeight();

            // 图片预处理
            long preprocessStart = System.currentTimeMillis();
            BufferedImage processedImage = preprocessImage(originalImage);
            stats.preprocessTime = System.currentTimeMillis() - preprocessStart;

            // 可选：保存预处理图片用于调试（生产环境可注释）
            if (logger.isDebugEnabled()) {
                saveDebugImage(processedImage, imageFile);
            }

            // 执行 OCR
            long ocrStart = System.currentTimeMillis();
            String result = tesseract.doOCR(processedImage);
            stats.ocrTime = System.currentTimeMillis() - ocrStart;

            logger.info("识别完成 - {}", stats);

            // 清理识别结果
            result = cleanResult(result);

            if (result.isEmpty()) {
                logger.warn("识别结果为空，可能需要调整预处理参数");
            }

            return result;

        } catch (TesseractException e) {
            logger.error("OCR识别失败: {}", e.getMessage(), e);
            return "OCR识别失败: " + e.getMessage();
        } catch (IOException e) {
            logger.error("读取图片失败: {}", e.getMessage(), e);
            return "读取图片失败: " + e.getMessage();
        }
    }

    /**
     * 保存调试图片（仅用于调试）
     */
    private void saveDebugImage(BufferedImage image, File originalFile) {
        try {
            String debugPath = originalFile.getParent() + "/debug_" + originalFile.getName();
            File debugFile = new File(debugPath);
            ImageIO.write(image, "png", debugFile);
            logger.debug("预处理图片已保存: {}", debugPath);
        } catch (IOException e) {
            logger.debug("保存调试图片失败: {}", e.getMessage());
        }
    }

    /**
     * 清理识别结果
     */
    private String cleanResult(String result) {
        if (result == null || result.isEmpty()) {
            return "";
        }

        // 去除首尾空白
        result = result.trim();

        // 规范化换行符
        result = result.replaceAll("\\r\\n", "\n")
                .replaceAll("\\r", "\n");

        // 合并多个空行
        result = result.replaceAll("\n{3,}", "\n\n");

        // 去除行首行尾空格
        String[] lines = result.split("\n");
        StringBuilder cleaned = new StringBuilder();
        for (String line : lines) {
            cleaned.append(line.trim()).append("\n");
        }

        return cleaned.toString().trim();
    }

    /**
     * 封装方法，接收上传的 MultipartFile
     */
    public String extractTextFromMultipartFile(MultipartFile file) {
        if (file == null || file.isEmpty()) {
            logger.warn("上传的文件为空");
            return "上传的文件为空";
        }

        // 验证文件大小
        if (file.getSize() > MAX_FILE_SIZE) {
            logger.warn("文件过大: {} bytes, 超过限制: {} bytes",
                    file.getSize(), MAX_FILE_SIZE);
            return String.format("文件过大，最大支持 %dMB", MAX_FILE_SIZE / 1024 / 1024);
        }

        // 验证文件格式
        String originalFilename = file.getOriginalFilename();
        if (originalFilename != null && !isAllowedImage(originalFilename)) {
            logger.warn("不支持的文件格式: {}", originalFilename);
            return "不支持的文件格式，仅支持: " + String.join(", ", ALLOWED_EXTENSIONS);
        }
        String sText = OcrUtil.ocrFile(file,tmpPath);
        return sText;
    }

    /**
     * 清理临时文件
     */
    private void cleanupTempFile(Path tempFile) {
        if (tempFile != null) {
            try {
                Files.deleteIfExists(tempFile);
                logger.debug("临时文件已删除: {}", tempFile);
            } catch (IOException e) {
                logger.warn("删除临时文件失败: {}", tempFile, e);
                // 注册JVM退出时删除
                tempFile.toFile().deleteOnExit();
            }
        }
    }

    /**
     * 批量识别（用于多张图片）
     */
    public List<String> batchExtractText(List<MultipartFile> files) {
        return files.stream()
                .map(this::extractTextFromMultipartFile)
                .collect(java.util.stream.Collectors.toList());
    }

    /**
     * 检查文件扩展名是否允许
     */
    private boolean isAllowedImage(String filename) {
        if (filename == null) {
            return false;
        }
        String lowerFilename = filename.toLowerCase();
        return ALLOWED_EXTENSIONS.stream()
                .anyMatch(lowerFilename::endsWith);
    }

    /**
     * 获取文件扩展名
     */
    private String getFileExtension(String filename) {
        if (filename == null || !filename.contains(".")) {
            return ".jpg";
        }
        return filename.substring(filename.lastIndexOf("."));
    }
}