Обучение YOLOv8s на Google Colab: детектим дорожные знаки13.08.2023 15:46

Всем привет! Решила я вернуться на Хабр с новым мини-проектом. Сегодня попробуем детектить дорожные знаки используя YOLOv8. Что ж, приступим!

1: Работа с Google Colab.

Первое что мы делаем это открываем Google colab и создаем New Notebook. После того как всё создалось, то сверху ищем <среду выполнения>. Дальше такие действия: Сменить среду выполнения → T4 GPU → Сохранить. (По незнанию мне пришлось запускать проект 2 раза, потому что у меня стояла среда выполнения CPU, поэтому сделаем сразу все важные моменты в самом начале).

P.s. меня забанили пока я писала статью, поэтому у вас есть 2 варианта (я не работала ещё ни разу с Google Colab, прошу не кидаться тапками):
1: переписать код и переключиться на GPU уже после написанного.
2: использовать готовый блокнот и переключиться на GPU, когда уже будете готовы запускать код.
Иначе получите это: Невозможно подключиться к ускорителю (GPU) из-за лимитов на использование в Colab.

Так-с, перейдём сразу к делу, проверим всё ли у нас правильно выбралось: 

!nvcc --version

nvcc: 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-smi

Sat 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_data

2: Установка утилит.

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 zipfile
archive = 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 FloatProgress

4.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.1

5.2: Импортим библитеку.

import yaml

5.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 ultralytics

7: Обучение.

Переходим непосредственно к обучению модели, наконец-то :)

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,                                                   
    )

тест 1

тест 1

тест 2

тест 2

тест 3

тест 3

Как итог могу сказать, что получилось у меня обучить модельку всего 4 эпохи (вышло практически 6 часов), на 5 эпохе Google colab решил, что я превысила время использования ресурсов и остановил обучение не дав скачать веса и удалив все файлы. Так что совет на будущее, если будете обучать на colab, то сохраняйте веса каждую эпоху, как это делала я, просто потом дообучите модельку и всё. Думаю, что я так и сделаю в ближайшее время, чтобы сравнить результаты.

Видно, что на 2 фотке некоторые знаки не получилось распознать из-за столба, но думаю что эту проблему можно решить. На 3 тесте тоже есть проблемы с распознаванием знака работы эвакуатора, хотя остальные 2 ей удалось найти. Итоги подведены, поэтому спасибо за прочтение статьи, надеюсь, она была полезной!

Я только начинаю изучать CV и ML, поэтому, если у вас есть рекомендации по материалам/проектам, то порекомендуйте что-нибудь, буду благодарна :)

Ну и оставлю ссылочки:  GitHub,  Telegram

© Habrahabr.ru