Simple Machine Learning Application to Classify Images with TensorFlow and Pre-Trained Dataset | by Ankit Arora | Dec, 2023

Let’s create a simple machine-learning application in Node.js using a popular library called TensorFlow.js to classify images.

TensorFlow.js allows you to build and train machine learning models directly in the browser or Node.js environment.

In this example, we’ll use TensorFlow.js and a pre-trained MobileNet model suitable for image classification tasks.

1. First, install the required npm packages:

npm install @tensorflow/tfjs-node express multer

2. Create a file named animalClassification.js:

const express = require('express');
const multer = require('multer');
const tf = require('@tensorflow/tfjs-node');
const app = express();
const port = 3000;
// Set up Multer for handling file uploads
const storage = multer.memoryStorage();
const upload = multer({ storage: storage });
// Load the MobileNet model
const loadModel = async () => {
const model = await tf.loadLayersModel('https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v1_0.25_224/model.json');
return model;
};
let mobilenetModel;
loadModel().then((model) => {
mobilenetModel = model;
// Start the server
app.listen(port, () => {
console.log(`Server is running on http://localhost:${port}`);
});
});
// Handle image classification requests
app.post('/classify', upload.single('image'), async (req, res) => {
try {
// Decode and preprocess the image
const imgBuffer = req.file.buffer;
const image = tf.node.decodeImage(imgBuffer);
const processedImage = tf.image.resizeBilinear(image, [224, 224]).toFloat().sub(127).div(127).expandDims();
// Make a prediction using the MobileNet model
const predictions = mobilenetModel.predict(processedImage);
const topPredictions = await getTopPredictions(predictions);
// Send the top predictions as JSON response
res.json(topPredictions);
} catch (error) {
console.error(error);
res.status(500).json({ error: 'An error occurred during image classification.' });
}
});
// Helper function to get the top predictions
const getTopPredictions = async (predictions, topK = 3) => {
const values = await predictions.data();
const indices = Array.from(Array(values.length).keys());
const sortedIndices = indices.sort((a, b) => values[b] - values[a]);
const topIndices = sortedIndices.slice(0, topK);
const topPredictions = topIndices.map(index => {
return {
className: `Class ${index}`,
probability: values[index],
};
});
return topPredictions;
};

3. Run the script using:

node animalClassification.js

Use a tool like Postman or CURL to test the image classification. Send a POST request to http://localhost:3000/classify with a form-data parameter named image containing an image file.

This example uses Express for creating a simple web server, Multer for handling file uploads, and TensorFlow.js to load a pre-trained MobileNet model for image classification.

When an image is uploaded, the server will respond with the top predictions for the image.

Note: Ensure that you have a reliable internet connection since the MobileNet model is loaded from an external URL.

Also, keep in mind that this example uses a small model (mobilenet_v1_0.25_224) for demonstration purposes.

For more accurate results, you might want to use a larger model or train a custom model on a specific dataset.

Play the snake game online https://symphonious-pithivier-ceca97.netlify.app/