TensorFlow Lite Model Android: Cannot find an axis to label. A valid axis to label should have size larger than 1

I am trying to use a pre-trained TensorFlow Lite model in an Android Application.

I have downloaded the Image Classification example application for TensorFlow Lite from here

I have changed following code in the all the four Model Classifier Files

protected String getModelPath() {
   // you can download this file from
   // see build.gradle for where to obtain this file. It should be auto
   // downloaded into assets.
   //return "mobilenet_v1_1.0_224_quant.tflite";
  return "model_23072020.tflite";
}

The TensorFlow Lite Model that I am using is a pre-trained model for image classification. Basically it scans the image and produces an output as 0 or 1. 0 indicates the image is not of good quality and 1 indicates the image is a good quality image.

Model has Dynamic Range Quantization.

When I am running the application and printing the value of outputProbabilityBuffer.getFloatArray(), I am getting following results

> I/tensorflow: Classifier: value of output [F@e08d0d3

I am trying to log value using below code

tflite.run(inputImageBuffer.getBuffer(), outputProbabilityBuffer.getBuffer().rewind());


Map<String, Float> labeledProbability = new HashMap<>();
labeledProbability.put("abc", 93.556f);

// Added logger for displaying value in console
LOGGER.i("value of output %s ", outputProbabilityBuffer.getFloatArray());

Update

I removed the logger above and now I am getting exception at this line

Map<String, Float> labeledProbability = new TensorLabel(labels, probabilityProcessor.process(outputProbabilityBuffer))
        .getMapWithFloatValue();

The error received is Cannot find an axis to label. A valid axis to label should have size larger than 1.

Complete stack trace is mentioned below

> java.lang.IllegalArgumentException: Cannot find an axis to label. A
> valid axis to label should have size larger than 1. at
> org.tensorflow.lite.support.label.TensorLabel.getFirstAxisWithSizeGreaterThanOne(TensorLabel.java:214)
> at
> org.tensorflow.lite.support.label.TensorLabel.<init>(TensorLabel.java:105)
> at
> org.tensorflow.lite.examples.classification.tflite.Classifier.recognizeImage(Classifier.java:263)
> at
> org.tensorflow.lite.examples.classification.ClassifierTest.classificationResultsShouldNotChange(ClassifierTest.java:67)
> at java.lang.reflect.Method.invoke(Native Method) at
> org.junit.runners.model.FrameworkMethod$1.runReflectiveCall(FrameworkMethod.java:50)
> at
> org.junit.internal.runners.model.ReflectiveCallable.run(ReflectiveCallable.java:12)
> at
> org.junit.runners.model.FrameworkMethod.invokeExplosively(FrameworkMethod.java:47)
> at
> org.junit.internal.runners.statements.InvokeMethod.evaluate(InvokeMethod.java:17)
> at
> androidx.test.internal.runner.junit4.statement.RunBefores.evaluate(RunBefores.java:80)
> at
> androidx.test.rule.ActivityTestRule$ActivityStatement.evaluate(ActivityTestRule.java:527)
> at org.junit.rules.RunRules.evaluate(RunRules.java:20) at
> org.junit.runners.ParentRunner.runLeaf(ParentRunner.java:325) at
> org.junit.runners.BlockJUnit4ClassRunner.runChild(BlockJUnit4ClassRunner.java:78)
> at
> org.junit.runners.BlockJUnit4ClassRunner.runChild(BlockJUnit4ClassRunner.java:57)
> at org.junit.runners.ParentRunner$3.run(ParentRunner.java:290) at
> org.junit.runners.ParentRunner$1.schedule(ParentRunner.java:71) at
> org.junit.runners.ParentRunner.runChildren(ParentRunner.java:288) at
> org.junit.runners.ParentRunner.access$000(ParentRunner.java:58) at
> org.junit.runners.ParentRunner$2.evaluate(ParentRunner.java:268) at
> org.junit.runners.ParentRunner.run(ParentRunner.java:363) at
> org.junit.runners.Suite.runChild(Suite.java:128) at
> org.junit.runners.Suite.runChild(Suite.java:27) at
> org.junit.runners.ParentRunner$3.run(ParentRunner.java:290) at
> org.junit.runners.ParentRunner$1.schedule(ParentRunner.java:71) at
> org.junit.runners.ParentRunner.runChildren(ParentRunner.java:288) at
> org.junit.runners.ParentRunner.access$000(ParentRunner.java:58) at
> org.junit.runners.ParentRunner$2.evaluate(ParentRunner.java:268) at
> org.junit.runners.ParentRunner.run(ParentRunner.java:363) at
> org.junit.runner.JUnitCore.run(JUnitCore.java:137) at
> org.junit.runner.JUnitCore.run(JUnitCore.java:115) at
> androidx.test.internal.runner.TestExecutor.execute(TestExecutor.java:56)
> at
> androidx.test.runner.AndroidJUnitRunner.onStart(AndroidJUnitRunner.java:389)
> at
> android.app.Instrumentation$InstrumentationThread.run(Instrumentation.java:2075)

You try to log an object. outputProbabilityBuffer is an output object, you need the value of it.

In the Classifier class, try to log entry.getKey() and entry.getValue() in line 332. This should be the label name and confidence.

there is also a getter defined: getTitle() e.g. which is used in the CameraActivity at line 522 where you can log the output.

The issue is that you are using a binary classifier, but the example TensorFlowLite image recognition code expects you to use a multi-class classifier (i.e., a classifier that spits out more than one probability).

The model you are using spits out only one probability (which you can get by doing outputProbabilityBuffer.getFloatArray()), where 0 is one class and 1 is the other class.

To fix the issue, don't use the code that the example uses to process the output and map labels to probabilities after execution - just do something like:

// after execution
float result = outputProbabilityBuffer.getFloatArray()[0];
if (result < 0.5) {
  // means prediction was for category corresponding to 0
} else {
  // means prediction was for category corresponding to 1
}