export default () => {
    const [model, setModel] = useState<LayersModel>();
    useEffect(() => {
        (async function init() {
            const model = await tf.loadLayersModel(MODEL_URL);
            setModel(() => model);
            model.summary();
        })();
    }, []);
    return model;
}

// freeze all layers of MobileNet
for (const layer of pretrainedModel.layers) {
    layer.trainable = false;
}

const truncatedLayer = pretrainedModel.getLayer('conv_pw_13_relu');
const truncatedLayerOutput = truncatedLayer.output as SymbolicTensor;

function buildNewHead(inputShape: Shape, numClasses: number) {
    // Creates a 2-layer fully connected model
    return tf.sequential({
        layers: [
            tf.layers.flatten({
                name: 'flatten',
                inputShape: inputShape.slice(1),
            }),
            tf.layers.dense({
                name: 'hidden_dense_1',
                units: 100,
                activation: 'relu'
            }),
            // Layer 2. The number of units of the last layer should
            // correspond to the number of classes we want to predict.
            tf.layers.dense({
                name: 'softmax_classification',
                units: numClasses,
                activation: 'softmax'
            })
        ]
    });
}

const transferHead = buildNewHead(truncatedLayerOutput.shape, numClasses);
const newOutput = transferHead.apply(truncatedLayerOutput) as SymbolicTensor;

// build new model
const model = tf.model({
    inputs: pretrainedModel.inputs,
    outputs: newOutput
});

model.compile({
    optimizer: tfC.train.adam(0.0001),
    loss: 'categoricalCrossentropy'
});

import {useEffect, useState} from 'react';
import {LayersModel, Shape, SymbolicTensor} from '@tensorflow/tfjs';
import * as tf from '@tensorflow/tfjs';
function buildNewHead(
      inputShape: Shape,
      numClasses: number): LayersModel {
    // Creates a 2-layer fully connected model
    return tf.sequential({
        layers: [
            tf.layers.flatten({
                name: 'flatten',
                inputShape: inputShape.slice(1),
            }),
            tf.layers.dense({
                name: 'hidden_dense_1',
                units: 100,
                activation: 'relu',
                kernelInitializer: 'varianceScaling',
                useBias: true
            }),
            // Layer 2. The number of units of the last layer
            // should correspond to the number of classes
            // we want to predict.
            tf.layers.dense({
                name: 'softmax_classification',
                units: numClasses,
                kernelInitializer: 'varianceScaling',
                useBias: false,
                activation: 'softmax'
            })
        ]
    });
}
export default (
    pretrainedModel: LayersModel | undefined,
    numClasses: number) => {
    const [model, setModel] = useState<LayersModel>();
    useEffect(() => {
        if (pretrainedModel) {
            // find last convolutional layer by its name
            // it is a layer followed by layers which are responsible
            // for classification
            const truncatedLayer =
                  pretrainedModel.getLayer('conv_pw_13_relu');
            const truncatedLayerOutput =
                  truncatedLayer.output as SymbolicTensor;
            // freeze all layers of MobileNet
            for (const layer of pretrainedModel.layers) {
                layer.trainable = false;
            }
            const transferHead = buildNewHead(
                  truncatedLayerOutput.shape, numClasses);
            const newOutput =
                  transferHead.apply(truncatedLayerOutput) as SymbolicTensor;
            const model = tf.model({
                inputs: pretrainedModel.inputs,
                outputs: newOutput
            });
            model.compile({
                optimizer: tf.train.adam(0.0001),
                loss: 'categoricalCrossentropy'
            });
            setModel(model);
        }
    }, [pretrainedModel]);
    return model;
}

function App() {
    // skipped for brevity
    const train = async () => {
        if (model
            && controllerDataset.current.xs !== null
            && controllerDataset.current.ys !== null) {
            const {current: {xs, ys}} = controllerDataset;
            model.fit(xs, ys, {
                batchSize: 24,
                epochs: 20,
                shuffle: true,
                callbacks: {
                    onBatchEnd: async (batch, logs) => {
                        setLoss(() => (logs?.loss as number).toFixed(5));
                    },
                    onTrainEnd: () => {
                        changeTrainingState(() => TRAINING_STATES.TRAINED);
                    }
                }
            });
        }
    };
    return (
        <div>
            {/* skipped for brevity*/}
             <button
                  disabled={trainingState === TRAINING_STATES.TRAINING}
                  onClick={() => {
                      if (trainingState !== TRAINING_STATES.TRAINING) {
                         changeTrainingState(() => TRAINING_STATES.TRAINING);
                         train();
                       }
              }}>
                 Train
              </button>
        </div>
    );
}

