let numClasses = 7;

​

let featureExtractor;

let classifier;

let video;

let loss;

let img;

let numImages;

let samples;

let label;

let buttons;

​

function setup() {

  createCanvas(500, 500);

​

  img = loadImage("./assets/guitar.png");

​

  samples = [];

  samples[1] = loadSound('./assets/C.mp3')

  samples[2] = loadSound('./assets/F.mp3')

  samples[3] = loadSound('./assets/G.mp3')

  samples[4] = loadSound('./assets/Am.mp3')

  samples[5] = loadSound('./assets/Dm.mp3')

  samples[6] = loadSound('./assets/Em.mp3')

​

  label = 0;

​

  numImages = Array(numClasses).fill(0);

​

  video = createCapture(VIDEO);

  video.parent('videoContainer');

​

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

  featureExtractor.numClasses = numClasses;

  classifier = featureExtractor.classification(video);

  window.c = classifier

​

  createButtons();

​

}

​

function modelReady() {

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

}

​

function addImage(label) {

  classifier.addImage(label);

}

​

function classify() {

  classifier.classify(gotResults);

}

​

function train_classifier() {

  classifier.train(function(lossValue) {

    if (lossValue) {

       loss = lossValue;

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

    } else {

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

    }

  });

}

​

function keyPressed() {

  if ((keyCode >= 48) && (keyCode <= 57)) {

    add_class(keyCode - 48);

  }

}

​

​

  function add_class(i) {

    addImage(i);

    numImages[i]++;

    select('#class' + i).html('class ' + i + ': ' + numImages[i]);

  }

​

  function createButtons() {

    buttons = [];

    for (var i = 0; i < numClasses; i++) {

      var button = createButton('class ' + i);

      button.id('class' + i);

      button.parent('buttons');

      button.mousePressed(add_class.bind(null, i))

      buttons.push(button);

    }

​

    // Train Button

    train = select('#train');

    train.mousePressed(train_classifier);

​

    // Predict Button

    buttonPredict = select('#buttonPredict');

    buttonPredict.mousePressed(classify);

  }

​

  function gotResults(err, nextLabel) {

    if (err) {

      console.error(err);

    }

    select('#result').html(nextLabel);

    if (label != nextLabel) {

      label = nextLabel;

      if (int(label) != 0) {

        samples[int(label)].play();

      }

    }

    classify();

  }

​

  function draw() {

    background(0);

    image(img, 50, 50);

  }