Spaces:
Sleeping
Sleeping
| document.addEventListener('DOMContentLoaded', () => { | |
| const videoElement = document.getElementById('userVideo'); | |
| const canvasElement = document.createElement('canvas'); // Offscreen canvas | |
| const context = canvasElement.getContext('2d'); | |
| const startButton = document.getElementById('startTrainerBtn'); | |
| const stopButton = document.getElementById('stopTrainerBtn'); | |
| const exerciseTypeSelect = document.getElementById('exerciseType'); | |
| // UI Displays for feedback | |
| const repsDisplay = document.getElementById('repsDisplay'); | |
| const stageDisplay = document.getElementById('stageDisplay'); | |
| const feedbackDisplay = document.getElementById('feedbackDisplay'); | |
| const angleDisplay = document.getElementById('angleDisplay'); | |
| const repsLeftHcDisplay = document.getElementById('repsLeftHc'); | |
| const repsRightHcDisplay = document.getElementById('repsRightHc'); | |
| const stageLeftHcDisplay = document.getElementById('stageLeftHc'); | |
| const stageRightHcDisplay = document.getElementById('stageRightHc'); | |
| const feedbackLeftHcDisplay = document.getElementById('feedbackLeftHc'); | |
| const feedbackRightHcDisplay = document.getElementById('feedbackRightHc'); | |
| const angleLeftCurlHcDisplay = document.getElementById('angleLeftCurlHc'); | |
| const angleRightCurlHcDisplay = document.getElementById('angleRightCurlHc'); | |
| const squatPushupUIDiv = document.getElementById('squatPushupUI'); | |
| const hammerCurlUIDiv = document.getElementById('hammerCurlUI'); | |
| const apiStatusDisplay = document.getElementById('apiStatus'); | |
| const sessionIdDisplay = document.getElementById('sessionIdDisplay'); | |
| let socket = null; | |
| let currentExerciseType = ''; | |
| let isSessionActive = false; | |
| let animationFrameId = null; | |
| let stream = null; | |
| const logger = { | |
| log: (message) => console.log(`[Trainer Log] ${new Date().toISOString()}: ${message}`), | |
| error: (message) => console.error(`[Trainer Error] ${new Date().toISOString()}: ${message}`), | |
| info: (message) => console.info(`[Trainer Info] ${new Date().toISOString()}: ${message}`) | |
| }; | |
| logger.log("WebSocket trainer script loaded. Waiting for DOM and user interaction."); | |
| function setupUIForExercise(exerciseType) { | |
| logger.info(`Setting up UI for exercise: ${exerciseType}`); | |
| if (exerciseType === 'hammer_curl') { | |
| if (squatPushupUIDiv) squatPushupUIDiv.style.display = 'none'; | |
| if (hammerCurlUIDiv) hammerCurlUIDiv.style.display = 'block'; | |
| } else { | |
| if (squatPushupUIDiv) squatPushupUIDiv.style.display = 'block'; | |
| if (hammerCurlUIDiv) hammerCurlUIDiv.style.display = 'none'; | |
| } | |
| } | |
| function updateGenericUI(data) { | |
| if (repsDisplay) repsDisplay.textContent = data.counter !== undefined ? data.counter : 'N/A'; | |
| if (stageDisplay) stageDisplay.textContent = data.stage || 'N/A'; | |
| if (feedbackDisplay) feedbackDisplay.textContent = data.feedback || 'N/A'; | |
| let angleText = 'Angle: -'; | |
| if (data.angle_left !== undefined && data.angle_right !== undefined) { | |
| angleText = `Angles L: ${Math.round(data.angle_left)}, R: ${Math.round(data.angle_right)}`; | |
| } else if (data.angle_body_left !== undefined && data.angle_body_right !== undefined) { | |
| angleText = `Body Angles L: ${Math.round(data.angle_body_left)}, R: ${Math.round(data.angle_body_right)}`; | |
| } | |
| if (angleDisplay) angleDisplay.textContent = angleText; | |
| } | |
| function updateHammerCurlUI(data) { | |
| if (repsLeftHcDisplay) repsLeftHcDisplay.textContent = data.counter_left !== undefined ? data.counter_left : 'N/A'; | |
| if (repsRightHcDisplay) repsRightHcDisplay.textContent = data.counter_right !== undefined ? data.counter_right : 'N/A'; | |
| if (stageLeftHcDisplay) stageLeftHcDisplay.textContent = data.stage_left || 'N/A'; | |
| if (stageRightHcDisplay) stageRightHcDisplay.textContent = data.stage_right || 'N/A'; | |
| if (feedbackLeftHcDisplay) feedbackLeftHcDisplay.textContent = data.feedback_left || 'N/A'; | |
| if (feedbackRightHcDisplay) feedbackRightHcDisplay.textContent = data.feedback_right || 'N/A'; | |
| let curlAngleText = 'Curl Angles: -'; | |
| if (data.angle_left_curl !== undefined && data.angle_right_curl !== undefined) { | |
| curlAngleText = `Curl Angles L: ${Math.round(data.angle_left_curl)}, R: ${Math.round(data.angle_right_curl)}`; | |
| } | |
| // Assuming you want to display curl angles in the general angleDisplay or have a specific one | |
| if (angleDisplay) angleDisplay.textContent = curlAngleText; | |
| } | |
| function clearUIFeedback() { | |
| logger.info("Clearing UI feedback elements."); | |
| if (repsDisplay) repsDisplay.textContent = '0'; | |
| if (feedbackDisplay) feedbackDisplay.textContent = 'N/A'; | |
| if (stageDisplay) stageDisplay.textContent = 'N/A'; | |
| if (angleDisplay) angleDisplay.textContent = 'Angle: -'; | |
| if (repsLeftHcDisplay) repsLeftHcDisplay.textContent = '0'; | |
| if (repsRightHcDisplay) repsRightHcDisplay.textContent = '0'; | |
| if (feedbackLeftHcDisplay) feedbackLeftHcDisplay.textContent = 'N/A'; | |
| if (feedbackRightHcDisplay) feedbackRightHcDisplay.textContent = 'N/A'; | |
| if (stageLeftHcDisplay) stageLeftHcDisplay.textContent = 'N/A'; | |
| if (stageRightHcDisplay) stageRightHcDisplay.textContent = 'N/A'; | |
| if (angleLeftCurlHcDisplay) angleLeftCurlHcDisplay.textContent = '0'; | |
| if (angleRightCurlHcDisplay) angleRightCurlHcDisplay.textContent = '0'; | |
| if (apiStatusDisplay) apiStatusDisplay.textContent = "Idle"; | |
| if (sessionIdDisplay) sessionIdDisplay.textContent = "-"; | |
| } | |
| function initializeSocket() { | |
| if (socket && socket.connected) { | |
| logger.info("Socket already connected."); | |
| return; | |
| } | |
| // Ensure the server URL is correct, especially if not running on the same host/port | |
| // For Hugging Face Spaces, it should connect to the same origin, so io() is fine. | |
| socket = io(); | |
| socket.on('connect', () => { | |
| logger.info(`Connected to server with SID: ${socket.id}`); | |
| if (apiStatusDisplay) apiStatusDisplay.textContent = "Connected to server."; | |
| }); | |
| socket.on('disconnect', (reason) => { | |
| logger.info(`Disconnected from server: ${reason}`); | |
| if (apiStatusDisplay) apiStatusDisplay.textContent = "Disconnected. Refresh if needed."; | |
| stopTrainer(false); | |
| }); | |
| socket.on('connect_error', (error) => { | |
| logger.error(`Connection Error: ${error.message}`); | |
| if (apiStatusDisplay) apiStatusDisplay.textContent = "Connection Error."; | |
| }); | |
| socket.on('connection_ack', (data) => { | |
| logger.info(`Server Acknowledged Connection: ${JSON.stringify(data)}`); | |
| if (sessionIdDisplay && data.sid) sessionIdDisplay.textContent = data.sid; | |
| }); | |
| socket.on('session_started', (data) => { | |
| logger.info(`Session started on server: ${data.session_id} for exercise: ${data.exercise_type}`); | |
| isSessionActive = true; | |
| // currentExerciseType should have been set by startButton click | |
| if (apiStatusDisplay) apiStatusDisplay.textContent = `Session for ${currentExerciseType} started.`; | |
| // Check if video is already playing and ready | |
| if (videoElement && videoElement.srcObject && !videoElement.paused && videoElement.readyState >= videoElement.HAVE_FUTURE_DATA) { | |
| logger.info("Video ready and session started, initiating frame loop."); | |
| startFrameLoop(); | |
| } else { | |
| logger.info("Session started, but video not yet playing/ready. Frame loop will start via video 'onplaying' or 'oncanplaythrough' event."); | |
| } | |
| }); | |
| socket.on('session_error', (data) => { | |
| logger.error(`Session error: ${data.error}`); | |
| if (apiStatusDisplay) apiStatusDisplay.textContent = `Session Error: ${data.error}`; | |
| alert(`Session Error: ${data.error}`); | |
| stopTrainer(false); | |
| }); | |
| socket.on('exercise_update', (data) => { | |
| if (!isSessionActive) return; | |
| if (data.success) { | |
| if (data.landmarks_detected) { | |
| // logger.log('Exercise Update:', data.data); | |
| if (currentExerciseType === 'hammer_curl') { | |
| updateHammerCurlUI(data.data); | |
| } else { | |
| updateGenericUI(data.data); | |
| } | |
| } else { | |
| if (currentExerciseType === 'hammer_curl') { | |
| if (feedbackLeftHcDisplay) feedbackLeftHcDisplay.textContent = `Feedback: ${data.message || 'No landmarks detected.'}`; | |
| if (feedbackRightHcDisplay) feedbackRightHcDisplay.textContent = ''; | |
| } else { | |
| if (feedbackDisplay) feedbackDisplay.textContent = `Feedback: ${data.message || 'No landmarks detected.'}`; | |
| } | |
| } | |
| } else { | |
| logger.error(`Exercise update indicated failure: ${data.message || 'Unknown error'}`); | |
| if (feedbackDisplay && currentExerciseType !== 'hammer_curl') feedbackDisplay.textContent = `Feedback: Error - ${data.message || 'Processing error'}`; | |
| else if (feedbackLeftHcDisplay && currentExerciseType === 'hammer_curl') feedbackLeftHcDisplay.textContent = `Feedback: Error - ${data.message || 'Processing error'}`; | |
| } | |
| }); | |
| socket.on('frame_error', (data) => { | |
| logger.error(`Frame processing error from server: ${data.error}`); | |
| if (feedbackDisplay && currentExerciseType !== 'hammer_curl') feedbackDisplay.textContent = `Error: ${data.error}`; | |
| else if (feedbackLeftHcDisplay && currentExerciseType === 'hammer_curl') feedbackLeftHcDisplay.textContent = `Error: ${data.error}`; | |
| }); | |
| } | |
| function startFrameLoop() { | |
| logger.info(`Attempting to start frame loop. isSessionActive: ${isSessionActive}, video.paused: ${videoElement.paused}, video.ended: ${videoElement.ended}, video.readyState: ${videoElement.readyState}`); | |
| if (isSessionActive && videoElement && videoElement.srcObject && !videoElement.paused && !videoElement.ended && videoElement.readyState >= videoElement.HAVE_ENOUGH_DATA) { //HAVE_ENOUGH_DATA is 4 | |
| logger.info("Conditions met, starting sendFrameLoop via requestAnimationFrame."); | |
| if (animationFrameId) { | |
| cancelAnimationFrame(animationFrameId); | |
| } | |
| sendFrameLoop(); | |
| } else { | |
| logger.warn(`Conditions not met to start frame loop. isSessionActive: ${isSessionActive}, srcObject: ${!!videoElement.srcObject}, paused: ${videoElement.paused}, ended: ${videoElement.ended}, readyState: ${videoElement.readyState}. Will try again if video becomes ready.`); | |
| } | |
| } | |
| async function sendFrameLoop() { | |
| if (!isSessionActive || !videoElement.srcObject || videoElement.paused || videoElement.ended || videoElement.readyState < videoElement.HAVE_FUTURE_DATA) { // HAVE_FUTURE_DATA is 3 | |
| logger.info(`[sendFrameLoop] Stopping. Conditions: isSessionActive=${isSessionActive}, srcObject=${!!videoElement.srcObject}, paused=${videoElement.paused}, ended=${videoElement.ended}, readyState=${videoElement.readyState}`); | |
| if (animationFrameId) { | |
| cancelAnimationFrame(animationFrameId); | |
| animationFrameId = null; | |
| } | |
| // If session is supposed to be active but video state is bad, consider stopping the session. | |
| if (isSessionActive && (videoElement.paused || videoElement.ended || videoElement.readyState < videoElement.HAVE_FUTURE_DATA)) { | |
| logger.warn("[sendFrameLoop] Video stream issue detected, stopping session from client-side."); | |
| // stopTrainer(true); // This might cause a loop if called from here. Better to just stop sending. | |
| } | |
| return; | |
| } | |
| try { | |
| if (canvasElement.width > 0 && canvasElement.height > 0) { | |
| context.drawImage(videoElement, 0, 0, canvasElement.width, canvasElement.height); | |
| const imageDataBase64 = canvasElement.toDataURL('image/jpeg', 0.7).split(',')[1]; | |
| if (socket && socket.connected) { | |
| // console.log(`[sendFrameLoop] Sending frame for exercise: ${currentExerciseType}`); // Can be too verbose | |
| socket.emit('process_frame', { | |
| image: imageDataBase64, | |
| exercise_type: currentExerciseType, | |
| frame_width: canvasElement.width, | |
| frame_height: canvasElement.height | |
| }); | |
| } else { | |
| logger.error("[sendFrameLoop] Socket not connected. Cannot send frame."); | |
| stopTrainer(false); | |
| } | |
| } else { | |
| // This might happen if video metadata hasn't loaded yet, or if the video dimensions are truly zero. | |
| logger.warn(`[sendFrameLoop] Canvas dimensions are zero (w:${canvasElement.width}, h:${canvasElement.height}) or video not ready, skipping frame draw/send.`); | |
| } | |
| } catch (error) { | |
| logger.error("Error in sendFrameLoop (drawing or emitting):", error); | |
| stopTrainer(false); | |
| } | |
| if (isSessionActive) { | |
| animationFrameId = requestAnimationFrame(sendFrameLoop); | |
| } | |
| } | |
| async function startTrainer() { | |
| if (isSessionActive) { | |
| logger.info("Session already active. Please stop current session first."); | |
| return; | |
| } | |
| currentExerciseType = exerciseTypeSelect.value; | |
| setupUIForExercise(currentExerciseType); | |
| logger.info(`Attempting to start exercise: ${currentExerciseType}`); | |
| if (apiStatusDisplay) apiStatusDisplay.textContent = "Initializing..."; | |
| if (startButton) startButton.disabled = true; | |
| if (stopButton) stopButton.disabled = false; | |
| if (exerciseTypeSelect) exerciseTypeSelect.disabled = true; | |
| if (!socket || !socket.connected) { | |
| logger.info("Socket not connected. Initializing socket connection..."); | |
| initializeSocket(); | |
| } | |
| // Emit start_exercise_session. The 'session_started' event will then trigger the frame loop. | |
| // Add a small delay to ensure socket might connect if it was just initialized. | |
| setTimeout(() => { | |
| if (socket && socket.connected) { | |
| logger.info(`Socket connected, emitting start_exercise_session for ${currentExerciseType}`); | |
| socket.emit('start_exercise_session', { exercise_type: currentExerciseType }); | |
| } else { | |
| logger.error("Socket not connected after delay. Cannot start session."); | |
| if (apiStatusDisplay) apiStatusDisplay.textContent = "Error: Cannot connect to server."; | |
| if (startButton) startButton.disabled = false; | |
| if (stopButton) stopButton.disabled = true; | |
| if (exerciseTypeSelect) exerciseTypeSelect.disabled = false; | |
| } | |
| }, 500); // 500ms delay to allow socket to connect if just initialized | |
| try { | |
| logger.info("Requesting user media..."); | |
| stream = await navigator.mediaDevices.getUserMedia({ video: { width: 640, height: 480 }, audio: false }); | |
| logger.info("User media obtained."); | |
| videoElement.srcObject = stream; | |
| videoElement.onloadedmetadata = () => { | |
| logger.info("Video metadata loaded."); | |
| canvasElement.width = videoElement.videoWidth; | |
| canvasElement.height = videoElement.videoHeight; | |
| logger.info(`Canvas dimensions set to: ${canvasElement.width}x${canvasElement.height}`); | |
| videoElement.play().catch(playError => { | |
| logger.error(`Error playing video: ${playError}`); | |
| if (apiStatusDisplay) apiStatusDisplay.textContent = "Error playing video."; | |
| alert(`Error playing video: ${playError.message}`); | |
| stopTrainer(false); | |
| }); | |
| }; | |
| videoElement.oncanplay = () => { | |
| logger.info("Video event: canplay. Video should be ready to play."); | |
| // Attempt to start frame loop if session is active and video is playable | |
| if(isSessionActive && videoElement.readyState >= videoElement.HAVE_FUTURE_DATA) { | |
| startFrameLoop(); | |
| } | |
| }; | |
| videoElement.oncanplaythrough = () => { | |
| logger.info("Video event: canplaythrough. Video has enough data to play without interruption."); | |
| if(isSessionActive) { | |
| startFrameLoop(); | |
| } | |
| }; | |
| videoElement.onplaying = () => { | |
| logger.info("Video event: playing. Video has started playing."); | |
| if(isSessionActive) { | |
| startFrameLoop(); | |
| } | |
| }; | |
| videoElement.onpause = () => { | |
| logger.info("Video event: pause. Stopping frame loop."); | |
| isSessionActive = false; | |
| if (animationFrameId) { | |
| cancelAnimationFrame(animationFrameId); | |
| animationFrameId = null; | |
| } | |
| }; | |
| videoElement.onended = () => { | |
| logger.info("Video event: ended. Stopping trainer."); | |
| stopTrainer(true); | |
| }; | |
| } catch (err) { | |
| logger.error(`Error accessing webcam: ${err}`); | |
| if (apiStatusDisplay) apiStatusDisplay.textContent = "Error accessing webcam."; | |
| alert(`Could not access webcam: ${err.message}`); | |
| if (startButton) startButton.disabled = false; | |
| if (stopButton) stopButton.disabled = true; | |
| if (exerciseTypeSelect) exerciseTypeSelect.disabled = false; | |
| } | |
| } | |
| function stopTrainer(emitToServer = true) { | |
| logger.info("Stop trainer called."); | |
| isSessionActive = false; | |
| if (animationFrameId) { | |
| cancelAnimationFrame(animationFrameId); | |
| animationFrameId = null; | |
| logger.info("Frame loop cancelled."); | |
| } | |
| if (stream) { | |
| stream.getTracks().forEach(track => track.stop()); | |
| stream = null; | |
| logger.info("Webcam tracks stopped."); | |
| } | |
| if (videoElement) { | |
| videoElement.srcObject = null; | |
| videoElement.pause(); // Explicitly pause | |
| } | |
| if (socket && socket.connected && emitToServer) { | |
| // No specific 'end_exercise_session' event needed as server handles 'disconnect' | |
| // but if you want explicit cleanup, you can add one. | |
| // socket.emit('stop_exercise_session'); | |
| logger.info("Client-side stop initiated. Server will handle via disconnect or if an explicit stop event is implemented."); | |
| } | |
| if (apiStatusDisplay) apiStatusDisplay.textContent = "Trainer stopped. Select an exercise to start."; | |
| if (startButton) startButton.disabled = false; | |
| if (stopButton) stopButton.disabled = true; | |
| if (exerciseTypeSelect) exerciseTypeSelect.disabled = false; | |
| clearUIFeedback(); | |
| if (sessionIdDisplay) sessionIdDisplay.textContent = '-'; | |
| } | |
| // Initialize Socket.IO connection when the script loads | |
| if (startButton && stopButton && exerciseTypeSelect) { // Ensure essential controls are present | |
| initializeSocket(); | |
| startButton.addEventListener('click', startTrainer); | |
| stopButton.addEventListener('click', () => stopTrainer(true)); | |
| exerciseTypeSelect.addEventListener('change', (event) => { | |
| if (!isSessionActive) { | |
| currentExerciseType = event.target.value; | |
| setupUIForExercise(currentExerciseType); | |
| } | |
| }); | |
| // Initial UI setup based on default selected exercise | |
| setupUIForExercise(exerciseTypeSelect.value); | |
| } else { | |
| logger.error("One or more essential UI elements (startButton, stopButton, exerciseTypeSelect) not found."); | |
| } | |
| clearUIFeedback(); | |
| }); | |