4_transfer-learning.html

<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@1.1.2"></script>
<script src="https://cdn.jsdelivr.net/npm/@tensorflow-models/mobilenet@1.0.1"></script>

<script src="https://cdn.jsdelivr.net/npm/@tensorflow-models/knn-classifier@1.1.0"></script>

<script src="../webcam.js"></script>

<video autoplay playsinline muted id="webcam" width="640" height="480"></video>
<h1 id="description"></h1>

<div>
    <input id="label" type="text">
    <button id="train">Train</button>
</div>


<script>
    let model;
    let webcamElement;
    let classifier;
    const classLabels = new Set();

    window.onload = async () => {
        const trainButton = document.getElementById('train');
        const description = document.getElementById('description');
        webcamElement = await setupWebcam(document.getElementById('webcam'));

        model = await mobilenet.load();
        classifier = knnClassifier.create();

        trainButton.addEventListener('click', addTrainingExample);

        while (true) {
            if (classifier.getNumClasses() > 0) {
                const activation = model.infer(webcamElement, 'conv_preds');
                const prediction = await classifier.predictClass(activation);

                if (prediction.confidences[prediction.classIndex] > 0.5) {
                    description.innerText = [...classLabels][prediction.classIndex];
                } else {
                    description.innerText = "Not sure"
                }
            } else {
                description.innerText = "Nothing trained yet"
            }

            await tf.nextFrame();
        }
    };

    function addTrainingExample() {
        const label = document.getElementById('label').value.toUpperCase();
        if (!label) return;

        classLabels.add(label);
        const classId = [...classLabels].indexOf(label);

        const activation = model.infer(webcamElement, 'conv_preds');
        classifier.addExample(activation, classId);

        console.log('added training example for ' + label)
    }
</script>