Привет, Хабр!

Этой статьей я открываю цикл материалов по разбору задач Всероссийского чемпионата "Цифровой Прорыв", связанных с Computer Vision. Решение, предлагаемое в статье, позволяет получить место в топ-5 лидерборда. Особенность данного подхода в том, что мы не будем использовать данные для обучения, предлагаемые организаторами соревнования. В конце даются советы по улучшению решения, а также идеи, которые могут привести к первому месту.

Введение

Автомобилестроение — одна из самых высокотехнологичных индустрий, которая вбирает в себя инновации из различных сфер, чтобы обеспечить водителям и пассажирам скорость, безопасность и комфорт при движении. Внимание многих автопроизводителей сейчас сфокусировано на создании беспилотных машин, что подразумевает внедрение целого комплекса программно-аппаратных решений, работающих, в том числе, на основе технологий искусственного интеллекта.

С каждым годом необходимость в системе автоматического распознавания дорожных знаков становится все актуальнее. Данные системы широко применяются в автопилотах и помощниках водителя для повышения безопасности движения автотранспортных средств. Системы могут помочь придерживаться установленного скоростного режима, соблюдать ограничения на проезд и обгон, что поможет существенно снизить аварийность на дорогах.

Условие задачи

Разработать решение, которое сможет распознавать дорожные знаки на кадрах записанных автомобильным видеорегистратором.

Описание входных данных 

• train/ — папка, содержит в себе 778 кадров снятых на видеорегистратор;

• train.csv — содержит перечисление знаков для каждой фотографии;

• test/ — содержит в себе 388 изображений на которых требуется определить автомобильные знаки;

• test.csv — содержит перечисление всех изображений тестового набора;

• sample_solution.csv — пример файла для отправки;

Пояснение к данным

Для удобство интерпретации результатов дорожные знаки были преобразованы в цифры от 1 до 70, где:

цифре 1 соответствует знак под ГОСТ ‘3.24',

цифре 2 соответствует знак под ГОСТ '1.16',

цифре 3 соответствует знак под ГОСТ '5.15.5',

и т.д. для следующих знаков: '5.19.1', '5.19.2', '1.20.1', '8.23', '2.1', '4.2.1', '8.22.1', '6.16', '1.22', '1.2', '5.16', '3.27', '6.10.1', '8.2.4', '6.12', '5.15.2', '3.13', '3.1', '3.20', '3.12', '7.14.2', '5.23.1', '2.4', '5.6', '4.2.3', '8.22.3', '5.15.1', '7.3', '3', '2.3.1', '3.11', '6.13', '5.15.4', '8.2.1', '1.34.3', '8.2.2', '5.15.3', '1.17', '4.1.1', '4.1.4', '3.25', '1.20.2', '8.22.2', '6.9.2', '3.2', '5.5', '5.15.7', '7.12', '8.2.3', '5.24.1', '1.25', '3.28', '5.9.1', '5.15.6', '8.1.1', '1.10', '6.11', '3.4', '6.10', '6.9.1', '8.2.5', '5.15', '4.8.2', '8.22', '5.21', '5.18'.

На что стоить обратить внимание

Важно отметить что на одном снимке может быть более одного знака, но максимально их число на одной фотографии для нашего набора - восемь.

Метрика качества

В задаче важна точность распознавания модели, поэтому к каждой строке набора будет применен Recall.

Решение задачи

Даже беглый просмотр информации о задаче сразу даёт понять, что данных, мягко говоря, недостаточно. У нас 70 различных классов, и всего 778 фотографий для обучения. Полезно построить гистограмму частот появления того или иного знака в обучающем наборе.

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

Альтернативный набор данных

Для наших нужд лучше всего подходит набор данных RTSD. Набор данных RTSD содержит кадры, предоставленные компанией "Геоцентр Консалтинг". Изображения получены с широкоформатного видеорегистратора, который снимает с частотой 5 кадров в секунду. Разрешения изображений от 1280×720 до 1920×1080. Фотографии были сделаны в разное время года (весна, осень, зима), в разное время суток (утро, день, вечер) и при различных погодных условиях (дождь, снег, яркое солнце). В наборе используется 155 знак дорожного движения, формат разметки - COCO.

