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

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

​

const FONT_SIZE = 24;

const TOP_MARGIN = 50;

const LEFT_MARGIN = 300;

const GRID_SIZE = 500;

​

let textGraphic;

let canvasGraphic;

let canClassify = false;

​

let nSamplesA = 0;

let nSamplesB = 0;

let nSamplesC = 0;

let featureExtractor;

let classifier;

let video;

​

let bMobileNetIsLoaded = false;

let trainingLoss = NaN;

let currentLabel = "A";

let currentConfidence = "";

​

const TRAINING_NOT_STARTED = -1;

const TRAINING_STARTED = 0;

const TRAINING_FINISHED = 1;

let systemStatus = TRAINING_NOT_STARTED;

​

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

function setup() {

  createCanvas(800, 600);

​

  textGraphic = createGraphics(800, 600);

  canvasGraphic = createGraphics(GRID_SIZE, GRID_SIZE);

​

  canvasGraphic.stroke('black');

  canvasGraphic.strokeWeight(2);

  canvasGraphic.rect(0, 0, GRID_SIZE, GRID_SIZE);

​

  // Extract the already learned features from MobileNet

  const options = {

    numLabels: 3

  };

  featureExtractor = ml5.featureExtractor('MobileNet', options, theModelReadyCallbackFunction);

​

  // Create a new classifier using those features, 

  // and provide it with the video object

  classifier = featureExtractor.classification();

}

​

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

function draw() {

  clear();

  drawOutline();

  drawCanvas();

}

​

function drawCanvas() {

  if (mouseIsPressed &&

    mouseX >= LEFT_MARGIN && mouseX <= LEFT_MARGIN + GRID_SIZE &&

    mouseY >= TOP_MARGIN && mouseY <= TOP_MARGIN + GRID_SIZE &&

    pmouseX >= LEFT_MARGIN && pmouseX <= LEFT_MARGIN + GRID_SIZE &&

    pmouseY >= TOP_MARGIN && pmouseY <= TOP_MARGIN + GRID_SIZE) {

    canvasGraphic.strokeWeight(5);

    canvasGraphic.stroke('black');

    canvasGraphic.line(mouseX - LEFT_MARGIN, mouseY - TOP_MARGIN, pmouseX - LEFT_MARGIN, pmouseY - TOP_MARGIN);

  }

​

  image(canvasGraphic, LEFT_MARGIN, TOP_MARGIN);

​

  if (canClassify) {

    classifyCurrentFrame();

  }

}

​

function drawOutline() {

  textGraphic.clear();

​

  //drawStatusIndicators();

​

  textGraphic.stroke(color(0, 0, 0, currentConfidence * 25));

  textGraphic.noStroke();

  textGraphic.textSize(FONT_SIZE);

  textGraphic.textAlign(CENTER);

​

  if (currentLabel != 'A') {

    var textToShow = currentLabel == 'B' ? 'bird' : 'plane';

    textGraphic.text(textToShow, LEFT_MARGIN + GRID_SIZE / 2, 10 + FONT_SIZE);

    textGraphic.text('It\'s a', LEFT_MARGIN / 2, TOP_MARGIN + GRID_SIZE / 2);

  }

​

  image(textGraphic, 0, 0);

}

​

function getCanvas() {

  let image = document.createElement("img");

  image.src = canvas.toDataURL();

  return image;

}

​

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

function drawStatusIndicators() {

  textGraphic.clear();

  textGraphic.noStroke();

  textGraphic.textSize(10);

  textGraphic.textAlign(LEFT);

​

  textGraphic.fill('black');

​

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

​

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

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

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

​

  if (systemStatus === TRAINING_NOT_STARTED) {

    textGraphic.text("TRAINING_NOT_STARTED", 15, 80);

  } else if (systemStatus === TRAINING_STARTED) {

    textGraphic.text("TRAINING_STARTED", 15, 80);

  } else if (systemStatus === TRAINING_FINISHED) {

    textGraphic.text("TRAINING_FINISHED", 15, 80);

  }

​

  if (systemStatus >= 0) {

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

  }

​

  if (systemStatus === TRAINING_FINISHED) {

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

    textGraphic.text("Confidence: " + currentConfidence, 15, 125);

  }

}

​

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

function keyPressed() {

  switch (key) {

    case 'r':

      canvasGraphic.clear();

      canvasGraphic.stroke('black');

      canvasGraphic.strokeWeight(2);

      canvasGraphic.rect(0, 0, GRID_SIZE, GRID_SIZE);

      break;

    case 't':

      classifier.train(function(lossValue) {

        if (lossValue) {

          trainingLoss = lossValue;

          systemStatus = TRAINING_STARTED;

          console.log("started training");

        } else {

          systemStatus = TRAINING_FINISHED;

          console.log("stopped training");

          classifyCurrentFrame();

        }

​

      });

      break;

    case 'a':

      classifier.addImage(getCanvas(), 'A', function() {

        console.log("A");

      });

      nSamplesA++;

      break;

    case 'b':

      classifier.addImage(getCanvas(), 'B', function() {

        console.log("B");

      });

      nSamplesB++;

      break;

    case 'c':

      classifier.addImage(getCanvas(), 'C', function() {

        console.log("C");

      });

      nSamplesC++;

  }

}

​

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

// Classify the current frame of video.

function classifyCurrentFrame() {

  canClassify = false;

  classifier.classify(getCanvas(), 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);

​

    console.log(currentLabel + " " + currentConfidence);

​

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

    // kick off the processing of the current frame.

    canClassify = true;

  }

}