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

async function loadMobilenet() {
  const mobilenet = await tf.loadModel(
      'https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v1_0.25_224/model.json');

  // Return a model that outputs an internal activation.
  const layer = mobilenet.getLayer('conv_pw_13_relu');
  return tf.model({inputs: mobilenet.inputs, outputs: layer.output});
}

mobilenet = loadMobilenet();

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

model = tf.sequential({
  layers: [
    tf.layers.flatten({inputShape: [7, 7, 256]}),
    // Layer 1
    tf.layers.dense({
      units: ui.getDenseUnits(),
      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({
      units: NUM_CLASSES,
      kernelInitializer: 'varianceScaling',
      useBias: false,
      activation: 'softmax'
    })
  ]
});

const optimizer = tf.train.adam(ui.getLearningRate());

model.compile({optimizer: optimizer, loss: 'categoricalCrossentropy'});

const batchSize =
    Math.floor(controllerDataset.xs.shape[0] * ui.getBatchSizeFraction());

// Train the model!
model.fit(controllerDataset.xs, controllerDataset.ys, {
  batchSize,
  epochs: ui.getEpochs(),
  callbacks: {
    onBatchEnd: async (batch, logs) => {
      ui.trainStatus('Loss: ' + logs.loss.toFixed(5));
      await tf.nextFrame();
    }
  }
});

const predictedClass = tf.tidy(() => {

  const img = webcam.capture();

  const activation = mobilenet.predict(img);

  const predictions = model.predict(activation);

  return predictions.as1D().argMax();
});

const classId = (await predictedClass.data())[0];