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

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

Инструменты

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

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

Как работает YOLO?

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

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

Модель Даркнет имеет ряд преимуществ перед системами на основе классификаторов. Он смотрит на все изображение во время тестирования, поэтому его прогнозы основываются на глобальном контексте изображения. Он также делает прогнозы с одной оценкой сети, в отличие от таких систем, как R-CNN, которые требуют тысячи оценок для одного изображения. Это делает его чрезвычайно быстрым, более чем в 1000 раз быстрее, чем R-CNN, и в 100 раз быстрее, чем Fast R-CNN.

Бэкэнд глубокого обучения

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

Чтобы видеокадр работал с библиотекой MXNet, нужно превратить массив numpy в массив mxnet.

right_frame = mxnet.nd.array(frame)

Обработка и поиск групп

Затем была выбрана реализация YOLO yolo3_darknet_coco, предварительно обученная на наборе данных COCO. В этом наборе уже есть класс «человек», который позволил не тратить время на обучение нейронной сети.

net = model_zoo.get_model(
     'yolo3_darknet53_coco', 
     pretrained=True
)
x, img = data.transforms.presets.yolo.transform_test(
    right_frame, 
    short=512, 
    max_size=1024
)
class_IDs, scores, bounding_boxes = net(x)

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

Далее были выполнены следующие действия: если ограничивающие рамки двух целей пересекаются, добавить эту пару значений в массив target_intersection. Затем, зацикливаясь, мы проверяем, есть ли цель в обоих массивах. Если существует, создайте или обновите их группу.

target_intersection = [
   sorted((k1,k2))
   for k1,v1 in targets.items()
   for k2, v2 in targets.items()
   if k1 != k2 and bounding_box_intersection(v1, v2)
]
groups = []
for pair in target_intersection:
   found_group = False
   for group in groups:
      if set(pair).intersection(group):
         group.update(pair)
         found_group = True
         break
   if not found_group:
      groups.append(set(pair))

Результат

На рисунке цели ИИ выделены через ограничивающие рамки вместе с именем класса и вероятностью попадания в этот класс. Результат программы:

groups = find_groups(persons)
print(f"Result = {groups}")
Result = [6,5]

Результат = [6, 5] означает две группы людей по шесть человек в первой и пять во второй. Если результат был равен [], это означает, что на изображении нет групп людей.

Вывод

Были разработаны концепции для получения групп целей ИИ на кадре с использованием глубокого обучения и OpenCV. Этот метод быстр и точен, а также может использоваться для систем видеонаблюдения, видеомонтажа и т. Д.