Teachable Machine > P5.js > Device (web serial)

function preload() {
  console.log('preload() started.');
  // We won't load the model here anymore.
  // Model loading will happen when the user enters a URL and clicks OK.
  console.log('preload() finished.');
}

function setup() {
  console.log('setup() started.');
  createCanvas(windowWidth, windowHeight); // Create a responsive canvas
  
  // Create the video capture from the webcam
  video = createCapture(VIDEO, videoReady); // Add a callback for video readiness
  video.hide(); // Hide the HTML video element as we'll draw it on the canvas
  console.log('Video capture object created in setup().');
  // Initialize custom ready flag to false
  video.ready = false;

  // Initial display message asking for model URL
  textSize(32);
  textAlign(CENTER, CENTER);
  fill(0);
  text('Please enter your Teachable Machine model URL below', width / 2, height / 2);

  // --- Model URL Input Setup ---
  modelURLInput = createInput();
  modelURLInput.attribute('placeholder', 'Enter Teachable Machine model URL here');
  modelURLInput.style('font-size', '16px');
  modelURLInput.style('padding', '8px');
  // FIX: videomargin is now defined, so this line will work.
  modelURLInput.style('width', 'calc(100% - 100px - 2 * ' + videoMargin + 'px)'); // Adjust width
  modelURLInput.style('border', '1px solid #ccc');
  modelURLInput.style('border-radius', '4px');

  modelURLLoadButton = createButton('OK');
  modelURLLoadButton.mousePressed(loadModelFromInput);
  modelURLLoadButton.style('font-size', '16px');
  modelURLLoadButton.style('padding', '10px 15px');
  modelURLLoadButton.style('background-color', '#007BFF');
  modelURLLoadButton.style('color', 'white');
  modelURLLoadButton.style('border', 'none');
  modelURLLoadButton.style('border-radius', '5px');
  modelURLLoadButton.style('cursor', 'pointer');

  // --- Web Serial Setup ---
  // Create p5 DOM elements for serial status and received data
  serialStatusP = createP('Web Serial: Not connected.');
  serialStatusP.style('font-size', '18px');
  serialStatusP.style('color', '#333');
  serialStatusP.style('font-family', 'sans-serif'); // Add a standard font for consistency

  receivedSerialDataP = createP('Received: '); // Renamed for clarity
  receivedSerialDataP.style('font-size', '18px');
  receivedSerialDataP.style('color', '#333');
  receivedSerialDataP.style('font-family', 'sans-serif'); // Add a standard font for consistency

  // Check if Web Serial API is supported
  if ('serial' in navigator) {
    // Create the "Connect Serial Port" button
    serialConnectButton = createButton('Connect Serial Port');
    serialConnectButton.mousePressed(connectSerial); // Call connectSerial() when button is pressed
    serialConnectButton.style('font-size', '16px');
    serialConnectButton.style('padding', '10px 15px');
    serialConnectButton.style('background-color', '#4CAF50');
    serialConnectButton.style('color', 'white');
    serialConnectButton.style('border', 'none');
    serialConnectButton.style('border-radius', '5px');
    serialConnectButton.style('cursor', 'pointer');

    // Add event listener for port disconnection (e.g., USB cable unplugged)
    navigator.serial.addEventListener('disconnect', portDisconnected);

  } else {
    // Display an error message if Web Serial is not supported
    serialStatusP.html('Web Serial API not supported in this browser. Try Chrome/Edge.');
    serialStatusP.style('color', 'red');
  }

  // Set initial visibility for UI elements
  setClassifierUIVisibility(false); // Hide classifier elements initially

  // Call windowResized once to set initial positions correctly
  windowResized();
  console.log('setup() finished. appState:', appState);
}

