TensorFlow.js チュートリアル + LocalStorage で画像分類Webアプリを作ってみた

Last updated at 2023-05-12Posted at 2023-05-12

はじめに

TensorFlow.jsの画像転移学習チュートリアルに学習データ保存機能を付けた簡単な画像分類webアプリを記述します

アプリ概要

webブラウザ上で画像学習ができ、それらを分類するアプリ
- 学習方法は事前トレーニング済みモデルを用いた転移学習
- 学習データは保存、読込み可能
  - localStorageを用いる

参考サイト

開発環境

macOS 13.0
Apple M2

チュートリアル要件

https://codelabs.developers.google.com/codelabs/tensorflowjs-teachablemachine-codelab?hl=ja#1

実装

実装内容の詳細については上記の参考サイトをご参照ください

index.html

<html>
  <head>
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs"></script>
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow-models/mobilenet"></script>
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow-models/knn-classifier"></script>
  </head>
  <body>
    <div id="console"></div>
    <video autoplay playsinline muted id="webcam" width="224" height="224"></video>
    <div>
      <button id="class-a">Add A</button>
      <button id="class-b">Add B</button>
      <button id="class-c">Add C</button>
    </div>
    <div>
      <button id="class-s">Data Save</button>
      <button id="class-l">Data Load</button>
      <button id="class-r">Data Remove</button>
    </div>
    <script src="index.js"></script>
  </body>
</html>

index.js

let net;
const webcamElement = document.getElementById('webcam');
const classifier = knnClassifier.create();

async function app() {
    console.log('Loading mobilenet..');
    net = await mobilenet.load();
    console.log('Successfully loaded model');

    const webcam = await tf.data.webcam(webcamElement);
  
    const addExample = async classId => {
        const img = await webcam.capture();
        const activation = net.infer(img, true);
        classifier.addExample(activation, classId);
        img.dispose();
    };

    //localStorageに保存
    function dataSave() {
        let str = JSON.stringify( Object.entries(classifier.getClassifierDataset()).map(([label, data])=>[label, Array.from(data.dataSync()), data.shape]) );
        if(str.length > 2){
            localStorage.setItem("任意のkey", str);
        }
    }

    //localStorageにデータがある場合はロードしてセットする
    function dataLoad() {
        if(localStorage.getItem('任意のkey') != null){
            let str = localStorage.getItem('任意のkey');
            classifier.setClassifierDataset( Object.fromEntries( JSON.parse(str).map(([label, data, shape])=>[label, tf.tensor(data, shape)]) ) );
        }
    }

    //localStorageに保存されているデータを削除する
    function dataRemove(){
        localStorage.removeItem('任意のkey');
    }
    
    document.getElementById('class-a').addEventListener('click', () => addExample(0));
    document.getElementById('class-b').addEventListener('click', () => addExample(1));
    document.getElementById('class-c').addEventListener('click', () => addExample(2));
    //転移学習のデータ保存、ロード、削除についてそれぞれId指定
    document.getElementById('class-s').addEventListener('click', () => dataSave());
    document.getElementById('class-l').addEventListener('click', () => dataLoad());
    document.getElementById('class-r').addEventListener('click', () => dataRemove());

    while (true) {
        if (classifier.getNumClasses() > 0) {
            const img = await webcam.capture();
            const result = await net.classify(img);
            const activation = net.infer(img, 'conv_preds');
            const resultTL = await classifier.predictClass(activation);
            const classes = ['A', 'B', 'C'];
            document.getElementById('console').innerText = `
                prediction: ${result[0].className}\n
                probability: ${result[0].probability}
                prediction TL: ${classes[resultTL.label]}\n
                probability TL: ${resultTL.confidences[resultTL.label]}
            `;
            img.dispose();
        }
        await tf.nextFrame();
    }
}

app();

使い方

カメラの使用を許可する
既に転移学習させたデータが存在し、そのデータを使用する場合は「Data Load」ボタンを押下
学習させたい対象物をカメラに捉え、「Add ...」ボタンを押下
同じAddボタンで上記を数回繰り返す
比較したい対象物をカメラに捉え、別のAddボタンで上記を繰り返す
学習データを保存したい場合は「Data Save」ボタンを押下

上記を行った後に判定させたい対象物をカメラに捉えると、どの学習データに一番近いかを判別して結果を出力します

↓参考に上記ソースコードのデモサイトを作成しました
デモサイト

以上、簡単でしたがTensorFlow.jsを用いた画像AI学習webアプリの例を記述しました
手元で試したところ判別精度はそこそこ高いように思えましたが、いかがだったでしょうか？
説明やソースコード等については最低限の記述のみとなりましたが、ご容赦いただけると幸いです

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up