Как построить детектор бейсбола с помощью Detectron2

Обучите обнаружение объектов Detectron2 с помощью пользовательских наборов данных.

Введение

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

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

В этом посте я собираюсь показать, как я создаю собственный набор данных о бейсболе, обучаю модель обнаружения объектов и применяю ее к видео о бейсболе с помощью Detectron2 (https://github.com/facebookresearch/detectron2).

Детектор бейсбола был построен в три этапа, которые будут подробно обсуждаться в этой публикации:
1. Создайте собственный набор данных бейсбола в формате COCO
2. Поиграйте с Detectron2 и обучите модель в Colab
3. Загрузите видео / изображение и примените обученную модель для обнаружения.

Создайте собственный набор данных бейсбола в формате COCO

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

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

Затем я использовал Labelimg, чтобы вручную наклеить этикетку на бейсбол. Labelimg - удобный инструмент для маркировки объектов. Следуя инструкции в репо, его можно легко установить и использовать.

Выберите «PascalVOC», это формат аннотации по умолчанию.
Откройте папку, содержащую изображения
Установите папку для сохранения аннотации
Обозначьте мяч и сохраните аннотацию

Мне потребовалось около 20 минут, чтобы вручную пометить 120 изображений, после чего аннотации были сохранены в папке, которую вы установили в xml.

Затем я использовал этот скрипт voice2coco.py из репозитория Tony607 на GitHub для преобразования файлов PascalVOC xml в один файл JSON в формате COCO.

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

Обучите модель с помощью Detectron2 в Colab

Я модифицировал красивый и понятный учебник Colab notebook от Detectron2 и обучил модель.

Сначала я подключил свой Google Диск к ноутбуку и загрузил в него созданный набор данных.

from google.colab import drive
drive.mount('/content/gdrive', force_remount=True)
root_dir = 'gdrive/My Drive/Colab Notebooks/'
base_dir = root_dir + 'baseball_detection'

Затем я последовал инструкциям в ноутбуке по установке и настройке detectron2.

# install dependencies:
!pip install pyyaml==5.1 pycocotools>=2.0.1
import torch, torchvision
print(torch.__version__, torch.cuda.is_available())
!gcc --version
# opencv is pre-installed on colab
# install detectron2: (colab has CUDA 10.1 + torch 1.5)
# See https://detectron2.readthedocs.io/tutorials/install.html for instructions
assert torch.__version__.startswith("1.5")
!pip install detectron2==0.2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.5/index.html
# Some basic setup:
# Setup detectron2 logger
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()
# import some common libraries
import numpy as np
import os, json, cv2, random
from google.colab.patches import cv2_imshow
# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog

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

from detectron2.data.datasets import register_coco_instances
register_coco_instances("ball", {}, path/'image'/'output.json', path/'image')

затем позвонил

ball_metadata = MetadataCatalog.get("ball")
dataset_dicts = DatasetCatalog.get("ball")

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

for d in random.sample(dataset_dicts, 3):
   img = cv2.imread(d["file_name"])
   visualizer = Visualizer(img[:, :, ::-1], metadata=ball_metadata,    scale=0.5)
   out = visualizer.draw_dataset_dict(d)
   cv2_imshow(out.get_image()[:, :, ::-1])

Теперь я готов к тренировкам!

Я настраиваю предварительно обученную модель R50-FPN Mask R-CNN COCO на наборе данных бейсбола. Конфиг и вес получил методом model_zoo. не забудьте изменить cfg.MODEL.ROI_HEADS.NUM_CLASSES на 1, так как у меня сейчас только один класс. Мне потребовалось несколько минут, чтобы запустить 2000 итераций на Colab.

from detectron2.engine import DefaultTrainer
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("ball",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")   # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR
cfg.SOLVER.MAX_ITER = 2000    # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  # only has one class (ball)
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()

Я вижу результаты тренировки в тензорной доске

# Look at training curves in tensorboard:
%load_ext tensorboard
%tensorboard --logdir output

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

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5   # set a custom testing threshold for this model
cfg.DATASETS.TEST = ("ball", )
predictor = DefaultPredictor(cfg)

выбрали несколько тестовых изображений для тестирования модели

im = cv2.imread('gdrive/My Drive/Colab Notebooks/baseball_detection/test/1.png')
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1],
  metadata=ball_metadata,
  scale=1)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2_imshow(out.get_image()[:, :, ::-1])

и посмотрите результат!

Загрузите видео / изображение и примените обученную модель для обнаружения

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

Я написал сценарий, чтобы читать видео с помощью OpenCV кадр за кадром и обнаруживать бейсбольный мяч, используя модель для каждого кадра. Итак, теперь у меня есть информация об обнаруженном бейсболе, такая как местоположение в каждом кадре, для дальнейшего анализа. Я также могу вывести видео об обнаружении бейсбольного мяча!

Будущая работа:

Применить отслеживание объекта, чтобы отслеживать мяч

Обнаружение и отслеживание бейсбольных мячей с помощью Detectron2 и SORT
Как я отслеживаю бейсбольные мячи, когда на видео показано несколько бейсбольных мячей. todatascience.com

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

Спасибо за чтение. Отзывы и предложения приветствуются!