Electron-vue3-ts-offline/server/utils/textRecognizer.js

// server/utils/textRecognizer.js
import { Tensor } from 'onnxruntime-node';
import sharp from 'sharp';
import fse from 'fs-extra';
import * as path from 'path';

class TextRecognizer {
    constructor() {
        this.recSession = null;
        this.config = null;
        this.characterSet = [];
        this.debugDir = path.join(process.cwd(), 'temp', 'debug');
        fse.ensureDirSync(this.debugDir);
    }

    initialize(recSession, config) {
        this.recSession = recSession;
        this.config = config;
    }

    async loadCharacterSet(keysPath) {
        try {
            const keysContent = await fse.readFile(keysPath, 'utf8');
            this.characterSet = [];
            const lines = keysContent.split('\n');

            for (const line of lines) {
                const trimmed = line.trim();
                if (trimmed && !trimmed.startsWith('#')) {
                    for (const char of trimmed) {
                        if (char.trim() && !this.characterSet.includes(char)) {
                            this.characterSet.push(char);
                        }
                    }
                }
            }

            if (this.characterSet.length === 0) {
                throw new Error('字符集文件为空或格式不正确');
            }

            console.log(`✅ 字符集加载完成，共 ${this.characterSet.length} 个字符`);

        } catch (error) {
            console.error('❌ 加载字符集失败，使用默认字符集:', error.message);
            this.characterSet = this.getDefaultCharacterSet();
        }
    }

    getDefaultCharacterSet() {
        const defaultSet = [];
        for (let i = 0; i <= 9; i++) defaultSet.push(i.toString());
        for (let i = 97; i <= 122; i++) defaultSet.push(String.fromCharCode(i));
        for (let i = 65; i <= 90; i++) defaultSet.push(String.fromCharCode(i));
        defaultSet.push(...' ，。！？；：""（）【】《》…—·'.split(''));

        const commonChinese = '的一是不了在人有的我他这个们中来就时大地为子中你说道生国年着就那和要她出也得里后自以会家可下而过天去能对小多然于心学么之都好看起发当没成只如事把还用第样道想作种开美总从无情已面最女但现前些所同日手又行意动方期它头经长儿回位分爱老因很给名法间斯知世什两次使身者被高已亲其进此话常与活正感';
        for (const char of commonChinese) {
            defaultSet.push(char);
        }

        console.log(`📝 使用默认字符集，共 ${defaultSet.length} 个字符`);
        return defaultSet;
    }

    getCharacterSetSize() {
        return this.characterSet.length;
    }

    async recognizeText(textRegionBuffer) {
        console.log('🔠 === 开始文本识别流程 ===');

        try {
            console.log('📥 1. 准备识别输入...');
            console.log(`   - 输入图像大小: ${textRegionBuffer.length} 字节`);

            const inputTensor = await this.prepareRecognitionInput(textRegionBuffer);
            console.log('✅ 输入张量准备完成');
            console.log(`   - 张量形状: [${inputTensor.dims.join(', ')}]`);
            console.log(`   - 张量类型: ${inputTensor.type}`);
            console.log(`   - 数据长度: ${inputTensor.data.length}`);

            // 数据验证
            const tensorData = inputTensor.data;
            let minVal = Infinity;
            let maxVal = -Infinity;
            let sumVal = 0;
            let validCount = 0;

            for (let i = 0; i < Math.min(100, tensorData.length); i++) {
                const val = tensorData[i];
                if (!isNaN(val) && isFinite(val)) {
                    minVal = Math.min(minVal, val);
                    maxVal = Math.max(maxVal, val);
                    sumVal += val;
                    validCount++;
                }
            }

            console.log(`   - 数据范围: ${minVal.toFixed(4)} ~ ${maxVal.toFixed(4)}`);
            console.log(`   - 数据均值: ${(sumVal / validCount).toFixed(4)}`);

            console.log('🧠 2. 执行模型推理...');
            const startInference = Date.now();
            const outputs = await this.recSession.run({ [this.recSession.inputNames[0]]: inputTensor });
            const inferenceTime = Date.now() - startInference;
            console.log(`✅ 模型推理完成 (${inferenceTime}ms)`);

            const outputNames = this.recSession.outputNames;
            console.log(`   - 输出数量: ${outputNames.length}`);

            outputNames.forEach((name, index) => {
                const output = outputs[name];
                if (output) {
                    console.log(`   - 输出 ${index + 1} (${name}): 形状 [${output.dims.join(', ')}]`);
                }
            });

            console.log('🔍 3. 后处理识别结果...');
            const result = this.postprocessRecognition(outputs);
            console.log('✅ 后处理完成');
            console.log(`   - 识别文本: "${result.text}"`);
            console.log(`   - 置信度: ${result.confidence.toFixed(4)}`);
            console.log(`   - 文本长度: ${result.text.length} 字符`);

            console.log('🎉 === 文本识别流程完成 ===');
            return result;

        } catch (error) {
            console.error('❌ 文本识别失败:');
            console.error(`   - 错误信息: ${error.message}`);
            return { text: '', confidence: 0 };
        }
    }

