ds-faceapi/commands/train.js

A DroidScript FaceAPI plugin trainer

const fs = require("fs")
const path = require("path")
const faceapi = require("face-api.js")
const canvas = require('canvas')
const ora = require("ora")
const chalk = require("chalk")

const { Canvas, Image, ImageData } = canvas
faceapi.env.monkeyPatch({ Canvas, Image, ImageData })

async function init( fast )
{
    const modelsPath = path.join(__dirname, "../", "models")
    await faceapi.nets.faceLandmark68Net.loadFromDisk(modelsPath)
    await faceapi.nets.faceRecognitionNet.loadFromDisk(modelsPath)
    if( fast ) await faceapi.nets.tinyFaceDetector.loadFromDisk(modelsPath)
    else await faceapi.nets.ssdMobilenetv1.loadFromDisk(modelsPath)
}

const train = async function( opt )
{
    let spinner1 = ora({
        text: "Loading face-api and models...",
        color: "green"
    }).start()

    try {
        await init( opt.fast )
        spinner1.succeed("Loading complete")
    }
    catch( err ) {
        spinner1.fail("Loading failed")
        return
    }

    const dir = process.cwd()
    const src = opt.source || "models"
    
    let options = null
    if( opt.fast ) options = new faceapi.TinyFaceDetectorOptions()
    else options = new faceapi.SsdMobilenetv1Options()

    const folderPath = path.join(dir, src)
    const jsonPath = path.join(dir, (opt.name||src)+".json")

    if( fs.existsSync(folderPath) ) {
        let folders = fs.readdirSync(folderPath).filter(file => {
            const fullPath = path.join(folderPath, file)
            return fs.statSync(fullPath).isDirectory()
        })

        let spinner2 = ora({
            text: "Training models from given images...",
            color: "green"
        }).start()

        const trained_models = await Promise.all(
            folders.map(async name => {
                const modelPath = path.join(folderPath, name)
                let images = fs.readdirSync(modelPath)
                images.filter(img => {
                    const ext = path.extname(img).toLowerCase()
                    return ['.jpg', '.jpeg', '.png', '.gif'].includes(ext)
                })
                const descriptions = []
                for (let i = 0; i < images.length; i++) {
                    const imgPath = path.join(modelPath, images[i])
                    const img = new Image()
                    img.src = imgPath
                    const detections = await faceapi.detectSingleFace(img, options).withFaceLandmarks().withFaceDescriptor()
                    descriptions.push(Object.values(detections.descriptor).map(Number))
                }
                return {name, descriptions}
            })
        )
        fs.writeFile(jsonPath, JSON.stringify(trained_models), 'utf8', (err) => {
            if (err) {
                console.error('Error writing file:', err)
                spinner2.fail("Training failed")
            } else {
                // console.log('File written successfully!')
                spinner2.succeed('Training successful. Model is save as "'+src+'.json"')
            }
        })
    }
    else {
        console.log(chalk.red('"' + src +'" folder does not exist.'))
    }
}

module.exports = train