Мое решение состоит из двух нейросетей. Первая нейросеть нужна для детекции людей, а вторая -- для детекции нарушений техники безопасности.
Пример работы первой нейросети:
Пример работы второй нейросети:
Обе нейросети основаны на популярной архитектуре Yolo V5.
Поскольку я использую архитектуру YOLO, а вся разметка дана в формате COCO, то их надо было конвертировать между собой. Для этого у меня есть 2 скрипта в соответствующих директориях: coco_to_yolo и yolo_to_coco. Краткая инструкция по их использованию:
- coco_to_yolo. Положите .json файл с разметкой в формате COCO в директорию coco_to_yolo (рядом с файлом coco_to_yolo.py), укажите в скрипте coco_to_yolo.py в переменной input_file название вашего .json файла. Запустите скрипт: python coco_to_yolo.py. В директории output_yolo появятся файлы с разметкой в формате YOLO.
- yolo_to_coco. Положите .txt файлы в формате YOLO в директорию input_yolo. Рядом с ними положите соответствующие им изображения. В скрипте yolo_to_coco.py можете указать имя выходного .json файла в переменной output_file. Запустите скрипт: python yolo_to_coco.py. В директории с coco_to_yolo.py появится выходной .json файл с разметкой в формате COCO.
Стоит указать, что в yolo_to_coco я использовал чужой код из репозитория, указанного в конце данного readme.
Основана на Yolo v5 (Large). Входная разметка - от организаторов хакатона. При обучении и инференсе используются изображения с разрешением 1280. Все что нужно для обучения лежит в person_detection/train. Все что нужно для инференса лежит в person_detection/inference. Подробности в README.md в соответствующих директориях репозитория.
Основана на Yolo v5 (Medium). Входная разметка - от меня лично, размечал руками. При обучении и инференсе используются изображения с разрешением 640. Все что нужно для обучения лежит в violation_detection/train. Все что нужно для инференса лежит в violation_detection/inference. Подробности в README.md в соответствующих директориях репозитория.
Как я понял, в Yolo v5 по дефолту фиксируются все seed'ы. Однако нужно проверить значение переменной окружения RANK в вашей системе:
printenv | grep RANK
Если вам ничего не вывелось (скорей всего будет так), значит у вас этой переменной нет, и все ок. В этом случае мои результаты должны быть воспроизводимы. Если же она у вас есть и не равна -1, то для воспроизводимости вам нужно будет подредактировать код Yolo v5 в файле train.py: строчку
init_seeds(1 + RANK)
надо будет заменить на
init_seeds(0)
Это можно сделать, если зайти в контейнер в -it режиме. Однако надеюсь, это не потребуется. Если что, пишите мне, я помогу.
Архитектура Yolo v5: https://github.com/ultralytics/yolov5. Лицензия: GNU General Public License v3.0 (разрешает коммерческое использование)
Скрипт для перегонки YOLO формата разметки в COCO формат: https://github.com/Taeyoung96/Yolo-to-COCO-format-converter. Лицензия: MIT License (разрешает коммерческое использование)
20.04.1-Ubuntu
NVIDIA: Driver Version: 470.74 CUDA Version: 11.4