p5(sketch => {
let mobileNet
let classifier
let video
let oneButton
let twoButton
let threeButton
let trainButton
let saveButton
let loadButton
let label = "-"
let loss = "-"
function onTraining(lossValue) {
if (lossValue === null) classifier.classify(video, onPrediction)
else loss = lossValue
}
function onPrediction(error, results) {
if (error) label = error.message
else {
results.sort((x, y) => y.confidence - x.confidence)
label = `${results[0].label} (${results[0].confidence})`
}
classifier.classify(video, onPrediction)
}
sketch.setup = function() {
sketch.createCanvas(640, 480)
sketch.background(0)
video = sketch.createCapture(sketch.VIDEO)
video.hide()
mobileNet = ml5.featureExtractor('MobileNet')
classifier = mobileNet.classification(video, { numLabels: 3 })
oneButton = sketch.createButton("One");
oneButton.mousePressed(() => classifier.addImage(video, 'one'))
twoButton = sketch.createButton("Two");
twoButton.mousePressed(() => classifier.addImage(video, 'two'))
threeButton = sketch.createButton("Three");
threeButton.mousePressed(() => classifier.addImage(video, 'three'))
trainButton = sketch.createButton("Train");
trainButton.mousePressed(() => classifier.train(onTraining))
saveButton = sketch.createButton("Save")
saveButton.mousePressed(() => classifier.save())
loadButton = sketch.createButton("Load")
loadButton.mousePressed(() => {
Promise
.all([modelFile.blob(), weightsFile.blob()])
.then(files => {
classifier.load(files)
classifier.classify(video, onPrediction)
})
})
};
sketch.draw = function() {
sketch.image(video, 0, 0)
sketch.fill(0)
sketch.rect(0, 0, 640, 54)
sketch.fill(255)
sketch.textSize(24)
sketch.text("Loss: " + loss, 10, 24)
sketch.text("Label: " + label, 10, 48)
}
})