Немного статистики

RSTD отличается по количеству и составу знаков от набора данных нашей задачи. Так что следующий вопрос назревает сам собой - какое количество знаков из исходной задачи покрывает наш набор данных ? Знаки из набора RTSD составляют 65.2 % от знаков дорожного движения в нашей задаче.

Как мы убедились выше, у знаков разная частота появления. Предположим, что соотношение знаков в обучающем и тестовом наборе одинаковое. Какой объём train-набора охватывают знаки, которые присутствуют в RTSD ? Знаки из набора RTSD охватывают 72.4 % от объема всех знаков дорожного движения в train-наборе. Таким образом, мы можем покрыть большую часть кейсов, вообще не используя train из нашей задачи. По-моему, это ????.

А теперь самое время импортировать все необходимые библиотеки.

import pandas as pd
from tqdm.notebook import tqdm
import os
from shutil import copyfile, move
import sys
import json

Загрузим набор данных с Kaggle.

!pip install kaggle
!kaggle datasets download -d watchman/rtsd-dataset
!7z x rtsd-dataset.zip

Детектор объектов

В качестве детектора будет выступать yolov5, а именно yolov5m6 с разрешением 1280 пикселей. По моему мнению, это оптимальный вариант, так как знаки дорожного движения маленькие, а само изображение большое; если мы будем обучать модель на разрешении 640, то можем пропустить значительное количество знаков дорожного движения. Конечно, вы всегда можете использовать другую архитектуру. Задача этой статьи не выбор оптимальной архитектуры, а демонстрация того, что можно занимать призовые места в соревновании, не прибегая к разметке организаторов. 

Преобразование набора данных в YOLO-формат

Есть несколько способов, как это можно сделать. Например, воспользоваться сервисом Roboflow, но тогда придётся загружать всю разметку в их сервис, что займёт довольно много времени. Как альтернатива - cvat, но все это очень долго, хотя и потребует лишь терпения и времени. Всегда можно включить hard-mode и написать всё самому. Мы же будем использовать готовый скрипт от Ultralytics, но внесём в него одно изменение.

!git clone https://github.com/ultralytics/JSON2YOLO

Нужно модифицировать 274 строку в файле general_json2yolo.py следующим образом:

h, w, f = img['height'], img['width'], img['file_name'].split('/')[1]

Перейдём непосредственно к конвертации СOCO-формата в YOLO-формат.

sys.path.append('./JSON2YOLO')
from JSON2YOLO.general_json2yolo import convert_coco_json

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

Создадим файл "trafic_signs.yaml" с описанием путей и классов, используемых в датасете. Это обязательное требование для yolov5.

train: /home/jovyan/train_annotation/images  # train images (relative to 'path') 128 images
val: /home/jovyan/test_annotation/images  # val images (relative to 'path') 128 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']

Обучение модели

!git clone https://github.com/ultralytics/yolov5
!cd "yolov5"
!pip install -r requirements.txt
!python train.py --img 1280 --batch -1 --epochs 40 --data "/home/jovyan/trafic_signs.yaml" --weights yolov5m6.pt --project "hackaton_trafic_signs" --name "yolov5m6"

Модель, определяющая знаки дорожного движения, у нас есть, перейдём к тестовому набору.

Предсказание на тестовом наборе данных

Применим наш детектор объектов к тестовому набору данных.

!python detect.py --source {путь к тестовому набору} --weights {путь к весам модели} --save-txt --save-conf --name "yolov5m6_signs_test" --imgsz 1280 --conf-thres 0.25

Замечание: в RTSD отсутствует знак '5.19.2', но он входит в топ-3 самых частых знаков в обучающем наборе нашей задачи. Посмотрим, как ведёт себя детектор при наличии знаков '5.19.1' и '5.19.2' на изображении. Получается, если мы встретили два раза '5.19.1' в результате работы детектора, то это не что иное, как '5.19.1' и '5.19.2'.

