Привет, Хабр! На связи команда «Чёрная уточка» из управления валидации Альфа-Банка (Светлана Хлыбова, Сергей Комаров, Буда Вампилов, Камиль Шакиров и Алексей Безручко). И сегодня мы расскажем о нашем первом (да ещё и вполне успешном) опыте участия в Дальневосточном федеральном окружном хакатоне проекта «Цифровой прорыв. Сезон: Искусственный интеллект. 2024». Нам приглянулся кейс от РЖД на тему компьютерного зрения, в котором было необходимо научить компьютер выявлять технологические нарушения по видеозаписям.

«Цифровой прорыв. Сезон: Искусственный интеллект. 2024» — это самый масштабный хакатон в стране, и в нём участвуют как банки (что уже традиционно), так и государственные учреждения (например, Министерство здравоохранения РФ и СГМУ им. В. И. Разумовского), Центральный банк РФ и различные крупные компании, такие как РЖД, Rutube, VK и другие. Например, СГМУ им. В. И. Разумовского совместно с министерством здравоохранения предоставили интересный кейс создания продукта для диагностики нарушения дыхания во время сна, но нас заинтересовало задание от РЖД — автоматизация выявления технологических нарушений.

Полные требования выглядят так: «Участникам необходимо создать программный модуль, который с помощью ИИ проанализирует видеозаписи с нагрудных регистраторов для определения правильного использования СИЗ, спецодежды и определения технологических изменений».

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

Как итог, ещё до начала хакатона у нас уже был первоначальный план действий и идея baseline-решения — некоторого несложного прототипа модели машинного обучения, который позволит более или менее удовлетворительно решить задачу. Мы решили анализировать отдельные кадры на предмет наличия на них нарушения (задача бинарной классификации) или, если потребуется, на те или иные типы нарушений (задача мультиклассовой классификации).

Для этого можно взять готовую нейросетевую модель (например, ResNet), обученную для мультиклассовой классификации изображений на очень большом наборе данных, и дообучить её, добавив вместо выходных слоёв сети несколько своих новых слоёв, параметры которых уже оптимизировались под конкретную задачу. Именно так мы и решили поступить, т.к. собрать такое решение на базе библиотеки TensorFlow достаточно просто. Ну а уже в ходе самого хакатона мы планировали придумать и реализовать другие варианты решений, которые можно ещё и ансамблировать, чтобы получить более высокую точность.

Старт кодинга

Мы получили для обучения около 30 видеофайлов, снятых на нагрудный видеорегистратор, и файл с разметкой, в котором указаны название видео и время начала нарушения (каждое нарушение идёт с новой строки, то есть в данном видео 4 нарушения).

Нарушения были зафиксированы лишь в 16 из предоставленных видеозаписей. «Не густо», — подумали мы, но с этим нужно было как-то работать.

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

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

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

3DCNN — продвинутый подход

Что ж, после того как стало понятно, что дополнительных данных не будет, мы убрали в стол разработку ещё одного, более сложного алгоритма — трёхмерной свёрточной нейронной сети (3DCNN), над которым некоторое время работали параллельно.

3DCNN основана на концепции сверточных нейронных сетей (CNN), но с добавлением временного измерения, что позволяет учитывать не только контекст отдельных кадров, но и контекст из их порядка в видеоряде. Архитектура 3DCNN представлена на картинке:

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

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

Однако такая сеть будет иметь много параметров, а для их качественной оптимизации в ходе обучения понадобится много данных, которыми мы не располагали. Можно было бы попробовать собрать очень много внешних данных, например, видео с железной дороги на YouTube, и обучить на них модель с данной архитектурой для решения какой-то абстрактной задачи, а добавив к модели дополнительных слоев и попробовать дообучить модель на наших данных под решаемую задачу. 

Но у нас не было достаточно вычислительных ресурсов, хотя важнее, что мы бы просто не успели, потому что мероприятие длилось только 48 часов. А наш хакатон проходил ещё и во Владивостоке, что внесло небольшие неудобства из-за разницы во времени относительно Москвы, поскольку по Москве старт кодинга был в 10 часов утра в пятницу, из-за чего мы потеряли около 7 часов, так как все работали, а стоп-кодинг был в 7 утра воскресенья. Стоит упомянуть, что данное мероприятие — самое масштабное соревнование по искусственному интеллекту в России, в ходе которого проводилась серия хакатонов в разных городах по всей стране.