function videoReady() {
  console.log('Video capture is ready! (videoReady() callback executed)'); // Critical log
  video.ready = true; // Custom flag to indicate video is ready
  console.log('Video element readyState:', video.elt.readyState); // Check HTML video element state
  console.log('Video dimensions: ', video.width, 'x', video.height); // Log actual video dimensions

  // If model is already loaded and we're in classifier state, start classifying
  if (appState === 'classifier' && classifier && classifier.ready) {
    console.log('Video ready and classifier ready. Starting classification from videoReady().');
    classifyVideo();
  }
}

/**
 * Hides/shows UI elements belonging to the classifier view.
 * @param {boolean} isVisible - True to show, false to hide.
 */
function setClassifierUIVisibility(isVisible) {
  const displayStyle = isVisible ? 'block' : 'none';
  
  if (serialStatusP) serialStatusP.style('display', displayStyle);
  if (receivedSerialDataP) receivedSerialDataP.style('display', displayStyle);
  if (serialConnectButton) serialConnectButton.style('display', displayStyle);

  // Input view elements are visible when classifier elements are hidden
  if (modelURLInput) modelURLInput.style('display', isVisible ? 'none' : 'block');
  if (modelURLLoadButton) modelURLLoadButton.style('display', isVisible ? 'none' : 'block');
  console.log('Classifier UI visibility set to:', isVisible);
}

function loadModelFromInput() {
  console.log('loadModelFromInput() called.');
  const newModelURL = modelURLInput.value().trim();
  if (newModelURL) {
    modelURL = newModelURL;
    console.log('Attempting to load model from:', modelURL); // Log attempt
    currentClassification.label = 'Loading model...';
    classifier = ml5.imageClassifier(modelURL, modelLoaded); // Re-initialize classifier
    classifier.ready = false; // Reset ready flag
    modelURLInput.attribute('disabled', ''); // Disable input during loading
    modelURLLoadButton.attribute('disabled', ''); // Disable button during loading
  } else {
    alert('Please enter a valid Teachable Machine model URL.');
    console.warn('No model URL entered.');
  }
}

// Modified modelLoaded function to handle errors
function modelLoaded(error) {
  console.log('modelLoaded() callback received.');
  if (error) {
    console.error('Error loading model:', error); // Log the error
    currentClassification.label = 'Model loading failed: ' + error.message;
    // Re-enable input and button so user can try again
    modelURLInput.removeAttribute('disabled');
    modelURLLoadButton.removeAttribute('disabled');
    // Keep appState as 'input'
    console.log('Model loading failed, staying in input state.');
    return; // Stop here if loading failed
  }
  console.log('Model Loaded successfully!'); // Confirm model is loaded
  classifier.ready = true; // Set a custom flag to indicate the model is ready
  currentClassification.label = 'Model Loaded. Starting classification...';
  // Check if video is also ready before starting classification
  if (video && video.ready) {
    console.log('Classifier ready and video ready. Starting classification from modelLoaded().');
    classifyVideo(); // Start classifying once the model is loaded AND video is ready
  } else {
    console.warn('Model loaded, but video not yet ready. Will start classification when video is ready.');
  }
  modelURLInput.removeAttribute('disabled'); // Enable input after loading
  modelURLLoadButton.removeAttribute('disabled'); // Enable button after loading
  appState = 'classifier'; // Switch to classifier view
  setClassifierUIVisibility(true); // Show classifier elements
  windowResized(); // Re-position elements for the classifier view
  console.log('modelLoaded() finished. appState:', appState);
}

function classifyVideo() {
  // Ensure the video is loaded and the model is ready before attempting to classify
  if (video && video.ready && classifier && classifier.ready) {
    console.log('Classifying video frame.');
    // Classify the current video frame and call gotResult when done
    // Reference: https://learn.ml5js.org/docs/#/reference/image-classifier?id=classify
    classifier.classify(video, gotResult);
  } else {
    // If not ready, try again after a short delay
    // This is important if video or classifier isn't ready simultaneously
    console.warn('classifyVideo() called, but video or classifier not ready. Retrying in 100ms.');
    setTimeout(classifyVideo, 100);
  }
}

