xxxxxxxxxx
130
// 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 and four classes
=== */
let featureExtractor;
let classifier;
let video;
let loss;
let imagesOfA = 0;
let imagesOfB = 0;
let imagesOfC = 0;
let imagesOfD = 0;
let classificationResult;
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
classifier = featureExtractor.classification(video, videoReady);
featureExtractor.numClasses = 4;
// Set up the UI buttons
setupButtons();
}
function draw() {
background(122);
image(video, 0, 0);
if (classificationResult == 'A') {
ellipse(100, 100, 100, 100);
} else if (classificationResult == 'B') {
rect(100, 100, 100, 100);
}
}
// A function to be called when the model has been loaded
function modelReady() {
select('#modelStatus').html('Base Model (MobileNet) loaded!');
}
// A function to be called when the video has loaded
function videoReady() {
select('#videoStatus').html('Video ready!');
}
// Classify the current frame.
function classify() {
classifier.classify(gotResults);
}
// A util function to create UI buttons
function setupButtons() {
// When the A button is pressed, add the current frame
// from the video with a label of "A" to the classifier
buttonA = select('#ButtonA');
buttonA.mousePressed(function() {
classifier.addImage('A');
select('#amountOfAImages').html(imagesOfA++);
});
// When the B button is pressed, add the current frame
// from the video with a label of "B" to the classifier
buttonB = select('#ButtonB');
buttonB.mousePressed(function() {
classifier.addImage('B');
select('#amountOfBImages').html(imagesOfB++);
});
// When the C button is pressed, add the current frame
// from the video with a label of "C" to the classifier
buttonC = select('#ButtonC');
buttonC.mousePressed(function() {
classifier.addImage('C');
select('#amountOfCImages').html(imagesOfC++);
});
// When the D button is pressed, add the current frame
// from the video with a label of "D" to the classifier
buttonD = select('#ButtonD');
buttonD.mousePressed(function() {
classifier.addImage('D');
select('#amountOfDImages').html(imagesOfD++);
});
// 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);
}
// Show the results
function gotResults(err, result) {
// Display any error
if (err) {
console.error(err);
}
select('#result').html(result);
classificationResult = result;
classify();
}