const predict = async () => {
    if (webcamIterator && model) {
        const image = await webcamIterator.capture();
        const processedImage = tf.tidy<Tensor4D>(
            () => image.expandDims().toFloat().div(127).sub(1)
        );
        const label =
               tf.argMax(model.predict(processedImage) as Tensor, 1)
                    .dataSync();
        setActiveLabel(() => label[0]);
        image.dispose();
        processedImage.dispose();
        requestAnimationFrame(() => predict());
    }
};
useEffect(() => {
    if (trainingState === TRAINING_STATES.TRAINED) {
        requestAnimationFrame(() => predict());
    }
}, [trainingState]);

import React, {useEffect, useRef, useState} from 'react';
import useUploadModel from './hook/useUploadModel';
import useWebcamIterator from './hook/useWebcamIterator';
import WorkingArea from './components/WorkingArea';
import useNewModelTopology from './hook/useNewModelTopology';
import ControllerDataset from './service/ControllerDataset';
import * as tf from '@tensorflow/tfjs';
import {Tensor4D, Tensor} from '@tensorflow/tfjs';
import {Container, Button, Grid, Box, TextField} from '@material-ui/core';
const NUM_CLASSES = 3;
enum TRAINING_STATES {
    NOT_TRAINED,
    TRAINING,
    TRAINED
}
function App() {
    const uploadedModel = useUploadModel();
    const model = useNewModelTopology(uploadedModel, NUM_CLASSES);
    const {videoRef, webcamIterator} = useWebcamIterator();
    const [loss, setLoss] = useState<string | null>(null);
    const controllerDataset =
               useRef(new ControllerDataset(NUM_CLASSES));
    const [trainingState, changeTrainingState] =
               useState<TRAINING_STATES>(TRAINING_STATES.NOT_TRAINED);
    const [activeLabel, setActiveLabel] = useState<number>();
    const train = async () => {
        if (model
            && controllerDataset.current.xs !== null
            && controllerDataset.current.ys !== null) {
            const {current: {xs, ys}} = controllerDataset;
            model.fit(xs, ys, {
                batchSize: 24,
                epochs: 20,
                shuffle: true,
                callbacks: {
                    onBatchEnd: async (batch, logs) => {
                        setLoss(() => (logs?.loss as number).toFixed(5));
                    },
                    onTrainEnd: () => {
                        changeTrainingState(
                              () => TRAINING_STATES.TRAINED
                        );
                    }
                }
            });
        }
    };
    const predict = async () => {
        if (webcamIterator && model) {
            const image = await webcamIterator.capture();
            const processedImage = tf.tidy<Tensor4D>(
                () => image.expandDims().toFloat().div(127).sub(1)
            );
            const label =
                  tf.argMax(model.predict(processedImage) as Tensor, 1)
                       .dataSync();
            setActiveLabel(() => label[0]);
            image.dispose();
            processedImage.dispose();
            requestAnimationFrame(() => predict());
        }
    };
    useEffect(() => {
        if (trainingState === TRAINING_STATES.TRAINED) {
            requestAnimationFrame(() => predict());
        }
    }, [trainingState]);
    return (
        <>
            <video ref={videoRef} width={224} height={224}/>
            <Button
                color="primary"
                variant="contained"
                disabled={trainingState === TRAINING_STATES.TRAINING}
                onClick={() => {
                    if (trainingState !== TRAINING_STATES.TRAINING) {
                        changeTrainingState(
                             () => TRAINING_STATES.TRAINING
                        );
                        train();
                    }
                }}>
                Train
            </Button>
            {loss && <div>Loss: {loss}</div>}
            <WorkingArea numClasses={NUM_CLASSES}
                         webcam={webcamIterator}
                         activeLabel={activeLabel}
                         controllerDataset={controllerDataset.current}
            />
        </>
    );
}