// Based on https://editor.p5js.org/ml5/sketches/FeatureExtractor_Image_Classification

// Image recognition using transfer learning. Uses ml5.js, p5.js, MobileNet. 

​

let featureExtractor;

let classifier;

let video;

​

let nSamplesA = 0;

let nSamplesB = 0;

let nSamplesC = 0;

let bMobileNetIsLoaded = false;

let trainingLoss = NaN;

let currentLabel = "";

let currentConfidence = "";

​

const TRAINING_NOT_STARTED = -1;

const TRAINING_STARTED = 0;

const TRAINING_FINISHED = 1;

let systemStatus = TRAINING_NOT_STARTED;

​

//-----------------------------------------

function setup() {

  createCanvas(320, 240);

​

  video = createCapture(VIDEO);

  video.size(320, 240);

  video.hide(); // hide the video's DOM object.

​

  // Extract the already learned features from MobileNet

  featureExtractor = ml5.featureExtractor('MobileNet',

    theModelReadyCallbackFunction);

​

  // Create a new classifier using those features, 

  // and provide it with the video object

  const options = {

    numLabels: 3

  };

  classifier = featureExtractor.classification(video, options);

}

​

//-----------------------------------------

function draw() {

  background('white');

  tint(255, 255, 255, 65);

  image(video, 0, 0, 320, 240);

​

  drawStatusIndicators();

}

​

//-----------------------------------------

function drawStatusIndicators() {

  fill('black');

​

  text("MobileNet done loading: " + bMobileNetIsLoaded, 15, 20);

​

  text("#A:" + nSamplesA, 15, 35);

  text("#B:" + nSamplesB, 15, 50);

  text("#C:" + nSamplesC, 15, 65);

​

  if (systemStatus === TRAINING_NOT_STARTED) {

    text("TRAINING_NOT_STARTED", 15, 80);

  } else if (systemStatus === TRAINING_STARTED) {

    text("TRAINING_STARTED", 15, 80);

  } else if (systemStatus === TRAINING_FINISHED) {

    text("TRAINING_FINISHED", 15, 80);

  }

​

  if (systemStatus >= 0) {

    text("trainingLoss = " + trainingLoss, 15, 95);

  } else {

    if ((nSamplesA > 0) && (nSamplesB > 0) && (nSamplesC > 0)) {

      text("Press t to start training", 15, 95);

    } else {

      text("Add samples by pressing keys a,b,c", 15, 95);

    }

  }

​

  if (systemStatus === TRAINING_FINISHED) {

    text("Label: " + currentLabel, 15, 110);

    text("Confidence: " + currentConfidence, 15, 125); //nf(currentConfidence,1,6)

  }

}

​

​

//-----------------------------------------

function keyPressed() {

  switch (key) {

    case 'a':

      classifier.addImage('A');

      nSamplesA++;

      break;

    case 'b':

      classifier.addImage('B');

      nSamplesB++;

      break;

    case 'c':

      classifier.addImage('C');

      nSamplesC++;

      break;

​

    case 't':

      classifier.train(function(lossValue) {

        if (lossValue) {

          trainingLoss = lossValue;

          systemStatus = TRAINING_STARTED;

        } else {

          systemStatus = TRAINING_FINISHED;

          classifyCurrentVideoFrame();

        }

      });

      break;

  }

}

​

//-----------------------------------------

// Classify the current frame of video.

function classifyCurrentVideoFrame() {

  classifier.classify(theGotResultsCallbackFunction);

}

​

//-----------------------------------------

// What to do when the model is finished loading.

function theModelReadyCallbackFunction() {

  bMobileNetIsLoaded = true;

}

​

//-----------------------------------------

// What we should do when the classifier sends back results

function theGotResultsCallbackFunction(err, results) {

  if (err) {

    console.error(err); // Display any error

  }

​

  if (results && results[0]) {

    currentLabel = results[0].label;

    currentConfidence = results[0].confidence.toFixed(2);

​

    // Now that we've completed processing yesterframe's results,

    // kick off the processing of the current frame.

    classifyCurrentVideoFrame();

  }

}