Как следствие, возвращаемся к первоначальной идее. И первый этап в ней — это, конечно же, подготовка данных. 

Baseline

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

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

Таким образом, процесс нарезки видео на кадры можно поделить на следующие шаги:

1. Получить метаданные о видео.

2. Задать количество кадров в секунду для сохранения.

3. Получить список моментов времени для сохранения с информацией о наличии нарушения в каждый момент.

4. В цикле пробежаться по каждому кадру видео: если время кадра соответствует времени для стоп-кадра, то сохраняем этот кадр как отдельное изображение (и не забываем записывать информацию по всем сохраненным кадрам в датафрейм).

Ниже функция для нарезки видео на кадры.

import pandas as pd
import imageio.v3 as iio
import os
from pathlib import Path

def video_slicer(video_file, video_name, video_num, df, delta, save_dir, 
                 make_dir=False):
  """
  Функция выделяет в видео отрезки с нарушениями и за delta секунд до нарушения, 
  создаёт кадры и сохраняет их в заданную папку.
  :param func:
          video_file - путь до видео
          video_name - название видео
          video_num - номер видео
          df - данные по нарушениям для конкретного видео
          delta - сколько секунд отступаем до нарушения
          make_dir - сохраняем в одну папку или в разные
          save_dir - куда сохраняем
      :return:
          file_names - дафафрейм с данными о сохраненных кадрах
  """
  file_names = pd.DataFrame()
  # Создаём папку для сохранения результатов.
  if not os.path.isdir(save_dir):
    os.mkdir(save_dir)  
  filename = save_dir  
  if make_dir:
      filename, _ = os.path.splitext(video_name)
      filename = os.path.join(save_dir, filename + "_iio")
      # создаем индивидуальную папку для сохранения результатов
      if not os.path.isdir(filename):
          os.mkdir(filename)
  # Получаем данные о видео.
  file_p = os.path.join(os.getcwd(), 'data', 'train', 'videos', video_name)
  metadata = iio.immeta(file_p, exclude_applied=False, plugin="pyav")
  # FPS видео.
  fps = metadata['fps']

  # Если SAVING_FRAMES_PER_SECOND выше видео FPS, то установите его на FPS (как максимум).
  saving_frames_per_second = min(fps, SAVING_FRAMES_PER_SECOND)

  # Получить список моментов времени для сохранения.
  saving_frames_durations, number_of_violent, flags, descr_list = \
        get_saving_frames_durations(metadata,
                                    saving_frames_per_second,
                                    df,
                                    delta)
  # Запускаем цикл.
  count = 0
  save_frame_count = 0
  for frame in iio.imiter(os.path.join(os.getcwd(), 'data', 'train', 'videos',
                          video_name), plugin="pyav"):
    # Получаем продолжительность, разделив количество кадров на FPS.
    frame_duration = count / fps
    try:
      # Получаем момент времени из списка для сохранения кадра.
      closest_duration = saving_frames_durations[0]
    except IndexError:
      # Список пуст, все нужные моменты времени сохранены.
      break
    if frame_duration >= closest_duration:
      # Если ближайший момент времени меньше или равен моменту времени кадра, 
      # сохраняем фрейм.
      frame_duration_formatted = format_timedelta(timedelta(seconds=frame_duration))
      # Получаем инфо по кадру, который будем сохранять.
      count_violation = number_of_violent[save_frame_count]
      flag = flags[save_frame_count]
      descr = descr_list[save_frame_count]
      frame_name = f"video_{video_num}_frame_{frame_duration_formatted}\
                   _flag_{flag}_violation_{count_violation}.jpg"
      # Сохраняем кадр.
      iio.imwrite(os.path.join(filename, frame_name), frame)

      # Добавляем инфо о сохраненном кадре в датафрейм.
      file_df = pd.DataFrame({'video_name': video_name,
                                    'video_num': video_num,
                                    'image_name': frame_name,
                                    'frame_duration': frame_duration_formatted,
                                    'target': flag,
                                    'description': descr}, index=[0])
      file_names = pd.concat([file_names, file_df], ignore_index=True)
      
      save_frame_count += 1

      # Удалить точку продолжительности удаляем время сохраненного кадра из списк
      try:
        saving_frames_durations.pop(0)
      except IndexError:
        pass  

    # Увеличить количество кадров.
    count += 1
  return file_names

