РЖД: Диспетчерский контроль

Описание проекта

Проект представляет собой асинхронный API для обработки видео, разработанный с использованием FastAPI. Разработка осуществляется в рамках хакатона “Цифровой Прорыв”, кейс “РЖД: Диспетчерский контроль”. Наша команда состоит из пяти человек. API позволяет загружать видео, обрабатывать их с использованием модели машинного обучения, получать аннотированные изображения и видео, а также анализировать возможные риски. Реализованы функции работы с видео через WebSocket в реальном времени и обработки загруженных файлов на сервере.

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

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

Видео проходит через несколько этапов обработки: начиная от загрузки и предобработки кадров, заканчивая анализом рисков и выдачей финальных результатов пользователю. Предобработка включает преобразование изображений в нужный формат и нормализацию данных для подачи в модель. Модель машинного обучения, Faster R-CNN с backbone ResNet50, выполняет распознавание объектов, а библиотека Pillow используется для добавления аннотаций на изображения.

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

Архитектура проекта

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

• Модуль загрузки видео: отвечает за приём видео от пользователя и передачу его в очередь задач.
• Модуль предобработки: выполняет подготовку кадров для подачи в модель машинного обучения.
• Модуль обработки и аннотирования: использует модель Faster R-CNN для обнаружения объектов и добавления аннотаций.
• Модуль анализа рисков: определяет потенциальные опасности на основе расположения объектов на кадрах.

Наша команда

• Клепцова Любовь — ML-разработчик
• Погудалов Никита — Fullstack-разработчик
• Скопин Кирилл — Data Scientist
• Роль Никита — Data Scientist
• Лянгузов Кирилл — Дизайнер

Особенности проекта

• FastAPI — построение RESTful и WebSocket API.
• Jinja2 — рендеринг HTML-страниц.
• CORS настроен для разрешения запросов с любого источника.
• Обработка видео с использованием библиотеки OpenCV.
• Поддержка загрузки видео для офлайн-обработки с последующим предоставлением отчётов и анализа.
• WebSocket — обработка видеопотока в реальном времени.
• Модель обнаружения объектов на основе Faster R-CNN с backbone ResNet50.
• Хранение файлов и обработка задач на сервере.
• Анализ рисков — определение потенциальных рисков на основе пересечения обнаруженных объектов, например, когда человек находится слишком близко к технике.
• Микросервисная архитектура — модульная структура для повышения надёжности и масштабируемости.

Технические детали реализации

Предобработка видео

Перед подачей в модель машинного обучения, видео проходит этап предобработки, который включает:

1. Преобразование цветового пространства: кадры преобразуются в цветовое пространство RGB для корректной работы модели.
2. Изменение размера: изображения масштабируются до необходимого размера, чтобы соответствовать входным требованиям модели.
3. Нормализация: значения пикселей нормализуются, чтобы улучшить производительность модели.

Модель машинного обучения

Для обнаружения объектов используется модель Faster R-CNN с backbone ResNet50, предобученная на наборе данных из 27 классов. Модель модифицируется для работы с пользовательскими классами, включающими объекты, которые необходимо обнаружить на железнодорожных станциях, например, “человек”, “техника”, “вагон” и другие.

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

Для обучения модели распознавания объектов используется предобученная модель Faster R-CNN с backbone ResNet50. Обучение модели включает следующие шаги:

1. Подготовка данных: используется набор данных COCO, который хранится в папке dataset с аннотациями в файле coco.json. Данные разделяются на обучающую, валидационную и тестовую выборки в пропорциях 80%, 10%, 10% соответственно.
2. Создание пользовательского датасета: реализован класс CustomDataset, основанный на CocoDetection. Этот класс позволяет обрабатывать изображения и аннотации, конвертируя их в формат, совместимый с моделью.
3. Преобразования данных: к изображениям применяются преобразования, такие как преобразование в тензоры и случайное горизонтальное отражение, что помогает улучшить обобщающую способность модели.
4. Определение модели: модель Faster R-CNN с backbone ResNet50 модифицируется для работы с количеством классов, соответствующим набору данных (включая фон). Для этого используется класс FastRCNNPredictor.
5. Обучение: модель обучается на протяжении 50 эпох с использованием стохастического градиентного спуска (SGD) в качестве оптимизатора. Обучение проводится на GPU. В процессе обучения рассчитываются потери (loss) на обучающей и валидационной выборках, и сохраняются контрольные точки модели.
6. Валидация и тестирование: после каждой эпохи проводится валидация модели, а по завершении обучения - тестирование на тестовой выборке. Для каждой эпохи рассчитываются и сохраняются значения потерь, которые затем визуализируются на графике.

Аннотирование и анализ рисков

После получения предсказаний от модели, кадры аннотируются с использованием библиотеки Pillow. Для каждого обнаруженного объекта рисуется рамка и добавляется метка с названием класса. Далее кадры анализируются на предмет потенциальных рисков, таких как пересечение опасных зон между объектами (например, человек рядом с техникой).

Установка и запуск

1. Клонируйте репозиторий:

   git clone https://gitflic.ru/project/dragonnp/dispatch-control-rzhd.git
   

2. Перейдите в папку проекта и установите зависимости:

   cd dispatch-control-rzhd
   pip install -r requirements.txt
   

3. Запустите Docker-контейнер:

   docker-compose up -d
   

4. Откройте браузер и перейдите по адресу http://127.0.0.1:8000/.

Эндпоинты

• GET / — Главная страница с HTML интерфейсом.
• GET /process-video-realtime — Страница для обработки видео в реальном времени.
• POST /upload/ — Загрузка видео для офлайн-обработки.
• GET /status/{task_id} — Проверка статуса задачи обработки.
• GET /report/{task_id} — Получение отчёта по задаче.
• GET /video/{task_id} — Получение обработанного видео.
• GET /risks/{task_id} — Получение анализа рисков по задаче.
• GET /download/edited/ — Скачивание обработанного видео.
• WEBSOCKET /video-stream — Эндпоинт для обработки видео в реальном времени.

Функции обработки видео

• decode_frame(base64_data) — Декодирование кадра из base64.
• process_frame(frame) — Обработка кадра, добавление рамки и аннотаций с использованием модели машинного обучения.
• encode_frame(frame) — Кодирование кадра в формат base64.

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

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

Зависимости

• FastAPI — фреймворк для создания веб-приложений.
• uvicorn — сервер ASGI для запуска FastAPI.
• OpenCV — библиотека для обработки изображений и видео.
• numpy — для обработки изображений в формате массива.
• torch и torchvision — для работы с моделью обнаружения объектов Faster R-CNN.
• PIL (Pillow) — для работы с изображениями.
• shapely — для геометрических операций, используемых при анализе рисков.
• Jinja2 — шаблонизатор для рендеринга HTML.
• uuid — для генерации уникальных идентификаторов задач.

Обученная модель

https://disk.yandex.ru/d/2TqQu74dPqjzGw

Поместить в папку ai_model