In the artificial intelligence (AI) era, we can see many uses of machine learning implement in our lives today. You may encounter a use case where you need to integrate machine learning with your backend. To integrate machine learning we need some configuration and it’s not surprising if someone struggles to integrate a machine learning model into the backend. In this article, I will share how to integrate a pre-trained machine learning model into a Node.js environment.
Prepare the project
The first step, we will set up the project. In this article, I will use JavaScript and Hapi Framework.
Create a new project and open the project in your text editor/IDE.
Initialize your project with npm init -y
Next, run npm install @hapi/hapi to install the latest version of Hapi. After successfully, create a project structure as shown in the following image.
Open your server.js file and write the code as shown below:
We write a Hapi plugin named predict and PredictService as a machine learning service.
Create a Hapi Plugin as HTTP Server
Open your index.js file as write the code as shown below:
After that, open yourroutes.js
Next, write the code as shown below:
We have a route ‘/predict’ with the method ‘POST’. We can process the image file with multipart/form-data.
And open your handler.js write the code as shown below:
We have a service as a constructor. The next step is to create the service. The handler returns a response with data disease label and confidence score.
So, the next step is to write our service. In the predict service, we integrate the machine learning model.
Load The Machine Learning Model
I don’t cover how to create a machine learning model, so I assume you already have a model. I use a pre-trained model to predict skin disease.
I created a machine learning model with TensorFlow, and I will use TensorFlow.js. To integrate the model into a RESTful API, you need to install some dependencies.
To add dependencies tfjs-node. you can write command as shown below:
npm install @tensorflow/tfjs-node.
After installing the dependencies, we can continue to the next step.
Open your service.js
First step, import the TensorFlow dependencies we installed before.
const tf = require('@tensorflow/tfjs-node')
After that, write a class PredictService and method predictImage.
class PredictService {
async predictImage(photo) {}
}
Inside the predictImage method, write a code to load the machine learning model.
const modelPath = 'file://model/model.json'
const model = await tf.loadLayersModel(modelPath);
We save the model.json in the folder model.
Next, process the image file. So, write an array with name buffers.
const buffers = [];
for await (const data of photo) {
buffers.push(data);
}const image = Buffer.concat(buffers)
We process chunks of data from the photo and then concatenate them into a single Buffer named image.
Next, we can predict the image with the following code:
const tensor = tf.node
.decodeImage(image)
.resizeNearestNeighbor([224, 224])
.expandDims()
.toFloat();const predict = await model.predict(tensor);
const score = await predict.data();
const confidenceScore = Math.max(...score);
const label = tf.argMax(predict, 1).dataSync()[0];
Resize the image with TensorFlow API resizeNearestNeighbor. And make predictions on the input tensor with a pre-trained machine learning model.
Variable score retrieves predicted values using the data() method. The results represent the confidence of the model in different classes.
Variable confidenceScore calculates the maximum confidence score from the array of scores.
Variable labeluses the argMax function to find the index of the maximum score, indicating the predicted class label.
In this case, we have three class labels. So, write an array of the class.
const diseaseLabels = ['Melanocytic nevus', 'Squamous cell carcinoma', 'Vascular lesion'];
const diseaseLabel = diseaseLabels[label];return { confidenceScore, diseaseLabel };
Variable diseaseLabel to return the name of label we predict. So, if label is 1, the disease label will be ‘Squamos cell carcinoma’.
And last, return the confidence score and disease label.
Here is the full code of service.js.
Testing
Last step, we test the REST API we made before. You can use Postman/Katalon/etc.
To test the REST API, make a request to your endpoint ‘/predict’ and make sure the endpoint returns a data confidence score and disease name.
We succeeded in integrating the machine learning model into the REST API. We usea pre-trained machine learning model and use TensorFlow.js to integrate into the REST API.
Want to connect?
https://www.linkedin.com/in/arasopraza/
https://github.com/arasopraza
https://twitter.com/arsyopraza
Be the first to comment