После того, как данные были готовы, приступили к обучению нашей базовой версии модели.

В качестве основы здесь взяли архитектуру ResNet50 (сверточная НС с 50-ю слоями), добавили слой GlobalAveragePooling2d для преобразования выходного тензора в вектор и два полносвязных слоя с функцией активации ReLu.

from tensorflow.keras import Sequential
from tensorflow.keras.layers import (Conv2D, 
                                     Flatten, 
                                     Dense, 
                                     AvgPool2D, 
                                     GlobalAveragePooling2D)
from tensorflow.keras.applications.resnet import ResNet50


backbone = ResNet50(input_shape=input_shape,
                    weights='imagenet',
                    include_top=False)

backbone.trainable = False

model = Sequential()
model.add(backbone)
# новая голова
model.add(GlobalAveragePooling2D())
model.add(Dense(units=1000, activation='relu'))
model.add(Dense(units=1000, activation='relu'))
model.add(Dense(units=1, activation='sigmoid'))
model.compile(loss=BinaryCrossentropy(), optimizer=optimizer,
              metrics=['accuracy'])

Получили accuracy = 0.77 на валидационной выборке. Уже что-то.

Сегментация специфичных объектов и эвристик

Конечно же, такой вариант нас не особо устроил, и мы решили, что можно сделать модель сегментации отдельных объектов в кадре (человек, вагон, рельсы, элементы одежды и т. д.), затем наложить на полученные результаты эвристики, исходя из которых определялись бы случаи нарушений. Например, если обнаружены объекты «человек» и «вагон» и человек находится под вагоном, при этом у колёс вагона нет специальных тормозных устройств (так называемых «башмаков»), то в кадре имеется нарушение.

Естественно, что для этого нам нужно больше данных и детальная разметка для обучения такой модели. Поэтому решили добавить картинок из интернета к имеющимся кадрам и детально разметить всё при помощи Roboflow. Кроме того, Roboflow предоставляет возможность обучить модель на собственном сервере, что на первый взгляд нам показалось отличным решением, так как с вычислительными ресурсами у нас была напряжёнка. После кропотливой работы и одной бессонной ночи мы получили данные и запустили на обучение модель YOLO. Результат нас порадовал.

«Как круто получилось!» — радостно возгласили мы, не ожидая подвоха. А он был. Оазалось, что с Roboflow нельзя взять и скачать обученную модель, можно лишь обратиться к ней через API, а как мы помним, одним из требований заказчика было отсутствие необходимости выхода в интернет. 

Не найдя способа достать полученную модель, мы решили на своем железе обучить ту же YOLO с помощью библиотеки Ultralytics (это Open Source библиотека для компьютерного зрения, которая включает в себя предобученные модели семейства YOLO с удобным и понятным API для быстрого запуска и решения задач сегментации и классификации изображений). Для запуска обучения потребуется написать всего лишь несколько строчек кода.

from roboflow import Roboflow
from ultralytics import YOLO

# Загружаем данные для обучения с Roboflow
rf = Roboflow(api_key="**************") # здесь нужно вставить свой ключ 
project = rf.workspace("buda-vampilov").project("violation_classification")
version = project.version(1)
dataset = version.download("folder")

# Загрузка модели
model = YOLO("yolov8n-cls.pt")

# Обучение модели
results = model.train(data="violation_classification-1", epochs=1000, 
		       imgsz=224,fliplr=0.5, degrees=7.5, 
		       perspective=0.000025)

Но время уже поджимало, оставалось меньше суток, поэтому самое логичное решение — обучаться просто на задачу бинарной классификации нарушений в кадре. 

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

Приложение

