503 lines
18 KiB
JavaScript
503 lines
18 KiB
JavaScript
/**
|
||
* =============================================================================
|
||
* 动态版 - 姿态识别与模型管理脚本 (v2.1)
|
||
* - 新增自动采集样本功能
|
||
* =============================================================================
|
||
* 功能列表:
|
||
* - 实时姿态检测 (MoveNet)
|
||
* - KNN 分类器训练
|
||
* - 实时姿态预测
|
||
* - 坐标完美对齐 (Canvas与Video重叠)
|
||
* - 动态添加/删除/重命名姿态类别
|
||
* - 模型导出为包含类别信息的 JSON 文件
|
||
* - 从 JSON 文件导入模型并恢复类别状态
|
||
* - ✅ 新增：自动采集10次样本，间隔0.3秒
|
||
* =============================================================================
|
||
*/
|
||
|
||
'use strict';
|
||
|
||
// --- 全局变量和常量 ---
|
||
const videoElement = document.getElementById('video');
|
||
const canvasElement = document.getElementById('canvas');
|
||
const canvasCtx = canvasElement.getContext('2d');
|
||
const statusElement = document.getElementById('status');
|
||
const resultElement = document.getElementById('result-text');
|
||
|
||
// UI元素
|
||
const poseClassesContainer = document.getElementById('pose-classes-container');
|
||
const addClassButton = document.getElementById('btn-add-class');
|
||
const predictButton = document.getElementById('btn-predict');
|
||
const exportButton = document.getElementById('btn-export');
|
||
const importButton = document.getElementById('btn-import');
|
||
const fileImporter = document.getElementById('file-importer');
|
||
|
||
let detector, classifier, animationFrameId;
|
||
let isPredicting = false;
|
||
let isAutoCollecting = false; // 新增：标记是否正在进行自动采集
|
||
|
||
// 📌 核心状态管理: 使用一个对象来管理所有动态状态
|
||
const appState = {
|
||
classMap: {}, // 存储 classId -> className 的映射, e.g., {0: '姿态 A', 1: '姿态 B'}
|
||
nextClassId: 0 // 用于生成唯一的 classId
|
||
};
|
||
|
||
// --- 主应用逻辑 ---
|
||
|
||
/**
|
||
* 初始化应用，加载模型并设置摄像头
|
||
*/
|
||
async function init() {
|
||
try {
|
||
classifier = knnClassifier.create();
|
||
const detectorConfig = { modelType: poseDetection.movenet.modelType.SINGLEPOSE_LIGHTNING };
|
||
detector = await poseDetection.createDetector(poseDetection.SupportedModels.MoveNet, detectorConfig);
|
||
|
||
await setupCamera();
|
||
setupEventListeners();
|
||
mainLoop();
|
||
|
||
statusElement.innerText = "模型和摄像头已就绪!";
|
||
enableControls();
|
||
addNewClass(); // 默认创建第一个类别
|
||
|
||
} catch (error) {
|
||
console.error("初始化失败:", error);
|
||
statusElement.innerText = "初始化失败，请检查摄像头权限或刷新。";
|
||
statusElement.style.backgroundColor = '#fce8e6';
|
||
statusElement.style.color = '#d93025';
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 设置和启动用户摄像头
|
||
*/
|
||
async function setupCamera() {
|
||
const stream = await navigator.mediaDevices.getUserMedia({ video: true });
|
||
videoElement.srcObject = stream;
|
||
return new Promise((resolve) => {
|
||
videoElement.onloadedmetadata = () => {
|
||
videoElement.play();
|
||
canvasElement.width = videoElement.videoWidth;
|
||
canvasElement.height = videoElement.videoHeight;
|
||
resolve();
|
||
};
|
||
});
|
||
}
|
||
|
||
/**
|
||
* 为所有交互式元素绑定事件监听器
|
||
*/
|
||
function setupEventListeners() {
|
||
addClassButton.addEventListener('click', addNewClass);
|
||
predictButton.addEventListener('click', togglePrediction);
|
||
exportButton.addEventListener('click', exportModel);
|
||
importButton.addEventListener('click', () => fileImporter.click());
|
||
fileImporter.addEventListener('change', importModel);
|
||
}
|
||
|
||
// --- 动态类别管理 ---
|
||
|
||
/**
|
||
* 动态创建一个新类别的UI元素并添加到页面
|
||
* @param {number} classId - 类别的唯一ID
|
||
* @param {string} className - 类别的名称
|
||
*/
|
||
function createClassUI(classId, className) {
|
||
const poseClassDiv = document.createElement('div');
|
||
poseClassDiv.className = 'pose-class';
|
||
poseClassDiv.dataset.classId = classId;
|
||
|
||
// 📌 修改这里：添加 btn-auto-sample 按钮
|
||
poseClassDiv.innerHTML = `
|
||
<div class="class-info">
|
||
<input type="text" class="class-name-input" value="${className}" data-class-id="${classId}">
|
||
<span class="sample-count">(0 样本)</span>
|
||
</div>
|
||
<div class="class-actions">
|
||
<button class="btn-auto-sample" data-class-id="${classId}">自动采集</button>
|
||
<button class="btn-sample" data-class-id="${classId}">采集样本</button>
|
||
<button class="btn-delete-class" title="删除类别" data-class-id="${classId}">×</button>
|
||
</div>
|
||
`;
|
||
|
||
poseClassesContainer.appendChild(poseClassDiv);
|
||
|
||
// 为新创建的元素绑定事件
|
||
const nameInput = poseClassDiv.querySelector('.class-name-input');
|
||
nameInput.addEventListener('change', (e) => {
|
||
appState.classMap[classId] = e.target.value;
|
||
});
|
||
|
||
const autoSampleButton = poseClassDiv.querySelector('.btn-auto-sample'); // 新增
|
||
autoSampleButton.addEventListener('click', () => toggleAutoCollection(classId, autoSampleButton)); // 新增
|
||
|
||
const sampleButton = poseClassDiv.querySelector('.btn-sample');
|
||
sampleButton.addEventListener('click', () => addExample(classId));
|
||
|
||
// 初始化时根据预测状态禁用按钮
|
||
if (isPredicting) {
|
||
sampleButton.disabled = true;
|
||
autoSampleButton.disabled = true; // 新增
|
||
}
|
||
|
||
const deleteButton = poseClassDiv.querySelector('.btn-delete-class');
|
||
deleteButton.addEventListener('click', () => deleteClass(classId));
|
||
}
|
||
|
||
/**
|
||
* 添加一个新的姿态类别
|
||
*/
|
||
function addNewClass() {
|
||
const classId = appState.nextClassId;
|
||
const className = `Class ${classId + 1}`;
|
||
appState.classMap[classId] = className;
|
||
appState.nextClassId++;
|
||
createClassUI(classId, className);
|
||
}
|
||
|
||
/**
|
||
* 删除一个指定的姿态类别
|
||
* @param {number} classId - 要删除的类别的ID
|
||
*/
|
||
function deleteClass(classId) {
|
||
if (confirm(`确定要删除类别 "${appState.classMap[classId]}" 吗？所有样本都将丢失。`)) {
|
||
// 从UI中移除
|
||
const elementToRemove = poseClassesContainer.querySelector(`[data-class-id="${classId}"]`);
|
||
if (elementToRemove) elementToRemove.remove();
|
||
|
||
// 从状态和分类器中移除
|
||
delete appState.classMap[classId];
|
||
classifier.clearClass(classId);
|
||
|
||
updateSampleCounts();
|
||
updatePredictionUI(); // 检查是否还有类别可以预测
|
||
checkExportAbility();
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 采集一个姿态样本并添加到KNN分类器
|
||
* @param {number} classId 类别的ID
|
||
*/
|
||
async function addExample(classId) {
|
||
const poses = await detector.estimatePoses(videoElement, { flipHorizontal: true });
|
||
if (poses && poses.length > 0) {
|
||
const poseTensor = flattenPose(poses[0]);
|
||
classifier.addExample(poseTensor, classId);
|
||
poseTensor.dispose();
|
||
|
||
updateSampleCounts();
|
||
checkExportAbility();
|
||
console.log(`为类别 ${appState.classMap[classId]} 采集1个样本。`);
|
||
return true; // 表示采集成功
|
||
} else {
|
||
console.warn(`为类别 ${appState.classMap[classId]} 采集样本失败，未检测到姿态。`);
|
||
return false; // 表示采集失败
|
||
}
|
||
}
|
||
|
||
// --- 新增：自动采集逻辑 ---
|
||
let autoCollectionIntervalId = null; // 用于存储 setInterval ID
|
||
let autoCollectionCount = 0; // 计数器
|
||
const AUTO_COLLECTION_TOTAL = 10; // 总共采集次数
|
||
const AUTO_COLLECTION_INTERVAL = 300; // 间隔时间 0.3 秒
|
||
|
||
async function toggleAutoCollection(classId, buttonElement) {
|
||
if (isAutoCollecting) {
|
||
// 如果正在自动采集，则停止
|
||
stopAutoCollection(buttonElement);
|
||
} else {
|
||
// 否则，开始自动采集
|
||
startAutoCollection(classId, buttonElement);
|
||
}
|
||
}
|
||
|
||
async function startAutoCollection(classId, buttonElement) {
|
||
isAutoCollecting = true;
|
||
autoCollectionCount = 0;
|
||
|
||
// 禁用其他采集和预测按钮
|
||
predictButton.disabled = true;
|
||
exportButton.disabled = true;
|
||
importButton.disabled = true;
|
||
addClassButton.disabled = true;
|
||
document.querySelectorAll('.btn-sample, .btn-auto-sample, .btn-delete-class, .class-name-input').forEach(btn => {
|
||
if (btn !== buttonElement) { // 不禁用当前自动采集按钮
|
||
btn.disabled = true;
|
||
}
|
||
if (btn.classList.contains('class-name-input')) btn.disabled = true;
|
||
});
|
||
|
||
buttonElement.innerText = `停止采集 (0/${AUTO_COLLECTION_TOTAL})`;
|
||
buttonElement.classList.add('stop'); // 添加停止样式
|
||
|
||
const performCollection = async () => {
|
||
if (autoCollectionCount < AUTO_COLLECTION_TOTAL) {
|
||
const success = await addExample(classId); // 调用手动采集功能
|
||
if (success) {
|
||
autoCollectionCount++;
|
||
}
|
||
buttonElement.innerText = `停止采集 (${autoCollectionCount}/${AUTO_COLLECTION_TOTAL})`;
|
||
} else {
|
||
stopAutoCollection(buttonElement);
|
||
alert(`类别 "${appState.classMap[classId]}" 自动采集完成！`);
|
||
}
|
||
};
|
||
|
||
// 立即执行一次，然后设置定时器
|
||
await performCollection();
|
||
if (autoCollectionCount < AUTO_COLLECTION_TOTAL) {
|
||
autoCollectionIntervalId = setInterval(performCollection, AUTO_COLLECTION_INTERVAL);
|
||
}
|
||
}
|
||
|
||
function stopAutoCollection(buttonElement) {
|
||
clearInterval(autoCollectionIntervalId);
|
||
autoCollectionIntervalId = null;
|
||
isAutoCollecting = false;
|
||
buttonElement.innerText = '自动采集';
|
||
buttonElement.classList.remove('stop'); // 移除停止样式
|
||
|
||
// 重新启用按钮（根据应用状态）
|
||
updatePredictionUI(); // 根据预测状态重新启用/禁用相关按钮
|
||
enableControls(); // 重新启用添加类别、导出、导入按钮
|
||
}
|
||
|
||
// --- 模型与预测逻辑 ---
|
||
|
||
/**
|
||
* 开始或停止姿态预测
|
||
*/
|
||
function togglePrediction() {
|
||
if (classifier.getNumClasses() === 0) {
|
||
alert("请先为至少一个姿态采集样本后再开始预测！");
|
||
return;
|
||
}
|
||
isPredicting = !isPredicting;
|
||
updatePredictionUI();
|
||
}
|
||
|
||
/**
|
||
* 应用的主循环
|
||
*/
|
||
async function mainLoop() {
|
||
const poses = await detector.estimatePoses(videoElement, { flipHorizontal: true });
|
||
canvasCtx.clearRect(0, 0, canvasElement.width, canvasElement.height);
|
||
|
||
if (poses && poses.length > 0) {
|
||
drawPose(poses[0]);
|
||
// 只有当不在自动采集状态时才进行预测
|
||
if (isPredicting && classifier.getNumClasses() > 0 && !isAutoCollecting) {
|
||
const poseTensor = flattenPose(poses[0]);
|
||
const result = await classifier.predictClass(poseTensor, 3);
|
||
poseTensor.dispose();
|
||
|
||
const confidence = Math.round(result.confidences[result.label] * 100);
|
||
const predictedClassName = appState.classMap[result.label] || '未知类别';
|
||
resultElement.innerText = `姿态: ${predictedClassName} (${confidence}%)`;
|
||
} else if (isAutoCollecting) {
|
||
resultElement.innerText = "自动采集中...";
|
||
}
|
||
} else {
|
||
resultElement.innerText = "未检测到姿态";
|
||
}
|
||
animationFrameId = requestAnimationFrame(mainLoop);
|
||
}
|
||
|
||
// --- 模型管理函数 (已更新以支持动态类别) ---
|
||
|
||
/**
|
||
* 导出KNN模型为包含类别信息的JSON文件
|
||
*/
|
||
function exportModel() {
|
||
if (classifier.getNumClasses() === 0) {
|
||
alert('模型中还没有任何样本，无法导出！');
|
||
return;
|
||
}
|
||
|
||
const dataset = classifier.getClassifierDataset();
|
||
const datasetObj = {};
|
||
Object.keys(dataset).forEach((key) => {
|
||
const data = dataset[key];
|
||
datasetObj[key] = data.arraySync();
|
||
});
|
||
|
||
const modelData = {
|
||
classMap: appState.classMap,
|
||
dataset: datasetObj
|
||
};
|
||
|
||
const jsonStr = JSON.stringify(modelData);
|
||
const blob = new Blob([jsonStr], { type: "application/json" });
|
||
const url = URL.createObjectURL(blob);
|
||
const a = document.createElement('a');
|
||
a.href = url;
|
||
a.download = `pose-knn-model.json`;
|
||
document.body.appendChild(a);
|
||
a.click();
|
||
document.body.removeChild(a);
|
||
URL.revokeObjectURL(url);
|
||
}
|
||
|
||
/**
|
||
* 从JSON文件导入KNN模型并恢复类别状态
|
||
* @param {Event} event
|
||
*/
|
||
function importModel(event) {
|
||
const file = event.target.files[0];
|
||
if (!file) return;
|
||
|
||
const reader = new FileReader();
|
||
reader.onload = (e) => {
|
||
try {
|
||
const modelData = JSON.parse(e.target.result);
|
||
|
||
if (!modelData.classMap || !modelData.dataset) {
|
||
throw new Error("无效的模型文件格式。");
|
||
}
|
||
|
||
// 1. 清理现有状态
|
||
classifier.clearAllClasses();
|
||
poseClassesContainer.innerHTML = '';
|
||
appState.classMap = {};
|
||
|
||
// 2. 加载新状态
|
||
appState.classMap = modelData.classMap;
|
||
const classIds = Object.keys(appState.classMap).map(Number);
|
||
appState.nextClassId = classIds.length > 0 ? Math.max(...classIds) + 1 : 0;
|
||
|
||
// 3. 恢复UI
|
||
classIds.forEach(id => {
|
||
createClassUI(id, appState.classMap[id]);
|
||
});
|
||
|
||
// 4. 加载模型数据
|
||
const newDataset = {};
|
||
Object.keys(modelData.dataset).forEach((key) => {
|
||
newDataset[key] = tf.tensor(modelData.dataset[key]);
|
||
});
|
||
classifier.setClassifierDataset(newDataset);
|
||
|
||
updateSampleCounts();
|
||
checkExportAbility();
|
||
alert('模型导入成功！');
|
||
|
||
} catch (error) {
|
||
console.error("导入模型失败:", error);
|
||
alert(`导入失败！请确保文件是正确的模型JSON文件。\n错误: ${error.message}`);
|
||
} finally {
|
||
fileImporter.value = '';
|
||
}
|
||
};
|
||
reader.readAsText(file);
|
||
}
|
||
|
||
|
||
// --- 辅助和UI更新函数 ---
|
||
|
||
function flattenPose(pose) {
|
||
const keypoints = pose.keypoints.map(p => [p.x / videoElement.videoWidth, p.y / videoElement.videoHeight]).flat();
|
||
return tf.tensor(keypoints);
|
||
}
|
||
|
||
function drawPose(pose) {
|
||
// 绘制关键点和骨骼...
|
||
if (pose.keypoints) {
|
||
// 绘制关键点
|
||
for (const keypoint of pose.keypoints) {
|
||
if (keypoint.score > 0.3) {
|
||
canvasCtx.beginPath();
|
||
canvasCtx.arc(keypoint.x, keypoint.y, 5, 0, 2 * Math.PI);
|
||
canvasCtx.fillStyle = '#1a73e8';
|
||
canvasCtx.fill();
|
||
}
|
||
}
|
||
// 绘制骨骼连接线
|
||
const adjacentPairs = poseDetection.util.getAdjacentPairs(poseDetection.SupportedModels.MoveNet);
|
||
adjacentPairs.forEach(([i, j]) => {
|
||
const kp1 = pose.keypoints[i];
|
||
const kp2 = pose.keypoints[j];
|
||
if (kp1.score > 0.3 && kp2.score > 0.3) {
|
||
canvasCtx.beginPath();
|
||
canvasCtx.moveTo(kp1.x, kp1.y);
|
||
canvasCtx.lineTo(kp2.x, kp2.y);
|
||
canvasCtx.strokeStyle = 'blue';
|
||
canvasCtx.lineWidth = 2;
|
||
canvasCtx.stroke();
|
||
}
|
||
});
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 更新所有类别UI上的样本数量
|
||
*/
|
||
function updateSampleCounts() {
|
||
const dataset = classifier.getClassifierDataset();
|
||
const allClassElements = document.querySelectorAll('.pose-class');
|
||
allClassElements.forEach(el => {
|
||
const classId = parseInt(el.dataset.classId, 10);
|
||
const classInfo = dataset[classId];
|
||
const count = classInfo ? classInfo.shape[0] : 0;
|
||
el.querySelector('.sample-count').innerText = `(${count} 样本)`;
|
||
});
|
||
}
|
||
|
||
/**
|
||
* 根据状态更新UI
|
||
*/
|
||
function updatePredictionUI() {
|
||
// 禁用所有采集按钮（包括手动和自动）和删除按钮
|
||
document.querySelectorAll('.btn-sample, .btn-auto-sample, .btn-delete-class').forEach(btn => btn.disabled = isPredicting || isAutoCollecting);
|
||
|
||
// 禁用添加类别和导入模型的按钮
|
||
addClassButton.disabled = isPredicting || isAutoCollecting;
|
||
importButton.disabled = isPredicting || isAutoCollecting;
|
||
|
||
// 禁用类别名称输入框
|
||
document.querySelectorAll('.class-name-input').forEach(input => input.disabled = isPredicting || isAutoCollecting);
|
||
|
||
if (isPredicting) {
|
||
predictButton.innerText = "停止预测";
|
||
predictButton.classList.add('stop');
|
||
resultElement.innerText = "正在分析...";
|
||
} else {
|
||
predictButton.innerText = "开始预测";
|
||
predictButton.classList.remove('stop');
|
||
resultElement.innerText = "已停止";
|
||
}
|
||
// 只有在有类别且有样本时才能预测
|
||
predictButton.disabled = isPredicting ? false : classifier.getNumClasses() === 0 || isAutoCollecting;
|
||
checkExportAbility();
|
||
}
|
||
|
||
/**
|
||
* 通用启用/禁用控件 (在自动采集停止后调用)
|
||
*/
|
||
function enableControls() {
|
||
// 重新评估所有按钮的状态
|
||
// 自动采集按钮的状态由其自身管理
|
||
predictButton.disabled = classifier.getNumClasses() === 0;
|
||
importButton.disabled = false; // 导入按钮总是可以手动启用
|
||
addClassButton.disabled = false;
|
||
checkExportAbility(); // 重新检查导出按钮
|
||
updatePredictionUI(); // 再次调用，确保其他按钮状态正确
|
||
}
|
||
|
||
/** 检查是否可以导出模型并更新按钮状态 */
|
||
function checkExportAbility() {
|
||
exportButton.disabled = isPredicting || classifier.getNumClasses() === 0 || isAutoCollecting;
|
||
}
|
||
|
||
function cleanup() {
|
||
if (detector) detector.dispose();
|
||
if (classifier) classifier.clearAllClasses();
|
||
if (animationFrameId) cancelAnimationFrame(animationFrameId);
|
||
if (autoCollectionIntervalId) clearInterval(autoCollectionIntervalId); // 清理自动采集定时器
|
||
}
|
||
|
||
// --- 启动应用 ---
|
||
window.onbeforeunload = cleanup;
|
||
init();
|