We are using TensorFlow.js to create and train a custom model. We use tf.browser.fromPixels() function to convert an image into tensor. We want to create a custom and train a custom model. In terms to achieve this we are created two different web pages (i.e: 1st page is for Create custom mode and train it with images and its associated labels, 2nd page is used to load trained model and by using pre-trained model we are trying to predict an image to get the associate label with this image) to achieve this functionality we are looked into the below properties:
- AddImage( HTML_Image_Element, 'Label'): - Add an imageElement with a label. Let's say we have three images for prediction i.e: img1, img2, img3 with three labels 'A', 'B' and 'C' respectively.
Train() / fit(): - Train this custom model with associated labels. So we want to create and train our model with these images and respective labels.
Save/Load model:- For making a pre-trained model as reusable, we want to save the trained model and the load this model whenever we want to prediction with the same dataset. By save() function we are getting two files i.e: 'model.json' and 'model.weights.bin'.
- Predict():- Once a successfully trained model got loaded, then we are able to load this model in another page so that user can Predict the images with their associated label, and it will return the predicted response with the attached label of every image. Let's say whenever a user wants to predict 'img1' then it shows the prediction as classname 'A', similarly, for 'img2' predicts with classname 'B' and for 'img3' predicts with classname 'C' with a confidence value as well.
Steps which are achieved: In the above mention requirement we are successfully implemented these following points:
In the very first web page, we are create a sequential model, and add images(By Convert it into a tensor) into the model. and after that, we train/fit this model with these images and associate Labels.
After training, we can easily Save() this custom pre-trained model for further predictions. with this meantime, a user can predict any particular image from the dataset with was used during training, the model gives the response with associated Label i.e: if a user wants the prediction for 'img1' the model response the prediction with the label as 'A'.
After saving the model, we can now second web page where a user can do prediction with images without doing any training by using a pre-trained model. in this phase, we are able to load our saved model and get the prediction in this manner:
prediction:::0.9590839743614197,0.0006004410679452121,0.002040663966909051,0.001962134148925543,0.008351234719157219,0.004203603137284517,0.010159854777157307,0.007813011296093464,0.0013025108492001891,0.004482310265302658
Still need to achieve: While training the custom model, we are adding images and Label both at once in our custom model. But when we save the model.json file then in this .json file we don't able to find the Labels which we associate with the images('A', 'B' and 'C') while add and train model. So when we are add this model.json in second page and try to predict then the model does't shows the associate label. We are did not able to add Labels in the model(model.json) while training/fit. Please find the code and attached web page screenshot for better understanding.
Following are some attachments/sample_code which helps to understand the requirement : Prediction Page(2nd WebPage) : Prediction Page With Pre-trained model
Find here both the model files(.json and .bin) : Custom model files
below are the code of both the pages :
//2nd Page which is for prediction -
<apex:page sidebar="false" >
<head>
<title>Predict with tensorflowJS</title>
<script src="https://ajax.googleapis.com/ajax/libs/jquery/3.3.1/jquery.min.js"></script>
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/[email protected]"> </script>
</head>
<div class="container mt-5">
<div class="row">
<div class="col-12">
<div class="progress progress-bar progress-bar-striped progress-bar-animated mb-2">Loading Model</div>
</div>
</div>
<div class="row">
<div class="col-3">
<select id="model-selector" class="custom-select" >
<option>mobilenet</option>
</select>
</div>
</div>
<input type="file" id="load" multiple="multiple" /><br/>
<label for="avatar">Load Model:</label>
<div class="row">
<input id ="image-selector" class="form-control border-0" type="file"/>
</div>
<div class="col-6">
<button id="predict-button" class="btn btn-dark float-right">Predict</button>
</div>
</div>
<div class="row">
<div class="col">
<h2>Prediction</h2>
<ol id="prediction-list"></ol>
</div>
</div>
<div class="row">
<div class="col-12">
<h2 class="ml-3">Image</h2>
<img id="selected-image" class="ml-3" src="" crossorigin="anonymous" width="400" height="300"/>
</div>
</div>
<script>
$(document).ready()
{
$('.progress-bar').hide();
}
$("#image-selector").change(function(){
let reader = new FileReader();
reader.onload = function(){
let dataURL = reader.result;
$("#selected-image").attr("src",dataURL);
$("#prediction-list").empty();
}
let file = $("#image-selector").prop('files')[0];
reader.readAsDataURL(file);
});
$("#model-selector").ready(function(){
loadModel($("#model-selector").val());
$('.progress-bar').show();
})
let model;
let cutomModelJson;
let cutomModelbin;
async function loadModel(name){
$("#load").change(async function(){
for (var i = 0; i < $(this).get(0).files.length; ++i) {
console.log('AllFiles:::'+JSON.stringify($(this).get(0).files[i]));
if($(this).get(0).files[i].name == 'my-model-1.json'){
cutomModelJson = $(this).get(0).files[i];
}else{
cutomModelbin = $(this).get(0).files[i];
}
}
console.log('cutomModelJson::'+cutomModelJson.name+'cutomModelbin::'+cutomModelbin.name);
model = await tf.loadModel(tf.io.browserFiles([cutomModelJson, cutomModelbin]));
console.log('model'+JSON.stringify(model));
});
}
$("#predict-button").click(async function(){
let image= $('#selected-image').get(0);
console.log('image',image);
let tensor = preprocessImage(image,$("#model-selector").val());
const resize_image = tf.reshape(tensor, [1, 224, 224, 3],'resize');
console.log('tensor',tensor);
console.log('resize_image',resize_image);
console.log('model1',model);
let prediction = await model.predict(tensor).data();
console.log('prediction:::'+ prediction);
let top5 = Array.from(prediction)
.map(function(p,i){
return {
probability: p,
className: prediction[i]
};
}).sort(function(a,b){
return b.probability-a.probability;
}).slice(0,1);
$("#prediction-list").empty();
top5.forEach(function(p){
$("#prediction-list").append(`<li>${p.className}:${p.probability.toFixed(6)}</li>`);
});
});
function preprocessImage(image,modelName)
{
let tensor=tf.browser.fromPixels(image)
.resizeNearestNeighbor([224,224])
.toFloat();
console.log('tensor pro', tensor);
if(modelName==undefined)
{
return tensor.expandDims();
}
if(modelName=="mobilenet")
{
let offset=tf.scalar(127.5);
console.log('offset',offset);
return tensor.sub(offset)
.div(offset)
.expandDims();
}
else
{
throw new Error("UnKnown Model error");
}
}
</script>
//1st Page which is for Create and train model -
<apex:page sidebar="false">
<head>
<title>Add image and train model with tensorflowJS</title>
<script src="https://ajax.googleapis.com/ajax/libs/jquery/3.3.1/jquery.min.js"></script>
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/[email protected]"> </script>
<script src="https://unpkg.com/@tensorflow-models/mobilenet"></script>
</head>
<div class="container mt-5">
<div class="row">
<div class="col-12">
<div class="progress progress-bar progress-bar-striped progress-bar-animated mb-2">Loading Model</div>
</div>
</div>
<div class="row">
<div class="col-3">
<select id="model-selector" class="custom-select" >
<option>mobilenet</option>
</select>
</div>
</div>
<div class="row">
<input id ="image-selector" class="form-control border-0" type="file"/>
</div>
<div class="col-6">
<button id="predict-button" class="btn btn-dark float-right">Predict</button>
</div>
</div>
<div class="row">
<div class="col">
<h2>Prediction></h2>
<ol id="prediction-list"></ol>
</div>
</div>
<div class="row">
<div class="col-12">
<h2 class="ml-3">Image</h2>
<img id="selected-image" src="{!$Resource.cat}" crossorigin="anonymous" width="400" height="300" />
</div>
</div>
<script>
$("#model-selector").ready(function(){
loadModel($("#model-selector").val());
})
let model;
async function loadModel(name){
model = tf.sequential();
console.log('model::'+JSON.stringify(model));
}
$("#predict-button").click(async function(){
let image= $('#selected-image').get(0);
console.log('image:::',image);
let tensor = preprocessImage(image,$("#model-selector").val());
const resize_image = tf.reshape(tensor, [1, 224, 224, 3],'resize');
console.log('tensorFromImage:::',resize_image);
// Labels
const label = ['cat'];
const setLabel = Array.from(new Set(label));
const ys = tf.oneHot(tf.tensor1d(label.map((a) => setLabel.findIndex(e => e === a)), 'int32'), 10)
console.log('ys:::'+ys);
model.add(tf.layers.conv2d({
inputShape: [224, 224 , 3],
kernelSize: 5,
filters: 8,
strides: 1,
activation: 'relu',
kernelInitializer: 'VarianceScaling'
}));
model.add(tf.layers.maxPooling2d({poolSize: 2, strides: 2}));
model.add(tf.layers.maxPooling2d({poolSize: 2, strides: 2}));
model.add(tf.layers.flatten({}));
model.add(tf.layers.dense({units: 64, activation: 'relu'}));
model.add(tf.layers.dense({units: 10, activation: 'softmax'}));
model.compile({
loss: 'meanSquaredError',
optimizer : 'sgd'
})
// Train the model using the data.
model.fit(resize_image, ys, {epochs: 100}).then((loss) => {
const t = model.predict(resize_image);
console.log('Prediction:::'+t);
pred = t.argMax(1).dataSync(); // get the class of highest probability
const labelsPred = Array.from(pred).map(e => setLabel[e])
console.log('labelsPred:::'+labelsPred);
const saveResults = model.save('downloads://my-model-1');
console.log(saveResults);
}).catch((e) => {
console.log(e.message);
})
});
function preprocessImage(image,modelName)
{
let tensor = tf.browser.fromPixels(image)
.resizeNearestNeighbor([224,224])
.toFloat();
console.log('tensor pro:::', tensor);
if(modelName==undefined)
{
return tensor.expandDims();
}
if(modelName=="mobilenet")
{
let offset=tf.scalar(127.5);
console.log('offset:::',offset);
return tensor.sub(offset)
.div(offset)
.expandDims();
}
else
{
throw new Error("UnKnown Model error");
}
}
</script>
Please let us know if we are going wrong any where or need to implement any additional step to achieve this task.