let classifier;

let video;

let loss;

let imagesOfA = 0;

let imagesOfB = 0;

let imagesOfC = 0;

let classificationResult;

let confidence = 0;

let pressedA = false;

let pressedB = false;

let pressedC = false;

​

function setup() {

  createCanvas(640, 480);

  // Create a video element

  video = createCapture(VIDEO);

  video.size(640, 480);

  video.hide();

  // Append it to the videoContainer DOM element

  //video.parent('videoContainer');

  // Extract the already learned features from MobileNet

  featureExtractor = ml5.featureExtractor('MobileNet', modelReady);

  // Create a new classifier using those features and give the video we want to use

  const options = { numLabels: 3 }; //Specify the number of classes/labels

  classifier = featureExtractor.classification(video, options);

  // Set up the UI buttons

  setupButtons();

  textAlign(CENTER);

  textSize(64);

  fill(0,255,0);

}

​

function draw() {

  background(122);

  image(video, 0, 0);

  if (classificationResult == 'A') {

    text("A", width/2, height/2);

    //rect(100,100,100,100);

  } else if (classificationResult == 'B') {

    text("B", width/2, height/2);

    //ellipse(100,100,100,100);

  } else if (classificationResult == 'C') {

    text("C", width/2, height/2);

  }

  if (pressedA) {

    classifier.addImage('A');

    select('#amountOfAImages').html(imagesOfA++); 

  }

  if (pressedB) {

    classifier.addImage('B');

    select('#amountOfBImages').html(imagesOfB++); 

  }

  if (pressedC) {

    classifier.addImage('C');

    select('#amountOfCImages').html(imagesOfC++); 

  }

}

​

// A function to be called when the model has been loaded

function modelReady() {

  select('#modelStatus').html('Base Model (MobileNet) loaded!');

}

​

​

​

// Classify the current frame.

function classify() {

  classifier.classify(gotResults);

}

​

// A util function to create UI buttons

function setupButtons() {

  buttonA = select('#ButtonA');

  buttonA.mousePressed (() => pressedA = true).mouseReleased(() => pressedA = false);

  buttonB = select('#ButtonB');

  buttonB.mousePressed (() => pressedB = true).mouseReleased(() => pressedB = false);

  buttonC = select('#ButtonC');

  buttonC.mousePressed (() => pressedC = true).mouseReleased(() => pressedC = false);

  // Train Button

  train = select('#train');

  train.mousePressed(function() {

    classifier.train(function(lossValue) {

      if (lossValue) {

        loss = lossValue;

        select('#loss').html('Loss: ' + loss);

      } else {

        select('#loss').html('Done Training! Final Loss: ' + loss);

      }

    });

  });

​

  // Predict Button

  buttonPredict = select('#buttonPredict');

  buttonPredict.mousePressed(classify);

  // Save model

  saveBtn = select('#save');

  saveBtn.mousePressed(function() {

    classifier.save();

  });

​

  // Load model

  loadBtn = select('#load');

  loadBtn.changed(function() {

    classifier.load(loadBtn.elt.files, function(){

      select('#modelStatus').html('Custom Model Loaded!');

    });

  });

}

​

// Show the results

function gotResults(err, result) {

  // Display any error

  if (err) {

    console.error(err);

  }

  select('#result').html(result[0].label);

  select('#confidence').html(result[0].confidence);

​

  classificationResult = result[0].label;

  confidence = result[0].confidence;

​

  classify();

}