script.js

const LABEL_TO_DIGIT = {
  zero: 0,
  one: 1,
  two: 2,
  three: 3,
  four: 4,
  five: 5,
  six: 6,
  seven: 7,
  eight: 8,
  nine: 9,
};

let nn;
let canvas;
let input;
let hasDrawn = false;

// DOM handles, filled in on setup()
const ui = {};

function setup() {
  canvas = createCanvas(400, 400);
  canvas.parent("stage");
  background(0);

  cacheElements();
  buildScoreRows();

  ui.clear.addEventListener("click", clearCanvas);

  nn = ml5.neuralNetwork({
    inputs: [56, 56, 4],
    task: "imageClassification",
    debug: true,
  });

  const modelDetails = {
    model: "model/model.json",
    metadata: "model/model_meta.json",
    weights: "model/model.weights.bin",
  };

  nn.load(modelDetails, () => {
    console.log("Model Loaded!");
    setStatus("ready", "Model ready");
    classifyCanvas();
  });
}

function cacheElements() {
  ui.stage = document.getElementById("stage");
  ui.hint = document.getElementById("hint");
  ui.clear = document.getElementById("clear");
  ui.status = document.querySelector(".status");
  ui.statusText = document.querySelector(".status__text");
  ui.guessDigit = document.getElementById("guess-digit");
  ui.guessConfidence = document.getElementById("guess-confidence");
  ui.scores = document.getElementById("scores");
  ui.rows = {};
}

function buildScoreRows() {
  for (let digit = 0; digit <= 9; digit++) {
    const row = document.createElement("li");
    row.className = "score";
    row.innerHTML = `
      <span class="score__digit">${digit}</span>
      <span class="score__value">0.0%</span>`;
    ui.scores.appendChild(row);
    ui.rows[digit] = {
      row,
      value: row.querySelector(".score__value"),
    };
  }
}

function setStatus(state, text) {
  ui.status.dataset.status = state;
  ui.statusText.textContent = text;
}

function draw() {
  if (mouseIsPressed && isOnCanvas()) {
    if (!hasDrawn) {
      hasDrawn = true;
      ui.hint.hidden = true;
    }
    stroke(255);
    strokeWeight(40);
    strokeCap(ROUND);
    line(mouseX, mouseY, pmouseX, pmouseY);
  }
}

function isOnCanvas() {
  return mouseX >= 0 && mouseX <= width && mouseY >= 0 && mouseY <= height;
}

function clearCanvas() {
  background(0);
  hasDrawn = false;
  ui.hint.hidden = false;
}

function keyPressed() {
  if (keyCode === 67) clearCanvas();
}

function classifyCanvas() {
  input = createGraphics(28, 28);
  input.copy(canvas, 0, 0, width, height, 0, 0, 28, 28);

  nn.classify({ image: input }, (error, result) => {
    if (error) {
      console.error(error);
      setStatus("error", "Classification failed");
      return;
    }

    renderResults(result);
    classifyCanvas();
  });
}

function renderResults(result) {
  const top = result[0];
  const topDigit = LABEL_TO_DIGIT[top.label];

  ui.guessDigit.textContent = hasDrawn ? topDigit : "–";
  ui.guessConfidence.textContent = hasDrawn
    ? `${(top.confidence * 100).toFixed(1)}%`
    : "—";

  for (const { label, confidence } of result) {
    const digit = LABEL_TO_DIGIT[label];
    const row = ui.rows[digit];
    if (!row) continue;

    row.value.textContent = `${(confidence * 100).toFixed(1)}%`;
    row.row.classList.toggle("score--top", hasDrawn && digit === topDigit);
  }
}