Привет, друзья! Вы не представляете, сколько я времени провел, мучаясь с переименованием классов в YOLO. Три дня — это не шутка, и благодаря нейросети, которая подсказала мне решение на сайте документации https://docs.ultralytics.com/models/yolov8/, я наконец-то выбрался из этого горячего ада!
Итак, давайте перейдем к делу! Чтобы поменять названия классов в модели YOLO, нужен код, которым я делюсь с вами ниже. Это как делать тосты: немного терпения, и все получится!
from ultralytics import YOLO
model = YOLO('models/best_original.pt')
# Get current names
print('Before: ')
print(model.model.names)
new_class_names = {
0: 'Helm',
1: 'Mask',
2: 'No_Helm',
3: 'No_Mask',
4: 'No_Vest',
5: 'Person',
6: 'Safety_Cone',
7: 'Vest',
8: 'Machinery',
9: 'Vehicle'
}
# Change name here
for key in new_class_names.keys():
model.model.names[key] = new_class_names[key]
print('After: ')
print(model.model.names)
if input('save?') == 'y':
model.save('models/rename_original.pt')
print('After in rename_original: ')
model = YOLO('models/rename_original.pt')
print(model.model.names)
Итак, вот краткий курс действия:
- Импортируем модель YOLO из библиотеки Ultralyics. Ничего сложного, просто как плевать в стену.
- Получаем текущие названия классов, чтобы посмотреть, что у нас тут за шедевры.
- Создаем словарь
new_class_names, который будет хранить новые названия классов. Звучит просто, правда? - Запускаем цикл, который проходит по ключам нашего словаря и заменяет старые названия на новые. Можно было бы даже устроить вечеринку, но лучше сохранить хоть какую-то серьезность.
- Выводим список классов после замен и спрашиваем, не хотим ли мы сохранить изменения. Если да, то сохраняем в файл, чтобы больше не париться с этим вопросом!
После всего этого, вы увидите «After» с вашими новыми, блестящими названиями классов. Прямо как у меня:
{0: 'Helm', 1: 'Mask', 2: 'No_Helm', 3: 'No_Mask', 4: 'No_Vest', 5: 'Person', 6: 'Safety_Cone', 7: 'Vest', 8: 'Machinery', 9: 'Vehicle'}
В общем, не делайте, как я, и не теряйте три дня — просто следуйте этой инструкции, и у вас все получится с первого раза! Удачи вам в переименовании классов, и пусть нейросети всегда будут на вашей стороне!