function gotResult(error, results) {
  // If there's an error during classification, log it
  if (error) {
    console.error('Error during classification:', error);
    currentClassification.label = 'Error: ' + error.message;
    currentClassification.confidence = 0;
    // Don't stop classifying, try again
    classifyVideo();
    return;
  }

  // If no error, results will be an array of objects
  // Each object has a 'label' and 'confidence'
  // Sort results by confidence in descending order to get the top prediction
  results.sort((a, b) => b.confidence - a.confidence);
  const topResult = results[0];

  // Update the current classification result
  currentClassification.label = topResult.label;
  currentClassification.confidence = topResult.confidence;

  // Keep classifying the video feed recursively
  classifyVideo();
}

function draw() {
  background(220); // Clear the canvas each frame

  if (appState === 'input') {
    // Display input view message
    textSize(32);
    textAlign(CENTER, CENTER);
    fill(0);
    text('Please enter your Teachable Machine model URL below', width / 2, height / 2 - 50); // Slightly above input field
  } else if (appState === 'classifier') {
    if (video && video.ready && classifier && classifier.ready) { // Check video.ready here too
      console.log('Attempting to draw video. Video ready:', video.ready, 'Classifier ready:', classifier.ready); // Critical log
      // Calculate preview size while maintaining aspect ratio
      // Let's make the preview height about half the canvas height, minus margins.
      let previewHeight = height / 2 - 2 * videoMargin;
      let previewWidth = (previewHeight / video.height) * video.width;

      // Ensure previewWidth doesn't exceed available space if canvas is very narrow
      // (e.g., if canvas width is less than previewWidth + videoMargin * 2)
      let availableWidthForVideo = width - 2 * videoMargin;
      if (previewWidth > availableWidthForVideo) {
        previewWidth = availableWidthForVideo;
        previewHeight = (previewWidth / video.width) * video.height;
      }

      // Draw the video feed at the top-left with margin
      image(video, videoMargin, videoMargin, previewWidth, previewHeight);

      // Classification text and background rectangle to the right of the video preview.
      let textRectX = videoMargin + previewWidth + videoMargin;
      let textRectWidth = width - textRectX - videoMargin;
      let textRectHeight = previewHeight; // Make it the same height as the video preview
      let textRectY = videoMargin; // Align with the top of the video

      fill(255, 200); // White with 200 alpha
      noStroke();
      rect(textRectX, textRectY, textRectWidth, textRectHeight);

      // Display the current classification result
      fill(0);
      textSize(24);
      textAlign(LEFT, CENTER);
      // Center text vertically within its rectangle
      let textYOffset = textRectY + textRectHeight / 2;
      text(`Label: ${currentClassification.label}`, textRectX + 10, textYOffset - 20); // Adjust y for two lines
      // Use nf() to format confidence as a percentage with 2 decimal places
      text(`Confidence: ${nf(currentClassification.confidence * 100, 0, 2)}%`, textRectX + 10, textYOffset + 20);

      // --- Web Serial: Send classification data ---
      let numberToSend = -1; // Default to -1 for unknown labels
      const currentLabel = currentClassification.label.toLowerCase(); // Use lowercase for comparison
      switch (currentLabel) {
        case 'ihminen':
          numberToSend = 1;
          break;
        case 'koira':
          numberToSend = 2;
          break;
        case 'kissa': // Assuming 'kissa' for number 3 based on common classification tasks
          numberToSend = 3;
          break;
        default:
          numberToSend = -1; // For any other label, or if label is empty/null
          break;
      }

      // Send the classification number if connected,
      // the classification label has changed,
      // enough time has passed, AND a valid number was mapped.
      if (port && writer && currentLabel !== lastSentLabel.toLowerCase() && millis() - lastClassificationSendTime > classificationSendInterval) {
        if (numberToSend !== -1) {
          const dataToSend = numberToSend + '\n'; // Convert number to string and add newline for microcontroller parsing
          writeSerial(dataToSend);
          lastSentLabel = currentLabel; // Update lastSentLabel with the *string* classification (lowercase)
          lastClassificationSendTime = millis();
        } else {
          console.warn(`Unrecognized or default label "${currentLabel}". Not sending data.`);
        }
      }
    } else if (!classifier.ready) {
      // Display loading message if the model is not yet ready (should only happen briefly after OK)
      textSize(32);
      textAlign(CENTER, CENTER);
      fill(0);
      text('Loading model...', width / 2, height / 2);
    }
  }
}