    async prepareRecognitionInput(textRegionBuffer) {
        console.log('   📝 准备识别输入详情:');

        try {
            const targetHeight = 48;
            const targetWidth = 320;

            const metadata = await sharp(textRegionBuffer).metadata();
            console.log(`     - 原始图像尺寸: ${metadata.width}x${metadata.height}`);

            // 保存原始图像用于调试
            const originalPath = path.join(this.debugDir, `original-${Date.now()}.png`);
            await fse.writeFile(originalPath, textRegionBuffer);

            // 关键修复：正确的预处理流程
            let processedBuffer = textRegionBuffer;

            // 1. 分析图像特性
            const stats = await sharp(processedBuffer)
                .grayscale()
                .stats();
            const meanBrightness = stats.channels[0].mean;
            const stdDev = stats.channels[0].stdev;

            console.log(`     - 图像统计: 均值=${meanBrightness.toFixed(1)}, 标准差=${stdDev.toFixed(1)}`);

            // 2. 改进的预处理策略
            if (meanBrightness > 200 && stdDev < 30) {
                console.log('     - 检测到高亮度图像，进行对比度增强');
                processedBuffer = await sharp(processedBuffer)
                    .linear(1.5, -50)
                    .normalize()
                    .grayscale()
                    .toBuffer();
            } else if (meanBrightness < 80) {
                console.log('     - 检测到低亮度图像，进行亮度调整');
                processedBuffer = await sharp(processedBuffer)
                    .linear(1.2, 30)
                    .normalize()
                    .grayscale()
                    .toBuffer();
            } else {
                console.log('     - 使用标准化灰度处理');
                processedBuffer = await sharp(processedBuffer)
                    .normalize()
                    .grayscale()
                    .toBuffer();
            }

            // 3. 保持宽高比的resize
            const originalAspectRatio = metadata.width / metadata.height;
            const targetAspectRatio = targetWidth / targetHeight;

            let resizeWidth, resizeHeight;

            if (originalAspectRatio > targetAspectRatio) {
                // 宽度限制
                resizeWidth = targetWidth;
                resizeHeight = Math.round(targetWidth / originalAspectRatio);
            } else {
                // 高度限制
                resizeHeight = targetHeight;
                resizeWidth = Math.round(targetHeight * originalAspectRatio);
            }

            // 确保尺寸有效
            resizeWidth = Math.max(1, Math.min(resizeWidth, targetWidth));
            resizeHeight = Math.max(1, Math.min(resizeHeight, targetHeight));

            processedBuffer = await sharp(processedBuffer)
                .resize(resizeWidth, resizeHeight, {
                    fit: 'contain',
                    background: { r: 255, g: 255, b: 255 }
                })
                .extend({
                    top: 0,
                    bottom: targetHeight - resizeHeight,
                    left: 0,
                    right: targetWidth - resizeWidth,
                    background: { r: 255, g: 255, b: 255 }
                })
                .png()
                .toBuffer();

            const processedMetadata = await sharp(processedBuffer).metadata();
            console.log(`     - 处理后尺寸: ${processedMetadata.width}x${processedMetadata.height}`);

            // 保存预处理后的图像用于调试
            const processedPath = path.join(this.debugDir, `processed-${Date.now()}.png`);
            await fse.writeFile(processedPath, processedBuffer);

            // 4. 转换为张量 - 关键修复：正确的归一化
            console.log('     - 转换为张量数据...');
            const imageData = await sharp(processedBuffer)
                .ensureAlpha()
                .raw()
                .toBuffer({ resolveWithObject: true });

            const inputData = new Float32Array(3 * targetHeight * targetWidth);
            const data = imageData.data;
            const channels = imageData.info.channels;

            // 使用正确的归一化方法
            for (let i = 0; i < data.length; i += channels) {
                const pixelIndex = Math.floor(i / channels);
                const y = Math.floor(pixelIndex / targetWidth);
                const x = pixelIndex % targetWidth;

                // 对每个位置，三个通道使用相同的灰度值
                const grayValue = data[i] / 255.0;

                for (let c = 0; c < 3; c++) {
                    const inputIndex = c * targetHeight * targetWidth + y * targetWidth + x;
                    if (inputIndex < inputData.length) {
                        inputData[inputIndex] = grayValue;
                    }
                }
            }

            console.log(`     - 输入数据长度: ${inputData.length}`);

            // 数据验证
            let validCount = 0;
            let sumValue = 0;
            let minValue = Infinity;
            let maxValue = -Infinity;

            for (let i = 0; i < Math.min(100, inputData.length); i++) {
                const val = inputData[i];
                if (!isNaN(val) && isFinite(val)) {
                    validCount++;
                    sumValue += val;
                    minValue = Math.min(minValue, val);
                    maxValue = Math.max(maxValue, val);
                }
            }

            console.log(`     - 数据验证: 有效=${validCount}`);
            console.log(`     - 数据范围: ${minValue.toFixed(4)} ~ ${maxValue.toFixed(4)}`);
            console.log(`     - 数据均值: ${(sumValue / validCount).toFixed(4)}`);

            return new Tensor('float32', inputData, [1, 3, targetHeight, targetWidth]);

        } catch (error) {
            console.error(`     ❌ 准备输入失败: ${error.message}`);
            // 返回有效的默认张量
            return new Tensor('float32', new Float32Array(3 * 48 * 320).fill(0.5), [1, 3, 48, 320]);
        }
    }

