Search code examples
pythondetectron

What do the class labels output be detectron 2 refer to?


In the documentation for detectron2, it states that class labels are located in output_dict['Instances'].pred_classes. This is all fine, and I can access this easily, but at no point in the documentation (or the output dictionary, as far as I can tell) does it specify which integer label refers to which class. I am assuming that somewhere there is a dictionary that contains something like {0: 'Person', 1: 'Bicycle', 2: 'Car',... }, but I can't seem to find it. Does anyone know where to find this?

Here are the specs for the output dictionary: https://detectron2.readthedocs.io/tutorials/models.html


Solution

  • The dataset metadata class has a property thing_classes which has a list of class names of that dataset. Simply by passing the class id from pred_classes one may get the class names as seen here.

    pred_classes = output_dict['instances'].pred_classes.cpu().tolist()
    class_names = MetadataCatalog.get("mydataset").thing_classes
    pred_class_names = list(map(lambda x: class_names[x], pred_classes))