CODEWITHSUNDEEP
javascriptnode.jstensorflowtensorflow.js
Alex
Alex

Reputation: 68406

How to classify image using custom model in tensorflow.js?

The script that builds and trains the model looks like this:

const model = tf.sequential();

model.add(tf.layers.conv2d({
  inputShape: [160, 200, 3],
  filters: 32,
  kernelSize: 3,
  activation: 'relu',
}));

model.add(tf.layers.flatten());

model.add(tf.layers.dense({units: labels.length, activation: 'softmax'}));  

model.compile({
  optimizer: 'sgd',
  loss: 'categoricalCrossentropy',
  metrics: ['accuracy']
});    

const info = await model.fitDataset(ds, {epochs: 5});
console.log('Accuracy:', info.history.acc);

console.log('Saving...');
await model.save('file://' + MODEL_PATH);
console.log('Saved model');

ds consists of images and labels. For 100 images I get these results:

4478ms 135692us/step - acc=0.109 loss=14.37

and it produced a 20 MB weights.bin file...

Frankly, I have no idea if that's good or not because I don't know how to use it classify new images.

I know how to load the model:

const model = await tf.loadLayersModel('file://' + MODEL_PATH + '/model.json');

but that's it.

mobilenet has a .classify method to which I can just pass an image and it outputs the predicted laabel. But this is not available on the model object.. So how do I proceeed?

Upvotes: 2

Views: 1266

Answers (1)

edkeveked
edkeveked

Reputation: 18371

After training your model, to classify a new image, the predict method will be used. It will return the probability of each label given your input image(s).

output = model.predict(image) // It can be a tensor of one image or a batch of many 
// output is a tensor that contain the probability for each label
images.argMax([-1]) // contain the label of high probability for each input

Upvotes: 1

Related Questions