Shape of Output Tensor is 6 instead of 4

General Discussion
,
#1

I have trained a custom YOLOv5 model for custom object detection. Then exported the trained model to TFlite format. I want to draw bounding box around detected object. To do that, the Location Tensor (see image) should be of shape 4 i.e. return 4 float values as coordinates (as mentioned in Tensorflow docs) but my model’s location tensor is [1, 25200, 6] => 6 float values.
This is making my model incompatible for integration in Android studio.

How to modify my model so that Output Tensor (aka Location Tensor in this case) returns only 4 values ?

#3

But Yolo models in general will return that shape.

The reason is that they output an array of bounding boxes with probabilities i.e

  • [xcenter, ycenter, width, height, isObject, isClass1, isClass2,...]

In your case “6” is means you likely have a single class. @Mohit

Can you explain why does this make it incompatible with Android Studio ?

1 Like
#4

Thank you for clarifying what those 6 values are! And yes I have a single class.

Visit this link to understand why it is incompatible with Android Studio: https://www.tensorflow.org/lite/inference_with_metadata/task_library/object_detector#model_compatibility_requirements

Scroll down a bit and you will see this:

Output tensors must be the 4 outputs of a DetectionPostProcess op, i.e:

Locations tensor (kTfLiteFloat32)
tensor of size [1 x num_results x 4], the inner array representing bounding boxes in the form [top, left, right, bottom].

#5

One of the ways to solve this problem is to somehow extract the first 4 values out of those 6 values in the Output array. Any idea how I can do that ?

#7

I think you will need to add some code to be sure.

Because it may be possible one, both or none of these:

  • Add a layer to the model to slice output
  • Slice the output tensor
1 Like