Search code examples
javascriptmachine-learningml5.js

ML5 FeatureExtractor addImage not working


I'm following the basic example of the ML5.js featureExtractor. I am not using video. After loading a model I am adding new images to it, and then training again. I get the following error:

Mobilenet.js:323 Uncaught (in promise) Error: Batch size is 0 or NaN. Please choose a non-zero fraction. at t. (Mobilenet.js:323)

My code seems correct, according to the documentation - using video is optional, so I expect I should be able to re-train the model after just adding images manually. I used the callback for the addImage function, to make sure the images are really added before calling train

let added = 0
let classifier
let featureExtractor = ml5.featureExtractor('MobileNet', modelLoaded)
function modelLoaded() {
    classifier = featureExtractor.classification()
    classifier.addImage(document.getElementById('person1'), 'nomask', addedImage)
    classifier.addImage(document.getElementById('mask1'), 'mask', addedImage)
}

// this gets called twice, but then train goes wrong
function addedImage(){
    added++
    if(added == 2){
    classifier.train((lossValue) => {
        console.log('Loss is', lossValue);
    })
}

Solution

  • You need to add at least 3 images for training to work.

    The following code should work.

    let added = 0;
    let classifier;
    let featureExtractor = ml5.featureExtractor('MobileNet', modelLoaded);
    
    function modelLoaded() {
        classifier = featureExtractor.classification()
        classifier.addImage(document.getElementById('person1'), 'nomask', addedImage);
        classifier.addImage(document.getElementById('person2'), 'nomask', addedImage);
        classifier.addImage(document.getElementById('mask1'), 'mask', addedImage);
    }
    
    function addedImage(){
        added++;
        if(added == 3){
        classifier.train((lossValue) => {
            console.log('Loss is', lossValue);
        });
    }
    

    A working example: https://glitch.com/edit/#!/ml5-feature-extractor-addimage