실전: 손글씨 인식 앱

// 캔버스 요소 설정
const canvas = document.getElementById('drawCanvas');
const ctx = canvas.getContext('2d');
canvas.width = 280;
canvas.height = 280;

// 배경을 검은색으로 설정 (MNIST 형식)
ctx.fillStyle = 'black';
ctx.fillRect(0, 0, canvas.width, canvas.height);

// 그리기 설정
ctx.strokeStyle = 'white';
ctx.lineWidth = 15;
ctx.lineCap = 'round';
ctx.lineJoin = 'round';

let isDrawing = false;

canvas.addEventListener('mousedown', (e) => {
  isDrawing = true;
  ctx.beginPath();
  ctx.moveTo(e.offsetX, e.offsetY);
});

canvas.addEventListener('mousemove', (e) => {
  if (!isDrawing) return;
  ctx.lineTo(e.offsetX, e.offsetY);
  ctx.stroke();
});

canvas.addEventListener('mouseup', () => {
  isDrawing = false;
  predictDigit();
});

import * as tf from '@tensorflow/tfjs';

async function createAndTrainModel() {
  const model = tf.sequential();

  model.add(tf.layers.conv2d({
    inputShape: [28, 28, 1],
    filters: 32,
    kernelSize: 3,
    activation: 'relu',
  }));
  model.add(tf.layers.maxPooling2d({ poolSize: 2 }));
  model.add(tf.layers.conv2d({
    filters: 64,
    kernelSize: 3,
    activation: 'relu',
  }));
  model.add(tf.layers.maxPooling2d({ poolSize: 2 }));
  model.add(tf.layers.flatten());
  model.add(tf.layers.dropout({ rate: 0.25 }));
  model.add(tf.layers.dense({ units: 128, activation: 'relu' }));
  model.add(tf.layers.dense({ units: 10, activation: 'softmax' }));

  model.compile({
    optimizer: 'adam',
    loss: 'categoricalCrossentropy',
    metrics: ['accuracy'],
  });

  // MNIST 데이터로 학습
  // (실제로는 MNIST 데이터를 로드하는 유틸리티 사용)
  await model.fit(trainImages, trainLabels, {
    epochs: 10,
    batchSize: 128,
    validationSplit: 0.1,
    callbacks: {
      onEpochEnd: (epoch, logs) => {
        const status = document.getElementById('status');
        status.textContent =
          "학습 중... 에포크 " + (epoch + 1) + "/10" +
          " (정확도: " + (logs.val_acc * 100).toFixed(1) + "%)";
      },
    },
  });

  return model;
}

let model;

async function init() {
  const status = document.getElementById('status');
  status.textContent = '모델 로딩 중...';

  // 저장된 모델 불러오기 또는 새로 학습
  try {
    model = await tf.loadLayersModel('localstorage://mnist-model');
    status.textContent = '모델 로드 완료!';
  } catch (e) {
    model = await createAndTrainModel();
    await model.save('localstorage://mnist-model');
    status.textContent = '모델 학습 및 저장 완료!';
  }
}

async function predictDigit() {
  if (!model) return;

  const tensor = tf.tidy(() => {
    // 캔버스를 28x28 텐서로 변환
    const imageData = tf.browser.fromPixels(canvas, 1);
    const resized = tf.image.resizeBilinear(imageData, [28, 28]);
    const normalized = resized.div(255.0);
    return normalized.expandDims(0);
  });

  const prediction = model.predict(tensor);
  const probabilities = await prediction.data();

  // 결과 표시
  const resultDiv = document.getElementById('result');
  const maxProb = Math.max(...probabilities);
  const digit = probabilities.indexOf(maxProb);

  resultDiv.innerHTML =
    '<h2>예측: ' + digit + '</h2>' +
    '<p>확률: ' + (maxProb * 100).toFixed(1) + '%</p>';

  // 모든 숫자의 확률 바 차트 표시
  const chartDiv = document.getElementById('chart');
  chartDiv.innerHTML = '';
  for (let i = 0; i < 10; i++) {
    const pct = (probabilities[i] * 100).toFixed(1);
    const bar = document.createElement('div');
    bar.className = 'prob-bar';
    bar.innerHTML =
      '<span>' + i + '</span>' +
      '<div class="bar" style="width:' + pct + '%"></div>' +
      '<span>' + pct + '%</span>';
    chartDiv.appendChild(bar);
  }

  tensor.dispose();
  prediction.dispose();
}

function clearCanvas() {
  ctx.fillStyle = 'black';
  ctx.fillRect(0, 0, canvas.width, canvas.height);
  document.getElementById('result').innerHTML = '';
  document.getElementById('chart').innerHTML = '';
}

// 초기화
init();