ml5 Example Feature Extractor Image Classification

// Copyright (c) 2018 ml5

//

// This software is released under the MIT License.

// https://opensource.org/licenses/MIT

​

/* ===

ml5 Example

Image Classification using Feature Extraction with MobileNet. Built with p5.js

=== */

​

/*

Instructions

1. Make a pose in front of your webcam and click the 'Add Cat Image' button. Try adding around 15 images. 

2. Make a different pose in front of your webcam and click the 'Add Dog Image' button. Try adding around 15 images. 

3. Click 'Train' and wait until the training process is done. 

4. Click 'Start guessing!' and switch between the two poses you trained the model on! 

*/

​

let featureExtractor;

let classifier;

let video;

let loss;

let dogImages = 0;

let catImages = 0;

​

let catImagesP;

let dogImagesP;

let lossP;

let resultP;

​

function setup() {

  createCanvas(windowWidth, windowHeight);

  // Create a video element

  video = createCapture(VIDEO);

  video.hide();

  imageMode(CENTER);

  // 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

  classifier = featureExtractor.classification(video);

  // Create the UI buttons

  createButtons();

}

​

function draw() {

  image(video, width/2, height/3, 340, 280);

}

​

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

function modelReady() {

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

  print('Base Model (MobileNet) loaded!');

}

​

// Add the current frame from the video to the classifier

function addImage(label) {

  classifier.addImage(label);

}

​

// Classify the current frame.

function classify() {

  classifier.classify(gotResults);

}

​

// A util function to create UI buttons

function createButtons() {

  // When the Cat button is pressed, add the current frame

  // from the video with a label of "cat" to the classifier

  buttonA = createButton('Add Cat Image');

  buttonA.position(width/3, height-height/3-50);

  catImagesP = createP(catImages+" Cat Images");

  catImagesP.position(width/3+150, height-height/3-65);

  buttonA.mousePressed(function() {

    addImage('cat');

    catImages++;

    catImagesP.html(catImages+" Cat Images");

  });

​

  // When the Dog button is pressed, add the current frame

  // from the video with a label of "dog" to the classifier

  buttonB = createButton('Add Dog Image');

  buttonB.position(width/3, height-height/3);

  dogImagesP = createP(dogImages+" Dog Images");

  dogImagesP.position(width/3+150, height-height/3-15);

  buttonB.mousePressed(function() {

    addImage('dog');

    dogImages++;

    dogImagesP.html(dogImages+" Dog Images");

  });

​

  // Train Button

  train = createButton('Train');

  train.position(width/3, height-height/3+50);

  lossP = createP('');

  lossP.position(width/3+150, height-height/3+35);

  train.mousePressed(function() {

    classifier.train(function(lossValue) {

      if (lossValue) {

        loss = lossValue;

        lossP.html('Loss: ' + loss);

      } 

      else {

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

      }

    });

  });

​

  // Predict Button

  buttonPredict = createButton('Start guessing!');

  buttonPredict.position(width/3, height-height/3+100);

  buttonPredict.mousePressed(classify);

  resultP = createP('Your custom model labeled this as: ...')

  resultP.position(width/3, height-height/3+150);

}

​

// Show the results

function gotResults(result) {

  resultP.html('Your custom model labeled this as: '+result);

  classify();

}