Add Metadata doesn't work with YoloV5 model

[chainyo]

Hi, I'm trying to add metadata to a converted YoloV5 model.

The conversion from model.pt to model.tflite is working well (thanks to yolov5 scripts), but I can't add metadata with tflite support because the object detection template is waiting for 4 tensors in his output layer when there is only one tensor in the YoloV5 output layer. Anyone know how to add metadata to this kind of model ?

This is the error I get everytime I'm trying to run the script:

ValueError: The number of output tensors (1) should match the number of output tensor metadata (4)

This is my script to add metadata to my model:

from tflite_support.metadata_writers import object_detector
from tflite_support.metadata_writers import writer_utils
from tflite_support import metadata

ObjectDetectorWriter = object_detector.MetadataWriter
_MODEL_PATH = "runs/train/exp2/weights/model.tflite"
_LABEL_FILE = "runs/train/exp2/weights/labelmap.txt"
_SAVE_TO_PATH = "runs/train/exp2/model_metadata.tflite"
_INPUT_NORM_MEAN = 127.5
_INPUT_NORM_STD = 127.5

writer = ObjectDetectorWriter.create_for_inference(
    writer_utils.load_file(_MODEL_PATH), 
    [_INPUT_NORM_MEAN], 
    [_INPUT_NORM_STD], 
    [_LABEL_FILE]
)

print(writer.get_metadata_json())

writer_utils.save_file(writer.populate(), _SAVE_TO_PATH)

This is what is printed with writer.get_metadata.json():

{
  "name": "ObjectDetector",
  "description": "Identify which of a known set of objects might be present and provide information about their positions within the given image or a video stream.",
  "subgraph_metadata": [
    {
      "input_tensor_metadata": [
        {
          "name": "image",
          "description": "Input image to be detected.",
          "content": {
            "content_properties_type": "ImageProperties",
            "content_properties": {
              "color_space": "RGB"
            }
          },
          "process_units": [
            {
              "options_type": "NormalizationOptions",
              "options": {
                "mean": [
                  127.5
                ],
                "std": [
                  127.5
                ]
              }
            }
          ],
          "stats": {
            "max": [
              255.0
            ],
            "min": [
              0.0
            ]
          }
        }
      ],
      "output_tensor_metadata": [
        {
          "name": "location",
          "description": "The locations of the detected boxes.",
          "content": {
            "content_properties_type": "BoundingBoxProperties",
            "content_properties": {
              "index": [
                1,
                0,
                3,
                2
              ],
              "type": "BOUNDARIES"
            },
            "range": {
              "min": 2,
              "max": 2
            }
          },
          "stats": {
          }
        },
        {
          "name": "category",
          "description": "The categories of the detected boxes.",
          "content": {
            "content_properties_type": "FeatureProperties",
            "content_properties": {
            },
            "range": {
              "min": 2,
              "max": 2
            }
          },
          "stats": {
          },
          "associated_files": [
            {
              "name": "labelmap.txt",
              "description": "Labels for categories that the model can recognize.",
              "type": "TENSOR_VALUE_LABELS"
            }
          ]
        },
        {
          "name": "score",
          "description": "The scores of the detected boxes.",
          "content": {
            "content_properties_type": "FeatureProperties",
            "content_properties": {
            },
            "range": {
              "min": 2,
              "max": 2
            }
          },
          "stats": {
          }
        },
        {
          "name": "number of detections",
          "description": "The number of the detected boxes.",
          "content": {
            "content_properties_type": "FeatureProperties",
            "content_properties": {
            }
          },
          "stats": {
          }
        }
      ],
      "output_tensor_groups": [
        {
          "name": "detection_result",
          "tensor_names": [
            "location",
            "category",
            "score"
          ]
        }
      ]
    }
  ]
}

I need metadata in order to import the model in Android Studio.

Best regards.

[lu-wang-g]

@ChainYo, Yolo with one output tensor is not supported in Model Metadata and Task Library. From the other github issue, seems like there's a version of Yolo with 4 output tensors. If so, that will be supported.

[chainyo]

It seems to be only with tfjs models not tflite models. But thanks for the quick answer

[glenn-jocher]

@ChainYo @lu-wang-g yes YOLOv5 exported models generally concat outputs into a single output. TFLite models do not export with NMS, only TF.js and pipelined CoreML models contains NMS.

[Lochan5403]

hey, i have the same issue. Could you please tell me how you overcame it? Did you find a way to add 4 tensors in it's output layer in the metadata? if so, how?