Научимся преобразовывать знаки дорожного движения из RTSD в формат нашей задачи.

sings_rtsd = {"2_1": 1, "1_23": 2, "1_17": 3, "3_24": 4, "8_2_1": 5, "5_20": 6, "5_19_1": 7, "5_16": 8, "3_25": 9, "6_16": 10, "7_15": 11, "2_2": 12, "2_4": 13, "8_13_1": 14, "4_2_1": 15, "1_20_3": 16, "1_25": 17, "3_4": 18, "8_3_2": 19, "3_4_1": 20, "4_1_6": 21, "4_2_3": 22, "4_1_1": 23, "1_33": 24, "5_15_5": 25, "3_27": 26, "1_15": 27, "4_1_2_1": 28, "6_3_1": 29, "8_1_1": 30, "6_7": 31, "5_15_3": 32, "7_3": 33, "1_19": 34, "6_4": 35, "8_1_4": 36, "8_8": 37, "1_16": 38, "1_11_1": 39, "6_6": 40, "5_15_1": 41, "7_2": 42, "5_15_2": 43, "7_12": 44, "3_18": 45, "5_6": 46, "5_5": 47, "7_4": 48, "4_1_2": 49, "8_2_2": 50, "7_11": 51, "1_22": 52, "1_27": 53, "2_3_2": 54, "5_15_2_2": 55, "1_8": 56, "3_13": 57, "2_3": 58, "8_3_3": 59, "2_3_3": 60, "7_7": 61, "1_11": 62, "8_13": 63, "1_12_2": 64, "1_20": 65, "1_12": 66, "3_32": 67, "2_5": 68, "3_1": 69, "4_8_2": 70, "3_20": 71, "3_2": 72, "2_3_6": 73, "5_22": 74, "5_18": 75, "2_3_5": 76, "7_5": 77, "8_4_1": 78, "3_14": 79, "1_2": 80, "1_20_2": 81, "4_1_4": 82, "7_6": 83, "8_1_3": 84, "8_3_1": 85, "4_3": 86, "4_1_5": 87, "8_2_3": 88, "8_2_4": 89, "1_31": 90, "3_10": 91, "4_2_2": 92, "7_1": 93, "3_28": 94, "4_1_3": 95, "5_4": 96, "5_3": 97, "6_8_2": 98, "3_31": 99, "6_2": 100, "1_21": 101, "3_21": 102, "1_13": 103, "1_14": 104, "2_3_4": 105, "4_8_3": 106, "6_15_2": 107, "2_6": 108, "3_18_2": 109, "4_1_2_2": 110, "1_7": 111, "3_19": 112, "1_18": 113, "2_7": 114, "8_5_4": 115, "5_15_7": 116, "5_14": 117, "5_21": 118, "1_1": 119, "6_15_1": 120, "8_6_4": 121, "8_15": 122, "4_5": 123, "3_11": 124, "8_18": 125, "8_4_4": 126, "3_30": 127, "5_7_1": 128, "5_7_2": 129, "1_5": 130, "3_29": 131, "6_15_3": 132, "5_12": 133, "3_16": 134, "1_30": 135, "5_11": 136, "1_6": 137, "8_6_2": 138, "6_8_3": 139, "3_12": 140, "3_33": 141, "8_4_3": 142, "5_8": 143, "8_14": 144, "8_17": 145, "3_6": 146, "1_26": 147, "8_5_2": 148, "6_8_1": 149, "5_17": 150, "1_10": 151, "8_16": 152, "7_18": 153, "7_14": 154, "8_23": 155}
sings_rtsd = dict(zip(range(len(sings_rtsd)), [x.replace('_','.') for x in list(sings_rtsd.keys())]))