    postprocessRecognition(outputs) {
        console.log('   📝 后处理识别结果详情:');

        try {
            const outputNames = this.recSession.outputNames;
            const recognitionOutput = outputs[outputNames[0]];

            if (!recognitionOutput) {
                console.log('     ❌ 识别输出为空');
                return { text: '', confidence: 0 };
            }

            const data = recognitionOutput.data;
            const [batch, seqLen, vocabSize] = recognitionOutput.dims;

            console.log(`     - 序列长度: ${seqLen}, 词汇表大小: ${vocabSize}`);
            console.log(`     - 输出数据总数: ${data.length}`);
            console.log(`     - 字符集大小: ${this.characterSet.length}`);

            if (this.characterSet.length === 0) {
                console.log('     ❌ 字符集为空');
                return { text: '', confidence: 0 };
            }

            // 改进的CTC解码算法
            let text = '';
            let lastCharIndex = -1;
            let confidenceSum = 0;
            let charCount = 0;

            // 降低置信度阈值，提高召回率
            const confidenceThreshold = 0.05;

            console.log('     - 处理每个时间步:');
            for (let t = 0; t < seqLen; t++) {
                let maxProb = -1;
                let maxIndex = -1;

                // 找到当前时间步的最大概率字符
                for (let i = 0; i < vocabSize; i++) {
                    const prob = data[t * vocabSize + i];
                    if (prob > maxProb) {
                        maxProb = prob;
                        maxIndex = i;
                    }
                }

                // 改进的解码逻辑
                if (maxIndex > 0 && maxProb > confidenceThreshold) {
                    const char = this.characterSet[maxIndex - 1] || '';

                    // 放宽重复字符限制
                    if (maxIndex !== lastCharIndex || maxProb > 0.8) {
                        if (char && char.trim() !== '') {
                            text += char;
                            confidenceSum += maxProb;
                            charCount++;
                            console.log(`       [位置 ${t}] 字符: "${char}", 置信度: ${maxProb.toFixed(4)}`);
                        }
                        lastCharIndex = maxIndex;
                    }
                } else if (maxIndex === 0) {
                    // 空白符，重置lastCharIndex
                    lastCharIndex = -1;
                }
            }

            const avgConfidence = charCount > 0 ? confidenceSum / charCount : 0;

            console.log(`     - 识别结果: "${text}"`);
            console.log(`     - 字符数: ${charCount}, 平均置信度: ${avgConfidence.toFixed(4)}`);

            return {
                text: text,
                confidence: avgConfidence
            };

        } catch (error) {
            console.error(`     ❌ 后处理失败: ${error.message}`);
            return { text: '', confidence: 0 };
        }
    }
}

export default TextRecognizer;