Cannot copy from a TensorFlowLite tensor (Identity) with shape [1, 30, 8400] to a Java object with shape [1, 26]

i trained a model yolov8 with custom dataset containing 26 classess, but when i convert the model to tflite i noticed that it gives as output [1,30,8400] and this is what caused me errors when using my model with flutter.

the error

E/AndroidRuntime(18479): Caused by: java.lang.IllegalArgumentException: Cannot copy from a TensorFlowLite tensor (Identity) with shape [1, 30, 8400] to a Java object with shape [1, 26].

how can i modify the output shape of my model ?

this is how trained my model :

 from ultralytics import YOLO

 model = YOLO('yolov8s.pt')

 results = model.train(data='/kaggle/input/my- 
 dataset/my_dataset/data.yaml', epochs=100, imgsz=640)

and this is the content of file the data.yaml

train: /kaggle/input/dataset-asl/ASL/train/images
val: /kaggle/input/dataset-asl/ASL/valid/images

nc: 26
names: ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 
'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 
'Y', 'Z']

and this is how convert my model to format tflite :

from ultralytics import YOLO

# Load a model
model = YOLO('best.pt')

# Export the model
model.export(format='tflite')

and this is output after convert :

Ultralytics YOLOv8.2.4  Python-3.9.0 torch-2.2.1+cpu CPU (Intel 
Core(TM) i5-7300U 2.60GHz)
Model summary (fused): 168 layers, 3010718 parameters, 0 
gradients, 8.1 GFLOPs

PyTorch: starting from 'best.pt' with input shape (1, 3, 640, 640) 
BCHW and output shape(s) (1, 30, 8400) (6.0 MB)

TensorFlow SavedModel: starting export with tensorflow 2.15.0...
WARNING  tensorflow<=2.13.1 is required, but tensorflow==2.15.0 is 
currently installed 
https://github.com/ultralytics/ultralytics/issues/5161

ONNX: starting export with onnx 1.15.0 opset 17...
ONNX: simplifying with onnxsim 0.4.36...
ONNX: export success  1.8s, saved as 'best.onnx' (11.7 MB)
TensorFlow SavedModel: starting TFLite export with onnx2tf 
1.17.5...
TensorFlow SavedModel: export success  14.3s, saved as 
'best_saved_model' (29.5 MB)

TensorFlow Lite: starting export with tensorflow 2.15.0...
TensorFlow Lite: export success  0.0s, saved as 
'best_saved_model\best_float32.tflite' (11.7 MB)

Export complete (16.8s)
Results saved to C:\Users\Bachir\Desktop\api\tflite
Predict:         yolo predict task=detect 
model=best_saved_model\best_float32.tflite imgsz=640  
Validate:        yolo val task=detect 
model=best_saved_model\best_float32.tflite imgsz=640 
data=/kaggle/input/my-dataset/my_dataset/data.yaml  
Visualize:       https://netron.app
'best_saved_model\\best_float32.tflite'
emphasized text

Solution

Your model output shape is correct. The 30 in your output shape indicate that [4 bounding box location + 26 class scores]. 8400 indicates that there are 8400 possible bounding boxes. Tensorflow lite doesn't comes with non max suppression. So you will need to combine some of those boxes to get the optimum results.