function windowResized() {
  resizeCanvas(windowWidth, windowHeight); // Adjust canvas size on window resize
  console.log('windowResized() called. appState:', appState);

  if (appState === 'input') {
    // When in input state, position input/button centrally below the initial text
    let inputWidth = min(width - 2 * videoMargin, 400); // Max 400px wide, or less if canvas is narrow
    if (modelURLInput) modelURLInput.style('width', `${inputWidth - modelURLLoadButton.width - 20}px`); // Adjust for button width and spacing
    if (modelURLInput) modelURLInput.position(width / 2 - inputWidth / 2, height / 2 + 50); // Below initial text
    if (modelURLLoadButton) modelURLLoadButton.position(width / 2 + inputWidth / 2 - modelURLLoadButton.width, height / 2 + 50); // To the right of input
  } else { // In classifier state, use the standard layout
    // Recalculate video preview size based on new window dimensions
    let previewHeight = height / 2 - 2 * videoMargin;
    let previewWidth = (previewHeight / video.height) * video.width;
    let availableWidthForVideo = width - 2 * videoMargin;
    if (previewWidth > availableWidthForVideo) {
      previewWidth = availableWidthForVideo;
      previewHeight = (previewWidth / video.width) * video.height;
    }

    // Calculate the new starting y-position for the DOM elements below the video
    // It's the bottom of the video (videoMargin + previewHeight) + an additional videoMargin for spacing
    let domElementsStartY = videoMargin + previewHeight + videoMargin;

    // Re-position model URL input and button (now below video, not centered)
    if (modelURLInput) modelURLInput.style('width', 'calc(100% - 100px - 2 * ' + videoMargin + 'px)'); // Reset width
    if (modelURLInput) modelURLInput.position(videoMargin, domElementsStartY);
    if (modelURLLoadButton) modelURLLoadButton.position(videoMargin + modelURLInput.width + 10, domElementsStartY); // 10px spacing

    // Re-position other p5 DOM elements
    // Adjust starting Y for other elements based on the model URL input/button height
    let otherDomElementsStartY = domElementsStartY + 40; // Assuming button height is around 30-35px + 10px spacing
    
    if (serialStatusP) serialStatusP.position(videoMargin, otherDomElementsStartY); // Adjusted y
    if (receivedSerialDataP) receivedSerialDataP.position(videoMargin, otherDomElementsStartY + 30); // Adjusted y
    if (serialConnectButton) serialConnectButton.position(videoMargin, otherDomElementsStartY + 70); // Adjusted y
  }
}

/**
 * Connects to a serial port selected by the user.
 */
