node-red-contrib-tfjs-pj1/tfjs-pj-util.js

custom Node-RED for classification

const tf = require('@tensorflow/tfjs-node');

// Load the Model from the specified URL
const loadModel = async function (modelUrl = 'http://192.168.1.18:8080/model.json') {
    console.log(`Loading model from ${modelUrl}`);
    const model = await tf.loadLayersModel(modelUrl);
    return model;
}

// Preprocess the input image
const processInput = function (imageBuffer) {
    console.log(`Preprocessing image`);
    const uint8array = new Uint8Array(imageBuffer);
    let tensor = tf.node.decodeImage(uint8array, 3);
    
    // Resize the image to the expected input size and rescale pixel values to [0,1].
    tensor = tensor.div(255.0).resizeBilinear([224, 224]).expandDims();
    return tensor;
}

// Process the model's output
const processOutput = function (prediction) {
    console.log('Processing output');
    const scoresArray = prediction[0];
    const maxScoreIndex = scoresArray.indexOf(Math.max(...scoresArray));

    return {
        classIndex: maxScoreIndex,
        label: labels[maxScoreIndex],  // This assumes you've loaded labels array somewhere
        probability: scoresArray[maxScoreIndex]
    };
}

module.exports = {
    loadModel: loadModel,
    processInput: processInput,
    processOutput: processOutput
}