Как выше описано, в требованиях кейсодержателя был пункт о реализации не только модели, но и программного интерфейса. Над ним мы, естественно, работали параллельно с разработкой алгоритмов обработки видео. Прототип приложения был максимально простым  — есть всего 3 кнопки: Выбор видео, Детектирование и Отчет. Здесь мы не прогадали — на защите пользовательский интерфейс судьи оценили по максимуму, поставив 3 балла.

В качестве детектора нарушений взяли модель бинарной классификации YOLO, так как ансамблирование со слабыми решениями не принесло существенного эффекта. В итоге получили время обработки в 7 раз быстрее длительности видео (и это на CPU). 

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

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

Финал

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

На защиту предоставляется всего 5 минут и 5 минут на вопросы кейсодержателей и технических экспертов. Каждую команду жюри оценивало по следующим критериям. За каждый аспект можно было получить максимум 3 балла. Каждый из членов жюри выставлял свои оценки. Итоговый балл формировался как сумма баллов от каждого из жюри. 

Мы подготовили презентацию, в которой коротко рассказали об этапах разработки и выборе итогового решения, акцентировали внимание на достоинствах нашего продукта и подготовили демо, чтобы показать, что приложение действительно работает и это именно то, что просили заказчики. Как итог — выступление команды все члены жюри оценили на максимально возможные 3 балла.

Каждую команду спрашивали о том, какие нарушения может детектировать их приложение и сколько видео за раз можно обработать. По части детектирования нарушений наше итоговое решение отчасти не соответствовало желанию заказчика, поскольку в качестве детектора нарушений мы взяли модель бинарной классификации YOLO, обладающей большей точностью из всех наших вариантов решений, но определяющей лишь факт наличия нарушения, а не его конкретный тип. Поэтому полноту выполнения требований к решению судьи оценили на 2 балла. Про количество обрабатываемых видео за раз в техническом задании ничего не говорилось, и мы реализовали выбор и обработку одного видео, однако некоторые команды реализовали множественную обработку видео, и на этом фоне наше решение выглядело менее выигрышным. 

Ещё одним слабым местом оказался наш репозиторий: мы сильно ударились в различные эксперименты и к моменту стоп-кодинга не успели сделать описание всех файлов, и в большинстве файлов не было комментариев. Поэтому за запускаемость кода мы получили всего 1 балл. Что ж, зато мы получили хороший урок на будущее: документировать все сразу, а не откладывать на потом!

Да, не всё пошло по плану, и мы не успели реализовать то, что хотели. Но мы честно рассказали про это, последовательно изложив ход наших мыслей, и сделали упор на достоинства выбранного решения и способы его доработки. Это и стало ключом к призовому второму месту. К слову, команды, занявшие первое и третье места, были сконцентрированы на реализации единственного варианта решения, которое оказалось очень схоже с одной из наших ключевых идей, а именно — детектирование отдельных объектов в кадре и определение нарушений по их комбинациям. 

Наверное, благодаря тому, что они не пробовали реализовать много разных идей, а все были сконцентрированы на одном, у них получилось довести решение до конца. Правда, у команды на третьем месте хромала точность, нарушение определялось правильно с вероятностью от 20 до 80%, и пользовательский интерфейс был не такой удобный как у нас. Ну а первое место было занято по праву, у ребят получился удобный пользовательский интерфейс и хорошая точность определения нарушений. Кроме того, они организовали отличный репозиторий. У них и правда есть чему поучиться.

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

И здесь мы сделали для себя вывод, что в первую очередь стоит отталкиваться от вычислительных ресурсов: если нет мощного и стабильно работающего сервера, то не стоит тратить драгоценное время и силы на попытку реализации чего-то слишком сложного (как например наша попытка с 3DCNN). И однозначно точно не стоит слишком сильно закапываться в эксперименты. Но интересно узнать какую стратегию выбрали бы вы и какое решение?

Подводя итог, хотим сказать, что, в конце концов, самое главное это не призовое место (хотя, безусловно, это крайне приятно), а опыт, который мы получили. В особенности ценными являются допущенные ошибки, из которых мы усвоили ценные уроки. Надеюсь, что статья была полезной и вы выписали себе некоторые тактические приемы, которые используете, когда будете участвовать в хакатонах.

Теперь в планах только первое место! Подробности нашей разработки вы найдете здесь.