async function connectSerial() {
  try {
    console.log('connectSerial() called.');
    // Request a serial port from the user
    port = await navigator.serial.requestPort();
    serialStatusP.html('Web Serial: Connecting...');
    console.log('Serial port requested and granted.');

    // Open the port with a specified baud rate
    // IMPORTANT: This baud rate (9600) must match the one used by your microcontroller (e.g., Arduino Serial.begin(9600);)
    await port.open({ baudRate: 9600 });
    serialStatusP.html('Web Serial: Connected!');
    serialConnectButton.html('Disconnect Serial Port'); // Change button text
    serialConnectButton.mousePressed(disconnectSerial); // Change button action to disconnect
    console.log('Serial port opened with baudRate 9600.');

    // Create TextDecoderStream and TextEncoderStream for easy text handling
    const textDecoder = new TextDecoderStream();
    const textEncoder = new TextEncoderStream();
    reader = textDecoder.readable.getReader(); // Reader for incoming text data
    writer = textEncoder.writable.getWriter(); // Writer for outgoing text data
    console.log('TextDecoderStream and TextEncoderStream created.');

    // Pipe the port's readable stream through the decoder
    port.readable.pipeTo(textDecoder.writable);
    // TÄMÄ ON OIKEIN: Pipe the encoder's readable stream (containing encoded bytes) to the port's writable stream
    textEncoder.readable.pipeTo(port.writable); // THIS IS THE CORRECTION
    console.log('Serial streams piped.');

    readSerial(); // Start continuously reading data from the port
  } catch (error) {
    console.error('Error connecting to serial port:', error);
    serialStatusP.html('Web Serial: Connection failed. ' + error.message);
  }
}

/**
 * Disconnects from the currently connected serial port.
 */
async function disconnectSerial() {
  try {
    console.log('disconnectSerial() called.');
    if (reader) await reader.cancel(); // Cancel any pending read operations
    if (writer) await writer.close(); // Close the writer stream
    if (port) await port.close(); // Close the serial port

    portDisconnected(); // Reset variables and UI
    console.log('Serial port successfully disconnected.');
  } catch (error) {
    console.error('Error disconnecting serial port:', error);
    serialStatusP.html('Web Serial: Disconnection failed. ' + error.message);
  }
}

/**
 * Continuously reads data from the serial port.
 */
async function readSerial() {
  console.log('readSerial() started.');
  while (port && port.readable) {
    try {
      const { value, done } = await reader.read();
      if (done) {
        console.log('Reader closed.');
        break; // Exit loop if reader is closed
      }
      
      // Check if the received value is "1" (and trim whitespace for robustness)
      if (value.trim() === '1') {
        receivedSerialDataP.html('Received: remote device connected');
        // You might want to update serialStatusP too, but currentStatusP is already "Connected!"
      } else {
        receivedSerialDataP.html(`Received: ${value}`); // Display received data
      }
      
      // For more complex data, you might parse 'value' here (e.g., JSON.parse(value))
    } catch (error) {
      console.error('Error reading from serial port:', error);
      serialStatusP.html('Web Serial: Read error. ' + error.message);
      break; // Exit loop on error
    }
  }
  portDisconnected(); // Assume disconnected if loop exits (e.g., due to error or port closure)
  console.log('readSerial() stopped.');
}

/**
 * Writes data to the serial port.
 * @param {string} data - The string data to send.
 */
async function writeSerial(data) {
  if (port && writer) {
    try {
      await writer.write(data);
      console.log('Sent number:', data.trim()); // Log sent data (the number) without the trailing newline
    } catch (error) {
      console.error('Error writing to serial port:', error);
      serialStatusP.html('Web Serial: Write error. ' + error.message);
      portDisconnected(); // Assume disconnected on write error
    }
  } else {
    console.warn('Cannot write: Serial port not connected.');
    serialStatusP.html('Web Serial: Cannot write, not connected.');
  }
}

/**
 * Resets serial variables and UI when the port disconnects.
 */
function portDisconnected() {
  console.log('Serial port disconnected.');
  port = undefined;
  reader = undefined;
  writer = undefined;
  serialStatusP.html('Web Serial: Disconnected.');
  receivedSerialDataP.html('Received: ');
  if (serialConnectButton) {
    serialConnectButton.html('Connect Serial Port');
    serialConnectButton.mousePressed(connectSerial); // Change button action back to connect
  }
  lastSentLabel = ''; // Reset to ensure label is sent again on reconnect
  lastClassificationSendTime = 0; // Reset timer too
}