이미지 분류

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

// CNN 모델: 이미지 패턴 인식에 특화
const model = tf.sequential();

// 합성곱 레이어 1
model.add(tf.layers.conv2d({
  inputShape: [28, 28, 1],  // 28x28 그레이스케일 이미지
  filters: 32,              // 32개의 필터
  kernelSize: 3,            // 3x3 필터 크기
  activation: 'relu',
}));
model.add(tf.layers.maxPooling2d({ poolSize: 2 }));

// 합성곱 레이어 2
model.add(tf.layers.conv2d({
  filters: 64,
  kernelSize: 3,
  activation: 'relu',
}));
model.add(tf.layers.maxPooling2d({ poolSize: 2 }));

// Flatten: 2D -> 1D
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.dropout({ rate: 0.5 }));
model.add(tf.layers.dense({ units: 10, activation: 'softmax' }));

model.summary();

// 캔버스에서 이미지 데이터 가져오기
function getImageTensor(canvas) {
  return tf.tidy(() => {
    // 캔버스를 텐서로 변환
    const imageTensor = tf.browser.fromPixels(canvas, 1);
    // 그레이스케일, 1채널

    // 28x28로 리사이즈
    const resized = tf.image.resizeBilinear(imageTensor, [28, 28]);

    // 0~1 범위로 정규화
    const normalized = resized.div(255.0);

    // 배치 차원 추가: [28, 28, 1] -> [1, 28, 28, 1]
    return normalized.expandDims(0);
  });
}

// 파일 입력에서 이미지 로드
async function loadImage(file) {
  return new Promise((resolve) => {
    const reader = new FileReader();
    reader.onload = (e) => {
      const img = new Image();
      img.onload = () => resolve(img);
      img.src = e.target.result;
    };
    reader.readAsDataURL(file);
  });
}

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

// 학습 실행
const history = await model.fit(trainImages, trainLabels, {
  epochs: 10,
  batchSize: 64,
  validationSplit: 0.2,  // 20%를 검증 데이터로 사용
  callbacks: {
    onEpochEnd: (epoch, logs) => {
      console.log(
        "에포크 " + (epoch + 1) +
        " - accuracy: " + logs.acc.toFixed(4) +
        " - val_accuracy: " + logs.val_acc.toFixed(4)
      );
    },
  },
});

// 테스트 데이터로 평가
const evalResult = model.evaluate(testImages, testLabels);
console.log("테스트 정확도:", evalResult[1].dataSync()[0].toFixed(4));

// 단일 이미지 분류
async function classifyImage(model, canvas) {
  const tensor = getImageTensor(canvas);
  const predictions = model.predict(tensor);

  // 확률 배열 가져오기
  const probabilities = await predictions.data();
  tensor.dispose();
  predictions.dispose();

  // 클래스별 확률 정리
  const classNames = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'];
  const results = classNames.map((name, i) => ({
    className: name,
    probability: probabilities[i],
  }));

  // 확률 높은 순으로 정렬
  results.sort((a, b) => b.probability - a.probability);

  return results;
}

// 사용 예시
const results = await classifyImage(model, canvas);
console.log("예측:", results[0].className);
console.log("확률:", (results[0].probability * 100).toFixed(1) + "%");