Problems printing predictions using tensorflow.js

Question

I have retrained a mobilenet_v2 model using the make_image_classifier command line tool to retrain the model and the tfjs-converter to prepare the model for the browser.

make_image_classifier \
  --image_dir image_data \
  --tfhub_module https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4 \
  --saved_model_dir trained_models/1 \
  --labels_output_file class_labels.txt \
  --tflite_output_file trained_model.tflite

tensorflowjs_converter \
    --input_format=tf_saved_model \
    --output_format=tfjs_graph_model \
    --signature_name=serving_default \
    --saved_model_tags=serve \
    ./trained_models/1 \
    ./web_model

For testing the TF Lite model I used the tflite example code. I followed the instructions of the tools and hence used the code provided.

If I now try to predict an image in the browser I don't get the expected output. It looks like only the probabilities are getting printed without the labels.

const MODEL_URL = 'model/model.json';
const model = await tf.loadGraphModel(MODEL_URL);

var canvas = document.getElementById("canvas").getContext("2d");;
const img = canvas.getImageData(0,0, 224,224)

const tfImg = tf.browser.fromPixels(img).expandDims(0);
const smalImg = tf.image.resizeBilinear(tfImg, [224, 224]);

let result = await model.predict(smalImg);

console.log(result.print())

output: Tensor [[0.0022475, 0.0040588, 0.0220788, 0.0032885, 0.000126, 0.0030831, 0.8462179, 0.1188994],]

Testing the model with python works well and I got the expected ouput with labels and probabilities. Am I doing something wrong?

Answer 1

In the python code there is the following after the predition :

  # output_data is the predition result
  results = np.squeeze(output_data)
  top_k = results.argsort()[-5:][::-1]
  labels = load_labels(args.label_file)
  for i in top_k:
    if floating_model:
      print('{:08.6f}: {}'.format(float(results[i]), labels[i]))
    else:
      print('{:08.6f}: {}'.format(float(results[i] / 255.0), labels[i]))

The same processing needs to be done after the prediction in js. The array of labels contains all the labels of the dataset. It needs to be loaded in js as it was loaded in the python code from the labels.txt file.

const topkIndices = result.topk(5) 
const topkIndices = await topk.indices.data();
const categories = (Array.from(topkIndices)).map((p: number) => labels[p]);

const topkValues = await topk.values.data()
const data = Array.from(topkValues).map(p => p * 100);
// categories will contain the top5 labels and
// data will contain their corresponding probabilities