BallDetection

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

Подготовка датасета

Для обучения нейронной сети для детекции футбольного мяча был выбран датасет Open Images V6, так как в разметке присутствует класс футбольного мяча. Чтобы скачать изображения и разметку с нужным классом, из папки open_images запускаем download_annotations.sh, затем preprocess_dataset.py, затем download_images.py.

Для подготовки реального датасета изначально необходимо в папке real_data запустить скрипт create_dirs.sh, далее в папку real_data/data/annotations поместить разметку, а в real_data/data/videos видео. После этого можно проверить разметку, запустив check_annotations.py, получив видео с отрисованной разметкой в папке debug_videos.

После проверки разметки запускаем create_real_dataset.py, а затем replace_to_open_images.py для перемещения в папку с общим датасетом.

Обучение сети

Для воспроизведения эксперимента необходимо запустить скрипт train.sh

Экспорт модели в ONNX

Запускаем export.sh с указанием пути до весов модели, затем помещаем onnx модель в папку detection/model и называем её best.onnx

Инференс

Подготовка Docker образа

Запускаем build_docker.sh

Тестирование сети

Запускаем detect.sh с указанием пути до видео, на котором нужно детектировать мяч.

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