How to update class names after model has already been trained

[ashleyha]

❔Question

Is there a way to update the class names for a model after it has been trained?

Additional context

For example, let's say you made a typo in one of the class names and wanted to fix it without re-training the model. I believe in yolov3 you could specify a --cfg param to detect.py to override the class names, is there a way to do this with yolov5?

[glenn-jocher]

@ashleyha class names are model attributes. Once you load a model you're free to modify the names list:

model = ... # load model

print(model.names)

[github-actions]

👋 Hello, this issue has been automatically marked as stale because it has not had recent activity. Please note it will be closed if no further activity occurs.

Access additional YOLOv5 🚀 resources:

• Wiki – https://github.com/ultralytics/yolov5/wiki
• Tutorials – https://docs.ultralytics.com/yolov5
• Docs – https://docs.ultralytics.com

Access additional Ultralytics ⚡ resources:

• Ultralytics HUB – https://ultralytics.com
• Vision API – https://ultralytics.com/yolov5
• About Us – https://ultralytics.com/about
• Join Our Team – https://ultralytics.com/work
• Contact Us – https://ultralytics.com/contact

Feel free to inform us of any other issues you discover or feature requests that come to mind in the future. Pull Requests (PRs) are also always welcomed!

Thank you for your contributions to YOLOv5 🚀 and Vision AI ⭐!

[korhanpolat]

how do we save it after the modification

[vaibhavpanday]

@glenn-jocher how to save after modification

[glenn-jocher]

@vaibhavpanday after modifying the class names, you can save the modified model by calling the save method. Here's an example:

model = ... # load model

# Modify class names
model.names = ['class1', 'class2', 'class3']

# Save the modified model
model.save('path/to/save/model.pt')

This will save the modified model with the updated class names to the specified path.

[thachln]

@glenn-jocher I have tried and faced this error:
AttributeError: property 'names' of 'YOLO' object has no setter

[glenn-jocher]

@thachln ah, my apologies for the confusion. The names attribute does not directly support reassignment. Instead, you need to modify the list it points to. Try this approach:

model = ... # load model

# Modify class names directly in the existing list
model.names[0] = 'new_class_name1'  # example for modifying the first class name

Remember, this modifies the list in place. If you need to completely replace the class names, ensure you modify each index as required. For saving changes, the model structure itself does not directly support saving these specific modifications, as model.save() is intended for weights. Instead, you should manage class names separately in your application or consider re-exporting your dataset with corrected names and re-training if class name persistence is critical.

[Malginn]

This code above didn't work for me, so I reached under the hood and found attributes that allow you to change the names of classes to the inference of the model:
model = YOLO('path')
model.model.names = {0: 'new_name1', 1: 'new_name2', 2: 'new_name3'}
model.predct(...)

[glenn-jocher]

Hi @Malginn,

Thank you for diving into the code and sharing your findings! It's great to see such proactive troubleshooting. To ensure we can assist you effectively, could you please provide a minimum reproducible code example? This will help us understand the context and reproduce the issue on our end. You can refer to our guide on creating a minimum reproducible example here: Minimum Reproducible Example.

Additionally, please make sure you are using the latest versions of torch and the YOLOv5 repository. Sometimes, issues are resolved in newer releases, and updating might solve the problem.

Regarding your approach, modifying the model.model.names dictionary is indeed a valid way to change class names for inference. Here's a refined example to ensure clarity:

from yolov5 import YOLO

# Load the model
model = YOLO('path/to/your/model.pt')

# Modify class names
model.model.names = {0: 'new_name1', 1: 'new_name2', 2: 'new_name3'}

# Perform inference
results = model.predict('path/to/your/image.jpg')

If you encounter any further issues, please share the specific error messages or unexpected behaviors you are seeing. This will help us diagnose and address the problem more effectively.

Thank you for your collaboration and patience. The YOLO community and the Ultralytics team are here to support you!

[MoAbbasid]

@Malginn
I was facing the same issue, and your solution worked for me, Thank you so much for sharing it.
any idea why we its this way? why we have to do model.model.names?

[Eswararaokotakota]

This code above didn't work for me, so I reached under the hood and found attributes that allow you to change the names of classes to the inference of the model: model = YOLO('path') model.model.names = {0: 'new_name1', 1: 'new_name2', 2: 'new_name3'} model.predct(...)

Thank you for this
CostumModel.model.names = {0: 'new_name1', 1: 'new_name2', 2: 'new_name3'}

[glenn-jocher]

Hi @Eswararaokotakota,

Thank you for your proactive approach and for sharing your solution with the community! 🎉

The reason you need to use model.model.names is due to the way the YOLOv5 architecture is structured. The model object is a wrapper that contains the core model under model.model. This core model holds the names attribute, which is why you need to access it through model.model.names.

Your solution is indeed correct and helps in modifying the class names for inference without retraining the model. Here's a refined example for clarity:

from yolov5 import YOLO

# Load the model
model = YOLO('path/to/your/model.pt')

# Modify class names
model.model.names = {0: 'new_name1', 1: 'new_name2', 2: 'new_name3'}

# Perform inference
results = model.predict('path/to/your/image.jpg')

If you encounter any further issues or have more questions, feel free to ask. The YOLO community and the Ultralytics team are always here to help!

Best regards and happy coding! 🚀