Всем привет! Решила я вернуться на Хабр с новым мини-проектом. Сегодня попробуем детектить дорожные знаки используя YOLOv8. Что ж, приступим!
1: Работа с Google Colab
Первое что мы делаем это открываем Google colab и создаем New Notebook. После того как всё создалось, то сверху ищем <среду выполнения>. Дальше такие действия: Сменить среду выполнения -> T4 GPU -> Сохранить. (По незнанию мне пришлось запускать проект 2 раза, потому что у меня стояла среда выполнения CPU, поэтому сделаем сразу все важные моменты в самом начале).
P.s. меня забанили пока я писала статью, поэтому у вас есть 2 варианта (я не работала ещё ни разу с Google Colab, прошу не кидаться тапками):
1: переписать код и переключиться на GPU уже после написанного.
2: использовать готовый блокнот и переключиться на GPU, когда уже будете готовы запускать код.
Иначе получите это: Невозможно подключиться к ускорителю (GPU) из-за лимитов на использование в Colab.
Так-с, перейдём сразу к делу, проверим всё ли у нас правильно выбралось:
!nvcc --versionnvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Sep_21_10:33:58_PDT_2022
Cuda compilation tools, release 11.8, V11.8.89
Build cuda_11.8.r11.8/compiler.31833905_0
!nvidia-smiSat Aug 12 08:27:04 2023
NVIDIA-SMI 525.105.17 Driver Version: 525.105.17 CUDA Version: 12.0
...
Если получим такие выводы, то всё отлично. Можно продолжать работать :)
Чтобы не путаться в последующих файлах, сразу удалим папку sample_data, а то она немного будет нам мешаться.
!rm -rf ./sample_data2: Установка утилит
2.1: Клонируем репозиторий
!git clone https://github.com/ultralytics/JSON2YOLOЕсть несколько важных деталей, которые нужно сделать в репозитории:
1: В файле general_json2yolo.py меняем код в 274 строке на h, w, f = img['height'], img['width'], img['file_name'].split('/')[1] .
2: В general_json2yolo.py и labelbox_json2yolo.py, нужно добавить перед каждым utils точку, чтобы было так from .utils import ...
Если у вас возникли вопросы почему так надо сделать, то объясню:
1: Метод .split('/') разделяет строку img['file_name'] по символу / и создает список подстрок. Затем [1] выбирает второй элемент этого списка, то есть вторую подстроку после разделителя /. Строка img['file_name'] содержит путь к файлу с одним или несколькими слешами /, и нам нужно получить имя файла без пути.
2: Я долго не понимала, что я делаю не так и почему у меня не импортится файл, который лежит в той же папке с файлами, которые мы используем. Пришлось лезть в документацию импортов и смотреть какие варианты вообще есть. На своем железе такой ошибки не было, но в Google Colab она почему-то вылезла. Если у вас есть какие-то объяснения этой проблемы, то напишите в комментариях, буду благодарна :)
Так-с, идём дальше!
2.2: Переходим к установке библиотек
%pip install -r /content/JSON2YOLO/requirements.txt%pip install kaggle ipywidgets widgetsnbextensionЕсли всё successfully installed, то переходим к самому датасету!
3: Работа с датасетом
3.1: Скачиваем датасет
import osДальше нам потребуется Kaggle. Если нет аккаунта -> создаём его, а если есть то переходим в Settings -> API -> Create New Token. Автоматически скачается json файл, там уже будут данные.
os.environ['KAGGLE_USERNAME'] = "сюда вписываем из json username"
os.environ['KAGGLE_KEY'] = "а сюда key"
!kaggle datasets download -d watchman/rtsd-datasetНемного расскажу о самом датасете, пока он скачивается. Набор данных RTSD содержит кадры, предоставленные компанией "Геоцентр Консалтинг". Изображения получены с широкоформатного видеорегистратора, который снимает с частотой 5 кадров в секунду. Разрешения изображений от 1280×720 до 1920×1080. Фотографии сделаны в разное время года, в разное время суток и при различных погодных условиях . В наборе используется 155 знаков дорожного движения в формате разметки - COCO.
3.2: Распаковываем архив с данными
import zipfilearchive = zipfile.ZipFile('rtsd-dataset.zip', 'r')
archive.extractall('.')3.3: Ну и чтобы он нам не мешался, то можем его удалить
os.remove('rtsd-dataset.zip')4: Преобразование датасета в YOLO-формат
4.1: Импортим библиотеки
import pandas as pd
from tqdm.notebook import tqdm
from shutil import copyfile, move
import sys
import json
from ipywidgets import FloatProgress4.2: Переходим к конвертации
Мы будем использовать готовый скрипт от Ultralytics, который скачали в самом начале. Это папка JSON2YOLO. Перейдём к конвертации СOCO-формата в YOLO-формат. Советую делать всё постепенно, если вы используете Google Colab. У меня возникали ошибки при выполнении, если я использовала код целиком.
from JSON2YOLO.general_json2yolo import convert_coco_json
sys.path.append('./JSON2YOLO')
test_path = 'test_annotation'
train_path = 'train_annotation'
os.makedirs(train_path, exist_ok=True)
os.makedirs(test_path, exist_ok=True)move('train_anno.json', os.path.join(train_path, 'train_anno.json'))
move('val_anno.json', os.path.join(test_path, 'val_anno.json'))for folder in ['labels', 'images']:
for path in [test_path, train_path]:
os.makedirs(os.path.join(path, folder), exist_ok=True)convert_coco_json(train_path)
for file in tqdm(os.listdir(os.path.join('new_dir/labels/train_anno'))):
move(os.path.join('new_dir/labels/train_anno', file), os.path.join(train_path, 'labels', file))
convert_coco_json('./test_annotation/')
for file in tqdm(os.listdir(os.path.join('new_dir/labels/val_anno'))):
move(os.path.join('new_dir/labels/val_anno', file), os.path.join(test_path, 'labels', file))Отлично, разметка у нас есть, не хватает лишь изображений, которые ей соответствуют. Добавим их :)
test_labels = os.listdir(os.path.join(test_path, 'labels'))
train_labels = os.listdir(os.path.join(train_path, 'labels'))
test_labels = set(map(lambda x: x.split('.')[0], test_labels))
train_labels = set(map(lambda x: x.split('.')[0], train_labels))images = 'rtsd-frames/rtsd-frames'
for file in os.listdir(images):
name = file.split('.')[0]
if name in train_labels:
move(os.path.join(images, file), os.path.join(train_path,'images', file))
if name in test_labels:
move(os.path.join(images, file), os.path.join(test_path,'images', file))5: Работа с YAML файлом
Создадим файл "trafic_signs.yaml" с описанием путей и классов, используемых в датасете. Это обязательное требование для YOLOv8.
5.1: Установим библиотеку для записи данных yaml в файл.
%pip install PyYAML==5.15.2: Импортим библитеку
import yaml5.3: Опишем пути и классы для обучения
data = [{
'train': '/content/train_annotation/images',
'val': '/content/test_annotation/images',
'nc': 155,
'names': ['2_1', '1_23', '1_17', '3_24', '8_2_1', '5_20', '5_19_1', '5_16', '3_25',
'6_16', '7_15', '2_2', '2_4', '8_13_1', '4_2_1', '1_20_3', '1_25', '3_4', '8_3_2',
'3_4_1', '4_1_6', '4_2_3', '4_1_1', '1_33', '5_15_5', '3_27', '1_15', '4_1_2_1',
'6_3_1', '8_1_1', '6_7', '5_15_3', '7_3', '1_19', '6_4', '8_1_4', '8_8', '1_16',
'1_11_1', '6_6', '5_15_1', '7_2', '5_15_2', '7_12', '3_18', '5_6', '5_5', '7_4',
'4_1_2', '8_2_2', '7_11', '1_22', '1_27', '2_3_2', '5_15_2_2', '1_8', '3_13',
'2_3', '8_3_3', '2_3_3', '7_7', '1_11', '8_13', '1_12_2', '1_20', '1_12', '3_32',
'2_5', '3_1', '4_8_2', '3_20', '3_2', '2_3_6', '5_22', '5_18', '2_3_5', '7_5',
'8_4_1', '3_14', '1_2', '1_20_2', '4_1_4', '7_6', '8_1_3', '8_3_1', '4_3', '4_1_5',
'8_2_3', '8_2_4', '1_31', '3_10', '4_2_2', '7_1', '3_28', '4_1_3', '5_4', '5_3',
'6_8_2', '3_31', '6_2', '1_21', '3_21', '1_13', '1_14', '2_3_4', '4_8_3', '6_15_2',
'2_6', '3_18_2', '4_1_2_2', '1_7', '3_19', '1_18', '2_7', '8_5_4', '5_15_7', '5_14',
'5_21', '1_1', '6_15_1', '8_6_4', '8_15', '4_5', '3_11', '8_18', '8_4_4', '3_30',
'5_7_1', '5_7_2', '1_5', '3_29', '6_15_3', '5_12', '3_16', '1_30', '5_11', '1_6',
'8_6_2', '6_8_3', '3_12', '3_33', '8_4_3', '5_8', '8_14', '8_17', '3_6', '1_26',
'8_5_2', '6_8_1', '5_17', '1_10', '8_16', '7_18', '7_14', '8_23']
}]def write_yaml_to_file(py_obj, filename) :
with open(f'{filename}.yaml', 'w+',) as f:
yaml.dump_all(py_obj, f, sort_keys=False)
write_yaml_to_file(data, 'trafic_signs')6: Доп. установка библиотек
6.1: Устанавливаем библиотеки
%pip install ultralytics7: Обучение
Переходим непосредственно к обучению модели, наконец-то :)
7.1: Импортим библиотеки
from ultralytics import YOLO
import torch
import gcЕсли спросите зачем импортить torch и gc, то дам такой ответ. Они нужны для очистки графической памяти, если вам нужно будет это сделать, то пропишите:
gc.collect()
torch.cuda.empty_cache()7.2: Начинаем обучать модель
model = YOLO('yolov8s.pt')
results = model.train(
data='/content/trafic_signs.yaml',
imgsz=1280,
epochs=5,
batch=6,
device=0,
name='YOLOv8s'
)P.s. если хотите, то можете поколдовать над размером батчей, у меня стоит 5 для обучения на colab, т.к. там всего 15гб GPU_mem. Ну или можете взять и другую модель от YOLOv8, а можно и больше эпох сделать. Тут уже всё на ваш выбор :)
8: Итоги обучения
Самое время посмотреть на результаты после обучение, давайте потестируем модельку :)
model = YOLO("путь к лучшему весу -> best.pt")pre = model.predict(
source="путь к фото / видео, где хотите предсказать",
show=True,
imgsz=1280,
hide_labels=True,
save=True,
name="название папки для результатов",
conf=0.1,
)
Как итог могу сказать, что получилось у меня обучить модельку всего 4 эпохи (вышло практически 6 часов), на 5 эпохе Google colab решил, что я превысила время использования ресурсов и остановил обучение не дав скачать веса и удалив все файлы. Так что совет на будущее, если будете обучать на colab, то сохраняйте веса каждую эпоху, как это делала я, просто потом дообучите модельку и всё. Думаю, что я так и сделаю в ближайшее время, чтобы сравнить результаты.
Видно, что на 2 фотке некоторые знаки не получилось распознать из-за столба, но думаю что эту проблему можно решить. На 3 тесте тоже есть проблемы с распознаванием знака работы эвакуатора, хотя остальные 2 ей удалось найти. Итоги подведены, поэтому спасибо за прочтение статьи, надеюсь, она была полезной!
Я только начинаю изучать CV и ML, поэтому, если у вас есть рекомендации по материалам/проектам, то порекомендуйте что-нибудь, буду благодарна :)