sings_input = ['3.24', '1.16', '5.15.5', '5.19.1', '5.19.2', '1.20.1', '8.23',
'2.1', '4.2.1', '8.22.1', '6.16', '1.22', '1.2', '5.16', '3.27',
'6.10.1', '8.2.4', '6.12', '5.15.2', '3.13', '3.1', '3.20', '3.12',
'7.14.2', '5.23.1', '2.4', '5.6', '4.2.3', '8.22.3', '5.15.1',
'7.3', '3', '2.3.1', '3.11', '6.13', '5.15.4', '8.2.1', '1.34.3',
'8.2.2', '5.15.3', '1.17', '4.1.1', '4.1.4', '3.25', '1.20.2',
'8.22.2', '6.9.2', '3.2', '5.5', '5.15.7', '7.12', '8.2.3',
'5.24.1', '1.25', '3.28', '5.9.1', '5.15.6', '8.1.1', '1.10',
'6.11', '3.4', '6.10', '6.9.1', '8.2.5', '5.15', '4.8.2', '8.22',
'5.21', '5.18']

Определим вспомогательные функции.

def parse_labeltxt(path):
    with open(os.path.join(path), 'r') as file:
        lines = file.readlines()
        labels = [sings_rtsd[int(x.split(' ')[0])] for x in lines]
        if labels.count('5.19.1')>1:
            labels.append('5.19.2')
        labels = list(set(labels))
        return labels

def rtsd2predict(labels):
    int_labels = []
    for sign in labels:
        if sign in sings_input:
            int_labels.append(sings_input.index(sign) + 1)
    return int_labels

Преобразуем yolo-предсказания в знаки дорожного движения RTSD, а затем в метки нашей задачи.

test_csv = pd.read_csv('test.csv', delimiter=',')
sample_solution = pd.read_csv('sample_solution.csv', delimiter=',')
labels_path = '/home/jovyan/yolov5/runs/detect/yolov5m6_signs_test/labels' #заменить на путь, где у вас хранятся запуски yolov5

predicted_labels = {}
for label in tqdm(os.listdir(labels_path)):
    predicted_labels[label[:-3]+'jpg'] = rtsd2predict(parse_labeltxt(os.path.join(labels_path, label)))

Последний шаг

Наш файл с решением должен содержать не названия файлов, а их id, которые мы можем получить из test.csv. Так же учтём, что у нас может быть максимум 8 знаков дорожного движения, то есть наш массив с предсказанием знаков дорожного движения нужно дополнить нулями так, чтобы его длина стала равна 8.

img2id = {}
for index, row in test_csv.iterrows():
    img2id[row['img']] = row['id']

for img in predicted_final.keys():
    img_id = img2id[img]
    signs = predicted_final[img] + ((8-len(predicted_final[img]))*[0])
    sample_solution[sample_solution['id']==img_id] = [img_id] + signs

sample_solution.to_csv('solution.csv', index=False)

Идеи по улучшению решения

• Разметить данные из обучающего набора для тех знаков дорожного движения, которые отсутствуют в RTSD. Этот способ точно даёт получить +0.04-0.1 к скору.

• Если взглянуть на обучающий набор данных и его разметку, то станет очевидно, что иногда нужно оставлять знаки дорожного движения, которые находятся далеко, а иногда нет. Возможно, что здесь есть какая-то корреляция с тем, какая именно локация используется. То есть, есть смысл поэкспериментировать  с confidence, а также визуализировать весь набор данных, чтобы понимать, в каких случаях в разметке были допущены ошибки.

• Эту задачу имеет смысл решать с конца. Нужно понять, насколько хорошо размечены тестовые данные. В ручном режиме разметить все тестовые изображения, а затем отправлять посылки, что позволит выяснить, что же система ожидает получить от нас. Таким образом, можно "подогнать" наше решение под ответ.

Итоги

Сама задача довольно интересная, но объём и качество разметки лишает всякого смысла попытки улучшить решение. На мой взгляд, данный подход к решению задачи - это вызов организатором соревнований. Надеюсь, что это побудит их детальнее продумывать свои кейсы. Реализация данного решения у автора заняла менее одного дня, но даже этот подход позволяет попасть в топ-5 лидерборда. Хочется верить, что эта статья будет полезна всем тем, кто только начинает свой соревновательный путь.

Участвуйте и побеждайте, всем удачи на чемпионатах и хакатонах!

Весь код доступен в Github.