433 行
16 KiB
JavaScript
433 行
16 KiB
JavaScript
// 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');
|
||
this.preprocessedDir = path.join(process.cwd(), 'temp', 'preprocessed');
|
||
this.logger = {
|
||
info: (msg, ...args) => console.log(`🔤 [识别] ${msg}`, ...args),
|
||
error: (msg, ...args) => console.error(`❌ [识别] ${msg}`, ...args),
|
||
debug: (msg, ...args) => console.log(`🐛 [识别] ${msg}`, ...args),
|
||
warn: (msg, ...args) => console.warn(`🐛 [识别] ${msg}`, ...args)
|
||
};
|
||
|
||
// 确保目录存在
|
||
fse.ensureDirSync(this.debugDir);
|
||
fse.ensureDirSync(this.preprocessedDir);
|
||
}
|
||
|
||
initialize(recSession, config) {
|
||
this.recSession = recSession;
|
||
this.config = config;
|
||
this.logger.info('文本识别器初始化完成');
|
||
}
|
||
|
||
async loadCharacterSet(keysPath) {
|
||
try {
|
||
const keysContent = await fse.readFile(keysPath, 'utf8');
|
||
this.characterSet = [];
|
||
const lines = keysContent.split('\n');
|
||
|
||
// 使用提供的字符集文件
|
||
const uniqueChars = new Set();
|
||
|
||
for (const line of lines) {
|
||
const trimmed = line.trim();
|
||
// 跳过空行和注释行
|
||
if (trimmed && !trimmed.startsWith('#')) {
|
||
// 将每行作为一个完整的字符处理
|
||
uniqueChars.add(trimmed);
|
||
}
|
||
}
|
||
|
||
this.characterSet = Array.from(uniqueChars);
|
||
|
||
if (this.characterSet.length === 0) {
|
||
throw new Error('字符集文件为空或格式不正确');
|
||
}
|
||
|
||
this.logger.info(`字符集加载完成: ${this.characterSet.length}个字符`);
|
||
|
||
// 记录字符集统计信息
|
||
const charTypes = {
|
||
chinese: 0,
|
||
english: 0,
|
||
digit: 0,
|
||
punctuation: 0,
|
||
other: 0
|
||
};
|
||
|
||
this.characterSet.forEach(char => {
|
||
if (/[\u4e00-\u9fff]/.test(char)) {
|
||
charTypes.chinese++;
|
||
} else if (/[a-zA-Z]/.test(char)) {
|
||
charTypes.english++;
|
||
} else if (/[0-9]/.test(char)) {
|
||
charTypes.digit++;
|
||
} else if (/[，。！？；：""（）【】《》…—·]/.test(char)) {
|
||
charTypes.punctuation++;
|
||
} else {
|
||
charTypes.other++;
|
||
}
|
||
});
|
||
|
||
this.logger.debug(`字符集统计: 中文${charTypes.chinese}, 英文${charTypes.english}, 数字${charTypes.digit}, 标点${charTypes.punctuation}, 其他${charTypes.other}`);
|
||
this.logger.debug(`前20个字符: ${this.characterSet.slice(0, 20).join('')}`);
|
||
|
||
} catch (error) {
|
||
this.logger.error('加载字符集失败', error.message);
|
||
// 完全使用提供的字符集，失败时抛出错误
|
||
throw new Error(`字符集加载失败: ${error.message}`);
|
||
}
|
||
}
|
||
|
||
getCharacterSetSize() {
|
||
return this.characterSet.length;
|
||
}
|
||
|
||
async recognizeText(textRegionBuffer, regionIndex = 0) {
|
||
const startTime = Date.now();
|
||
this.logger.info(`开始文本识别 - 区域 ${regionIndex}`);
|
||
|
||
try {
|
||
const inputTensor = await this.prepareRecognitionInput(textRegionBuffer, regionIndex);
|
||
const outputs = await this.recSession.run({ [this.recSession.inputNames[0]]: inputTensor });
|
||
const result = this.postprocessRecognition(outputs);
|
||
|
||
const processingTime = Date.now() - startTime;
|
||
this.logger.info(`识别完成 - 区域 ${regionIndex}: "${result.text}", 置信度: ${result.confidence.toFixed(4)}, 耗时: ${processingTime}ms`);
|
||
|
||
return result;
|
||
|
||
} catch (error) {
|
||
this.logger.error(`文本识别失败 - 区域 ${regionIndex}`, error);
|
||
return { text: '', confidence: 0 };
|
||
}
|
||
}
|
||
|
||
async prepareRecognitionInput(textRegionBuffer, regionIndex = 0) {
|
||
this.logger.debug(`准备识别输入 - 区域 ${regionIndex}`);
|
||
|
||
const targetHeight = 48;
|
||
const targetWidth = 320; // 原始目标宽度
|
||
const finalWidth = targetWidth + 20; // 最终宽度（左右各加10像素）
|
||
const timestamp = Date.now();
|
||
|
||
try {
|
||
const metadata = await sharp(textRegionBuffer).metadata();
|
||
this.logger.debug(`原始区域 ${regionIndex}: ${metadata.width}x${metadata.height}`);
|
||
|
||
// 保存原始裁剪区域图像
|
||
const originalPath = path.join(this.preprocessedDir, `region-${regionIndex}-original-${timestamp}.png`);
|
||
await fse.writeFile(originalPath, textRegionBuffer);
|
||
this.logger.debug(`保存原始区域图像: ${originalPath}`);
|
||
|
||
// 图像分析
|
||
const stats = await sharp(textRegionBuffer).grayscale().stats();
|
||
const meanBrightness = stats.channels[0].mean;
|
||
const stdDev = stats.channels[0].stdev;
|
||
|
||
this.logger.debug(`图像统计 - 区域 ${regionIndex}: 亮度=${meanBrightness.toFixed(1)}, 对比度=${stdDev.toFixed(1)}`);
|
||
|
||
// 智能预处理
|
||
let processedBuffer = await this.applySmartPreprocessing(textRegionBuffer, meanBrightness, stdDev, regionIndex);
|
||
|
||
// 保存预处理后的图像（灰度+对比度调整后）
|
||
const processedPath = path.join(this.preprocessedDir, `region-${regionIndex}-processed-${timestamp}.png`);
|
||
await fse.writeFile(processedPath, processedBuffer);
|
||
this.logger.debug(`保存预处理图像: ${processedPath}`);
|
||
|
||
// 保持宽高比的resize，并在左右添加10像素空白
|
||
const resizedBuffer = await this.resizeWithAspectRatio(processedBuffer, targetWidth, targetHeight, regionIndex);
|
||
|
||
// 保存调整大小后的图像
|
||
const resizedPath = path.join(this.preprocessedDir, `region-${regionIndex}-resized-${timestamp}.png`);
|
||
await fse.writeFile(resizedPath, resizedBuffer);
|
||
this.logger.debug(`保存调整大小图像: ${resizedPath}`);
|
||
|
||
// 使用最终尺寸创建张量
|
||
const inputData = await this.bufferToTensor(resizedBuffer, finalWidth, targetHeight);
|
||
this.logger.debug(`识别输入张量准备完成 - 区域 ${regionIndex}`);
|
||
|
||
// 创建张量时使用最终尺寸
|
||
return new Tensor('float32', inputData, [1, 3, targetHeight, finalWidth]);
|
||
|
||
} catch (error) {
|
||
this.logger.error(`准备识别输入失败 - 区域 ${regionIndex}`, error);
|
||
return new Tensor('float32', new Float32Array(3 * targetHeight * finalWidth).fill(0.5), [1, 3, targetHeight, finalWidth]);
|
||
}
|
||
}
|
||
|
||
async applySmartPreprocessing(buffer, meanBrightness, stdDev, regionIndex = 0) {
|
||
let processedBuffer = buffer;
|
||
|
||
if (meanBrightness > 200 && stdDev < 30) {
|
||
this.logger.debug(`区域 ${regionIndex}: 应用高亮度图像增强`);
|
||
processedBuffer = await sharp(buffer)
|
||
.linear(1.5, -50)
|
||
.normalize()
|
||
.grayscale()
|
||
.toBuffer();
|
||
} else if (meanBrightness < 80) {
|
||
this.logger.debug(`区域 ${regionIndex}: 应用低亮度图像增强`);
|
||
processedBuffer = await sharp(buffer)
|
||
.linear(1.2, 30)
|
||
.normalize()
|
||
.grayscale()
|
||
.toBuffer();
|
||
} else if (stdDev < 20) {
|
||
this.logger.debug(`区域 ${regionIndex}: 应用低对比度增强`);
|
||
processedBuffer = await sharp(buffer)
|
||
.linear(1.3, -20)
|
||
.normalize()
|
||
.grayscale()
|
||
.toBuffer();
|
||
} else {
|
||
this.logger.debug(`区域 ${regionIndex}: 应用标准化灰度处理`);
|
||
processedBuffer = await sharp(buffer)
|
||
.normalize()
|
||
.grayscale()
|
||
.toBuffer();
|
||
}
|
||
|
||
return processedBuffer;
|
||
}
|
||
|
||
async resizeWithAspectRatio(buffer, targetWidth, targetHeight, regionIndex = 0) {
|
||
const metadata = await sharp(buffer).metadata();
|
||
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));
|
||
|
||
this.logger.debug(`区域 ${regionIndex}: 调整尺寸 ${metadata.width}x${metadata.height} -> ${resizeWidth}x${resizeHeight}`);
|
||
|
||
// 计算居中的偏移量
|
||
const offsetX = Math.floor((targetWidth - resizeWidth) / 2);
|
||
const offsetY = Math.floor((targetHeight - resizeHeight) / 2);
|
||
|
||
this.logger.debug(`区域 ${regionIndex}: 居中偏移 X=${offsetX}, Y=${offsetY}`);
|
||
|
||
// 先调整大小并居中
|
||
let resizedBuffer = await sharp(buffer)
|
||
.resize(resizeWidth, resizeHeight, {
|
||
fit: 'contain',
|
||
background: { r: 255, g: 255, b: 255 }
|
||
})
|
||
.extend({
|
||
top: offsetY,
|
||
bottom: targetHeight - resizeHeight - offsetY,
|
||
left: offsetX,
|
||
right: targetWidth - resizeWidth - offsetX,
|
||
background: { r: 255, g: 255, b: 255 }
|
||
})
|
||
.png()
|
||
.toBuffer();
|
||
|
||
// 在左右各添加10像素空白
|
||
const finalWidth = targetWidth + 20; // 左右各加10像素
|
||
const finalHeight = targetHeight;
|
||
|
||
resizedBuffer = await sharp(resizedBuffer)
|
||
.extend({
|
||
top: 0,
|
||
bottom: 0,
|
||
left: 10,
|
||
right: 10,
|
||
background: { r: 255, g: 255, b: 255 }
|
||
})
|
||
.png()
|
||
.toBuffer();
|
||
|
||
this.logger.debug(`区域 ${regionIndex}: 最终尺寸 ${finalWidth}x${finalHeight} (左右各加10像素空白)`);
|
||
|
||
return resizedBuffer;
|
||
}
|
||
|
||
async bufferToTensor(buffer, width, height) {
|
||
// 获取实际图像尺寸（因为现在宽度增加了20像素）
|
||
const metadata = await sharp(buffer).metadata();
|
||
const actualWidth = metadata.width;
|
||
const actualHeight = metadata.height;
|
||
|
||
const imageData = await sharp(buffer)
|
||
.ensureAlpha()
|
||
.raw()
|
||
.toBuffer({ resolveWithObject: true });
|
||
|
||
// 使用实际尺寸创建张量
|
||
const inputData = new Float32Array(3 * actualHeight * actualWidth);
|
||
const data = imageData.data;
|
||
|
||
for (let i = 0; i < data.length; i += 4) {
|
||
const pixelIndex = Math.floor(i / 4);
|
||
const y = Math.floor(pixelIndex / actualWidth);
|
||
const x = pixelIndex % actualWidth;
|
||
|
||
// 使用灰度值填充三个通道
|
||
const grayValue = data[i] / 255.0;
|
||
|
||
for (let c = 0; c < 3; c++) {
|
||
const inputIndex = c * actualHeight * actualWidth + y * actualWidth + x;
|
||
if (inputIndex < inputData.length) {
|
||
inputData[inputIndex] = grayValue;
|
||
}
|
||
}
|
||
}
|
||
|
||
return inputData;
|
||
}
|
||
|
||
postprocessRecognition(outputs) {
|
||
this.logger.debug('开始识别后处理');
|
||
|
||
try {
|
||
const outputNames = this.recSession.outputNames;
|
||
const recognitionOutput = outputs[outputNames[0]];
|
||
|
||
if (!recognitionOutput) {
|
||
this.logger.debug('识别输出为空');
|
||
return { text: '', confidence: 0 };
|
||
}
|
||
|
||
const data = recognitionOutput.data;
|
||
const [batch, seqLen, vocabSize] = recognitionOutput.dims;
|
||
|
||
this.logger.debug(`序列长度: ${seqLen}, 词汇表大小: ${vocabSize}, 字符集大小: ${this.characterSet.length}`);
|
||
|
||
if (this.characterSet.length === 0) {
|
||
this.logger.error('字符集为空');
|
||
return { text: '', confidence: 0 };
|
||
}
|
||
|
||
// 验证词汇表大小与字符集大小的匹配
|
||
if (vocabSize !== this.characterSet.length + 1) {
|
||
this.logger.warn(`词汇表大小(${vocabSize})与字符集大小(${this.characterSet.length})不匹配，可能影响识别效果`);
|
||
}
|
||
|
||
const { text, confidence } = this.ctcDecode(data, seqLen, vocabSize);
|
||
this.logger.debug(`解码结果: "${text}", 置信度: ${confidence.toFixed(4)}`);
|
||
|
||
return { text, confidence };
|
||
|
||
} catch (error) {
|
||
this.logger.error('识别后处理失败', error);
|
||
return { text: '', confidence: 0 };
|
||
}
|
||
}
|
||
|
||
ctcDecode(data, seqLen, vocabSize) {
|
||
let text = '';
|
||
let lastCharIndex = -1;
|
||
let confidenceSum = 0;
|
||
let charCount = 0;
|
||
|
||
// 动态阈值调整
|
||
const baseThreshold = 0.03;
|
||
let confidenceThreshold = baseThreshold;
|
||
|
||
// 先分析整个序列的置信度分布
|
||
let maxSequenceProb = 0;
|
||
for (let t = 0; t < seqLen; t++) {
|
||
for (let i = 0; i < vocabSize; i++) {
|
||
maxSequenceProb = Math.max(maxSequenceProb, data[t * vocabSize + i]);
|
||
}
|
||
}
|
||
|
||
// 如果整体置信度较低，降低阈值
|
||
if (maxSequenceProb < 0.5) {
|
||
confidenceThreshold = baseThreshold * 0.5;
|
||
}
|
||
|
||
this.logger.debug(`使用解码阈值: ${confidenceThreshold.toFixed(4)}`);
|
||
|
||
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;
|
||
}
|
||
}
|
||
|
||
// 改进的CTC解码逻辑
|
||
if (maxIndex > 0 && maxProb > confidenceThreshold) {
|
||
const charIndex = maxIndex - 1;
|
||
if (charIndex < this.characterSet.length) {
|
||
const char = this.characterSet[charIndex];
|
||
|
||
// 更智能的重复字符处理
|
||
const shouldAddChar = maxIndex !== lastCharIndex ||
|
||
maxProb > 0.8 ||
|
||
(maxIndex === lastCharIndex && charCount > 0 && text[text.length - 1] !== char);
|
||
|
||
if (shouldAddChar && char && char.trim() !== '') {
|
||
text += char;
|
||
confidenceSum += maxProb;
|
||
charCount++;
|
||
}
|
||
lastCharIndex = maxIndex;
|
||
} else {
|
||
this.logger.warn(`字符索引${charIndex}超出字符集范围(0-${this.characterSet.length-1})`);
|
||
}
|
||
} else if (maxIndex === 0) {
|
||
lastCharIndex = -1;
|
||
}
|
||
}
|
||
|
||
const avgConfidence = charCount > 0 ? confidenceSum / charCount : 0;
|
||
|
||
// 基本的文本清理（不包含错误模式修复）
|
||
const cleanedText = this.basicTextCleaning(text);
|
||
|
||
return {
|
||
text: cleanedText,
|
||
confidence: avgConfidence
|
||
};
|
||
}
|
||
|
||
basicTextCleaning(text) {
|
||
if (!text) return '';
|
||
|
||
let cleaned = text;
|
||
|
||
// 1. 移除过多的重复字符（保留合理的重复）
|
||
cleaned = cleaned.replace(/([^0-9])\1{2,}/g, '$1$1');
|
||
|
||
// 2. 修复标点符号
|
||
cleaned = cleaned.replace(/∶/g, '：')
|
||
.replace(/《/g, '（')
|
||
.replace(/》/g, '）');
|
||
|
||
// 3. 修复数字和百分号
|
||
cleaned = cleaned.replace(/(\d+)%%/g, '$1%');
|
||
|
||
return cleaned.trim();
|
||
}
|
||
}
|
||
|
||
export default TextRecognizer; |
