440 lines
18 KiB
JavaScript
440 lines
18 KiB
JavaScript
// audioClassifier.js (核心功能模块)
|
||
|
||
// TensorFlow.js 和 Speech Commands 库不需要在这里再次导入，
|
||
// 假定在使用此模块的 HTML 页面中已经通过 <script> 标签引入。
|
||
|
||
// ======================= 模块内部共享变量 =======================
|
||
let recognizer; // 基础的 SpeechCommands recognizer 实例
|
||
let transferRecognizer; // 用于迁移学习的 recognizer 实例
|
||
|
||
// 用户定义的类别标签数组 (包括背景噪音)。
|
||
// 这是模块的核心状态之一，外部可以通过 getLabels() 获取。
|
||
let labels = [];
|
||
const BACKGROUND_NOISE_LABEL = '_background_noise_'; // 内部使用的背景噪音标签
|
||
|
||
let isPredicting = false; // 预测状态标志
|
||
let isRecording = false; // 录音状态标志，防止重复点击
|
||
const recordDuration = 1000; // 每个样本的录音时长 (毫秒)
|
||
|
||
let isModelTrainedFlag = false; // 手动维护模型训练状态
|
||
let predictionStopFunction = null; // 存储 transferRecognizer.listen() 返回的停止函数
|
||
|
||
|
||
// ======================= 核心功能函数 =======================
|
||
|
||
/**
|
||
* 初始化音频分类器模型。
|
||
* 必须在使用任何其他功能之前调用。
|
||
* @returns {Promise<void>} resolve 表示成功，reject 表示失败。
|
||
*/
|
||
async function init() {
|
||
// 确保每次初始化时重置状态
|
||
recognizer = null;
|
||
transferRecognizer = null;
|
||
labels = [];
|
||
isPredicting = false;
|
||
isRecording = false;
|
||
isModelTrainedFlag = false;
|
||
if (predictionStopFunction) {
|
||
predictionStopFunction(); // 停止任何正在进行的预测
|
||
predictionStopFunction = null;
|
||
}
|
||
|
||
try {
|
||
// 创建基础识别器
|
||
recognizer = speechCommands.create('BROWSER_FFT');
|
||
await recognizer.ensureModelLoaded(); // 确保模型加载完成
|
||
|
||
// 创建迁移学习识别器
|
||
transferRecognizer = recognizer.createTransfer('my-custom-model');
|
||
|
||
// 初始化时，将背景噪音标签加入到我们的内部 labels 数组
|
||
labels.push(BACKGROUND_NOISE_LABEL);
|
||
|
||
console.log('AudioClassifier: 模型和迁移学习器初始化成功');
|
||
return Promise.resolve();
|
||
} catch (error) {
|
||
console.error('AudioClassifier: 初始化失败:', error);
|
||
return Promise.reject(new Error(`模型初始化失败或麦克风无法访问: ${error.message}`));
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 录制单个或批量音频样本。
|
||
* 在录制示例时会暂时禁用 isRecording 标志以防止多重录制。
|
||
*
|
||
* @param {string} label - 样本所属的类别标签。
|
||
* @param {number} countToRecord - 要录制样本的数量。
|
||
* @param {Object} [options] - 其他选项。
|
||
* @param {Function} [options.onProgress] - (optional) 进度回调函数，参数为 (currentCount, totalCount)。
|
||
* @returns {Promise<number>} resolve 时返回该类别当前的总样本数量，reject 时返回错误。
|
||
*/
|
||
async function recordMultipleExamples(label, countToRecord = 1, options = {}) {
|
||
if (isRecording) {
|
||
return Promise.reject(new Error('核心模块：正在录制中，请等待当前录音完成。'));
|
||
}
|
||
|
||
isRecording = true;
|
||
let currentLabelCount = 0;
|
||
|
||
try {
|
||
for (let i = 0; i < countToRecord; i++) {
|
||
// console.log(`AudioClassifier: 正在录制 "${label}" 样本... (第 ${i + 1} 个 / 共 ${countToRecord} 个)`);
|
||
await transferRecognizer.collectExample(label, {
|
||
amplitudeRequired: true,
|
||
durationMillis: recordDuration
|
||
});
|
||
|
||
const exampleCounts = transferRecognizer.countExamples();
|
||
currentLabelCount = exampleCounts[label] || 0;
|
||
|
||
if (options.onProgress && typeof options.onProgress === 'function') {
|
||
options.onProgress(i + 1, countToRecord, currentLabelCount);
|
||
}
|
||
|
||
// 在每次录音之间增加短暂延迟，以便更好地分离样本
|
||
if (i < countToRecord - 1) {
|
||
await new Promise(resolve => setTimeout(resolve, Math.max(200, recordDuration / 5)));
|
||
}
|
||
}
|
||
console.log(`AudioClassifier: 已为 "${label}" 收集了 ${currentLabelCount} 个样本。`);
|
||
return Promise.resolve(currentLabelCount);
|
||
} catch (error) {
|
||
console.error(`AudioClassifier: 录制 "${label}" 样本失败:`, error);
|
||
return Promise.reject(new Error(`录制 "${label}" 样本失败: ${error.message}`));
|
||
} finally {
|
||
isRecording = false;
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 添加一个自定义类别标签到模型。
|
||
* 标签将添加到内部的 `labels` 数组中。
|
||
* @param {string} categoryName - 要添加的类别名称。
|
||
* @returns {Promise<void>} resolve 表示成功，reject 表示失败（如名称为空或重复）。
|
||
*/
|
||
function addCustomCategory(categoryName) {
|
||
if (!categoryName || categoryName.trim() === '') {
|
||
return Promise.reject(new Error('类别名称不能为空。'));
|
||
}
|
||
const trimmedName = categoryName.trim();
|
||
if (labels.some(label => label.toLowerCase() === trimmedName.toLowerCase())) {
|
||
return Promise.reject(new Error(`类别 "${trimmedName}" 已经存在。`));
|
||
}
|
||
labels.push(trimmedName);
|
||
console.log(`AudioClassifier: 已添加新类别: "${trimmedName}"`);
|
||
return Promise.resolve();
|
||
}
|
||
|
||
/**
|
||
* 检查模型是否已准备好进行训练。
|
||
* @returns {Object} 包含 `ready` (boolean) 和 `details` (string) 的对象。
|
||
*/
|
||
function checkTrainingReadiness() {
|
||
let exampleCounts = {};
|
||
try {
|
||
exampleCounts = transferRecognizer.countExamples();
|
||
} catch (error) {
|
||
// 这是预期内的错误，当没有样本时 countExamples 会抛出。
|
||
// 将其捕获并视为空样本集处理。
|
||
console.warn("countExamples() 抛出异常 (无样本时正常行为)。", error.message);
|
||
}
|
||
|
||
// 检查是否有任何样本，用于更新导出按钮等逻辑
|
||
const totalSamples = Object.values(exampleCounts).reduce((acc, count) => acc + count, 0);
|
||
const hasAnyExamples = totalSamples > 0;
|
||
|
||
const backgroundNoiseCount = exampleCounts[BACKGROUND_NOISE_LABEL] || 0;
|
||
const backgroundNoiseReady = backgroundNoiseCount > 0;
|
||
|
||
let customCategoriesWithSamples = 0;
|
||
// 遍历内部 labels 数组，检查每个自定义类别是否有样本
|
||
// 从索引 1 开始，因为 0 是背景噪音
|
||
for (let i = 1; i < labels.length; i++) {
|
||
const customLabel = labels[i];
|
||
if ((exampleCounts[customLabel] || 0) > 0) {
|
||
customCategoriesWithSamples++;
|
||
}
|
||
}
|
||
|
||
// 训练就绪的条件：必须有背景噪音样本，并且至少一个自定义类别有样本
|
||
const canTrain = backgroundNoiseReady && customCategoriesWithSamples >= 1;
|
||
|
||
// 更详细的训练就绪判断，包括样本数量建议
|
||
const MIN_SAMPLES_PER_CLASS_FOR_TRAINING = 5;
|
||
let trainingDetails = [];
|
||
|
||
if (!backgroundNoiseReady) {
|
||
trainingDetails.push('需要录制背景噪音样本。');
|
||
} else if (backgroundNoiseCount < MIN_SAMPLES_PER_CLASS_FOR_TRAINING) {
|
||
trainingDetails.push(`建议背景噪音至少 ${MIN_SAMPLES_PER_CLASS_FOR_TRAINING} 个样本 (当前: ${backgroundNoiseCount})。`);
|
||
}
|
||
|
||
if (customCategoriesWithSamples === 0) {
|
||
trainingDetails.push('需要至少一个自定义类别并录制样本。');
|
||
} else {
|
||
let insufficientCustomSamples = [];
|
||
for (let i = 1; i < labels.length; i++) {
|
||
const customLabel = labels[i];
|
||
const count = exampleCounts[customLabel] || 0;
|
||
if (count > 0 && count < MIN_SAMPLES_PER_CLASS_FOR_TRAINING) {
|
||
insufficientCustomSamples.push(`类别 "${customLabel}" 样本不足 (当前: ${count}，建议: ${MIN_SAMPLES_PER_CLASS_FOR_TRAINING})。`);
|
||
}
|
||
}
|
||
if (insufficientCustomSamples.length > 0) {
|
||
trainingDetails = trainingDetails.concat(insufficientCustomSamples);
|
||
}
|
||
}
|
||
|
||
return {
|
||
ready: canTrain && trainingDetails.length === 0, // 只有完全满足条件才认为是 ready
|
||
details: trainingDetails.join('\n') || '模型已准备好进行训练。',
|
||
hasAnyExamples: hasAnyExamples // 额外返回是否有任何样本，供 UI 判断是否启用导出按钮
|
||
};
|
||
}
|
||
|
||
/**
|
||
* 训练模型。
|
||
* @param {Object} [config] - 训练配置。
|
||
* @param {number} [config.epochs=50] - 训练迭代次数。
|
||
* @param {number} [config.batchSize=16] - 批大小。
|
||
* @param {number} [config.validationSplit=0.1] - 验证集比例。
|
||
* @param {boolean} [config.shuffle=true] - 是否打乱数据。
|
||
* @param {Function} [config.onEpochEnd] - (可选) 每个 epoch 结束时的回调函数，参数为 (epoch, logs)。
|
||
* @returns {Promise<void>} resolve 表示训练完成，reject 表示训练失败。
|
||
*/
|
||
async function trainModel(config = {}) {
|
||
// 再次进行训练前检查，防止外部在不满足条件时调用
|
||
const readiness = checkTrainingReadiness();
|
||
if (!readiness.ready) {
|
||
return Promise.reject(new Error(`模型未准备好训练：\n${readiness.details}`));
|
||
}
|
||
|
||
const defaultConfig = {
|
||
epochs: 50,
|
||
batchSize: 16,
|
||
validationSplit: 0.1,
|
||
shuffle: true,
|
||
// 外部传入的回调函数将在 train 内部被调用
|
||
callbacks: {
|
||
onEpochEnd: (epoch, logs) => {
|
||
if (config.onEpochEnd && typeof config.onEpochEnd === 'function') {
|
||
config.onEpochEnd(epoch, logs);
|
||
}
|
||
}
|
||
}
|
||
};
|
||
|
||
// 合并默认配置和用户指定的配置
|
||
const trainingConfig = { ...defaultConfig, ...config };
|
||
|
||
try {
|
||
await transferRecognizer.train(trainingConfig);
|
||
isModelTrainedFlag = true;
|
||
console.log('AudioClassifier: 模型训练完成。');
|
||
return Promise.resolve();
|
||
} catch (error) {
|
||
isModelTrainedFlag = false;
|
||
console.error('AudioClassifier: 模型训练失败:', error);
|
||
return Promise.reject(new Error(`模型训练失败: ${error.message}. 请确保有足够的样本且类别均衡。`));
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 开始实时预测。
|
||
* @param {Function} onPrediction - 预测结果回调函数，参数为 {label: string, score: number, scores: number[], labels: string[]}。
|
||
* @param {Object} [listenOptions] - 监听选项，直接传递给 transferRecognizer.listen()。
|
||
* @returns {Promise<Function>} resolve 时返回停止预测的函数，reject 时返回错误。
|
||
*/
|
||
async function startPrediction(onPrediction, listenOptions = {}) {
|
||
if (isPredicting) {
|
||
return Promise.reject(new Error('核心模块：识别已经在进行中。'));
|
||
}
|
||
|
||
if (!isModelTrainedFlag) {
|
||
return Promise.reject(new Error('核心模块：模型尚未训练完成，请先训练模型。'));
|
||
}
|
||
|
||
isPredicting = true;
|
||
|
||
const defaultOptions = {
|
||
includeEmbedding: true,
|
||
probabilityThreshold: 0.75, // 预测置信度阈值
|
||
suppressionTimeMillis: 300, // 抑制时间，避免连续触发
|
||
overlapFactor: 0.50, // 帧重叠因子
|
||
};
|
||
|
||
const options = { ...defaultOptions, ...listenOptions };
|
||
|
||
try {
|
||
predictionStopFunction = await transferRecognizer.listen(result => {
|
||
if (!isPredicting) return; // 确保在停止后不再处理结果
|
||
|
||
const classLabels = transferRecognizer.wordLabels(); // 获取模型内部的标签列表
|
||
const scores = result.scores; // 预测分数
|
||
const maxScore = Math.max(...scores);
|
||
const predictedIndex = scores.indexOf(maxScore);
|
||
|
||
let predictedLabel = classLabels[predictedIndex];
|
||
// 对背景噪音标签进行友好化处理
|
||
if (predictedLabel === BACKGROUND_NOISE_LABEL) {
|
||
predictedLabel = '背景噪音';
|
||
}
|
||
|
||
// 调用外部传入的回调函数，传递结构化的预测结果
|
||
if (typeof onPrediction === 'function') {
|
||
onPrediction({
|
||
label: predictedLabel,
|
||
score: maxScore,
|
||
scores: scores, // 原始分数数组
|
||
predictedIndex: predictedIndex, // 预测索引
|
||
labels: classLabels.map(label => label === BACKGROUND_NOISE_LABEL ? '背景噪音' : label) // 友好化的标签列表
|
||
});
|
||
}
|
||
|
||
}, options);
|
||
|
||
console.log('AudioClassifier: 开始识别。');
|
||
return Promise.resolve(predictionStopFunction); // 返回停止函数给调用者
|
||
} catch (error) {
|
||
isPredicting = false;
|
||
console.error('AudioClassifier: 开始识别失败:', error);
|
||
return Promise.reject(new Error(`开始识别失败: ${error.message}`));
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 停止实时预测。
|
||
*/
|
||
function stopPrediction() {
|
||
if (isPredicting) {
|
||
if (typeof predictionStopFunction === 'function') {
|
||
predictionStopFunction(); // 调用 SpeechCommands 库返回的停止函数
|
||
predictionStopFunction = null;
|
||
} else {
|
||
console.warn('AudioClassifier: 停止预测函数未定义或不是函数。');
|
||
}
|
||
isPredicting = false;
|
||
console.log('AudioClassifier: 已停止识别。');
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 获取当前收集到的各类别样本数量。
|
||
* @returns {Object} 键为类别名称，值为样本数量。
|
||
*/
|
||
function getExampleCounts() {
|
||
try {
|
||
return transferRecognizer.countExamples();
|
||
} catch (error) {
|
||
// 捕获无样本时的错误，返回空对象
|
||
return {};
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 获取模型是否已训练的状态。
|
||
* @returns {boolean} 如果模型已训练则返回 true，否则 false。
|
||
*/
|
||
function isModelTrained() {
|
||
return isModelTrainedFlag;
|
||
}
|
||
|
||
/**
|
||
* 获取当前模块中所有标签的数组（包括背景噪音和自定义标签）。
|
||
* @returns {string[]} 标签数组。
|
||
*/
|
||
function getLabels() {
|
||
return [...labels]; // 返回副本，避免外部直接修改内部状态
|
||
}
|
||
|
||
/**
|
||
* 序列化所有收集到的样本数据到 ArrayBuffer。
|
||
* @returns {Promise<ArrayBuffer>} 包含样本数据的 ArrayBuffer。
|
||
*/
|
||
async function serializeExamples() {
|
||
try {
|
||
const serialized = await transferRecognizer.serializeExamples();
|
||
console.log('AudioClassifier: 样本数据序列化成功。');
|
||
return Promise.resolve(serialized);
|
||
} catch (error) {
|
||
console.error('AudioClassifier: 序列化样本失败:', error);
|
||
return Promise.reject(new Error(`序列化样本失败: ${error.message}`));
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 从 ArrayBuffer 导入样本数据。
|
||
* 导入后会更新内部的 labels 数组和模型状态。
|
||
* @param {ArrayBuffer} dataBuffer - 包含样本数据的 ArrayBuffer。
|
||
* @returns {Promise<Object>} resolve 时返回导入后的样本总数统计，reject 时返回错误。
|
||
*/
|
||
async function loadExamples(dataBuffer) {
|
||
if (!dataBuffer instanceof ArrayBuffer) {
|
||
return Promise.reject(new Error('核心模块：导入数据必须是 ArrayBuffer 类型。'));
|
||
}
|
||
|
||
// 首先清空当前的 transferRecognizer 中的所有样本，再加载新的。
|
||
// 这等同于创建一个新的 transferRecognizer 实例，并同步其内部状态。
|
||
// SpeechCommands库没有直接的`clearExamples`或`resetExamples`方法，
|
||
// 最稳妥的方法是重新创建transferRecognizer并加载。
|
||
// 但loadExamples本身就是设计来覆盖和更新现有样本的，
|
||
// 考虑到性能，我们假设直接loadExamples即可。
|
||
// 注意：loadExamples会清除原有数据。
|
||
|
||
try {
|
||
await transferRecognizer.loadExamples(dataBuffer);
|
||
|
||
// 成功加载后，需要更新内部的 `labels` 数组，以匹配导入的数据。
|
||
// `transferRecognizer.wordLabels()` 返回的是模型内部的实际标签。
|
||
const loadedLabels = transferRecognizer.wordLabels();
|
||
// 确保 BACKGROUND_NOISE_LABEL 还在第一个位置，并过滤掉 '_version_', '_numExamples_' 等内部标签
|
||
labels = loadedLabels.filter(label => label !== '_version_' && label !== '_numExamples_');
|
||
|
||
if (!labels.includes(BACKGROUND_NOISE_LABEL)) {
|
||
// 如果导入的数据中没有背景噪音（不常见，但作为健壮性处理）
|
||
labels.unshift(BACKGROUND_NOISE_LABEL);
|
||
} else {
|
||
// 确保背景噪音在第一位
|
||
const bgIndex = labels.indexOf(BACKGROUND_NOISE_LABEL);
|
||
if (bgIndex > 0) {
|
||
const bgLabel = labels.splice(bgIndex, 1);
|
||
labels.unshift(bgLabel[0]);
|
||
}
|
||
}
|
||
|
||
// 重置模型的训练状态，因为样本已更改，需要重新训练
|
||
isModelTrainedFlag = false;
|
||
|
||
const exampleCounts = transferRecognizer.countExamples();
|
||
console.log('AudioClassifier: 样本数据导入成功。', exampleCounts);
|
||
return Promise.resolve(exampleCounts);
|
||
} catch (error) {
|
||
console.error('AudioClassifier: 加载样本失败:', error);
|
||
return Promise.reject(new Error(`加载样本失败: ${error.message}. 请确保文件有效且未损坏。`));
|
||
}
|
||
}
|
||
|
||
|
||
// ======================= 导出的公共接口 =======================
|
||
// 这里使用 CommonJS 风格的 module.exports，因为这通常用于 Node.js 环境或 Webpack/Rollup 等打包工具。
|
||
// 如果你是在浏览器中直接作为 <script> 标签引入（没有打包工具），
|
||
// 并且希望它在全局作用域下可用，可以使用 `window.AudioClassifier = {...}`。
|
||
// 在这种情况下，我将提供 `window.AudioClassifier` 的方式。
|
||
|
||
// 为了 Web 浏览器环境，将核心功能挂载到全局对象上
|
||
window.AudioClassifier = {
|
||
init, // 初始化模型
|
||
recordMultipleExamples, // 录制样本
|
||
addCustomCategory, // 添加自定义类别
|
||
checkTrainingReadiness, // 检查模型是否可训练
|
||
trainModel, // 训练模型
|
||
startPrediction, // 开始预测
|
||
stopPrediction, // 停止预测
|
||
getExampleCounts, // 获取所有类别样本数量
|
||
isModelTrained, // 检查模型是否已训练
|
||
getLabels, // 获取当前所有标签
|
||
serializeExamples, // 导出样本数据
|
||
loadExamples // 导入样本数据
|
||
};
|