classifier/index.js

//get folders of images to use as classes
const glob = require('glob');
const library = require('image-classifier');
const fs = require('fs');
// const 
const path = require('path')
var shouldRebuildModel = false;

process.argv.forEach((val, index, array) => {
    console.log(val)
    if(val === '--rebuild'){
        shouldRebuildModel = true;
    }

})


function getClass(file) {
    return path.dirname(file).split('/').pop();
}
function pGlob(pattern) {
    return new Promise((resolve, reject) => {
        glob(pattern, (error, matches) => {
            if(error) return reject(error);
            resolve(matches);
        })
    })
}
var classifier;

async function runAsync() {

    if(shouldRebuildModel) {
        classifier = await library.create();

        const matches = await pGlob('training images/**/*');
        await Promise.all(matches.map(file => {
            if(fs.lstatSync(file).isDirectory()) {
                console.log('Skipping', file)
                return Promise.resolve();
            }
            const classA = getClass(file);
            return classifier.addExample(classA, file)
        })).catch(error => {
            console.error(error);
        });
        console.log('all examples loaded...');
        console.log('saving dataset...')
        await classifier.save('./export.json');
    } else {
        classifier = await library.load('./export.json');
    }
    const matches = await pGlob('test images/**/*');
    await Promise.all(matches.map(async file => {
        const baseFileName = path.parse(file).name;
        const extension = path.extname(file);
        const newClass= await classifier.predict(file);
        console.log(newClass);
        return new Promise((resolve, reject) => {
            fs.copyFile(file, path.resolve('sorted images', `${baseFileName}-${newClass.label}.${extension}`), resolve);
        })
    }))  
}
runAsync();