Для многих людей обнаружение объектов на изображении кажется каким-то колдовством, однако существует огромное количество способов распознать на изображении определенные объекты. Одно только перечисление займет всю статью, а уж их сравнение вполне тянет на серьезное исследование.
Для начала поставим задачу. Нужно посчитать количество людей, имеющихся на изображении. Делать это будем в бесплатном colab. Считать людей будем на публичном наборе изображений с людьми: https://www.cis.upenn.edu/~jshi/ped_html/PennFudanPed.zip.
Известные мне системы обнаружения объектов используют фреймворк Pytorch, поэтому право первой "брачной ночи" очереди за предобученными моделями этого фреймворка.
Объявляем необходимые библиотеки. Объяснение функций torchvision будет непосредственно в коде:
Загрузим датасет и распакуем его. Из всего датасета нам нужно только папка PennFudanPed/PNGImages:
Получим веса нашей модели (третья объявленная функция) и саму модель (вторая объявленная функция), в модель передадим веса и нижний уровень точности определения объекта:
Инициализируем модель с помощью функции eval:
Создадим функцию трансформации изображения под требования предобученной модели:
Первой объявленной функцией была read_image, которая загружает изображение из файла, и конвертирует его в тензор (что-то похожее на массив в numpy). Вот как она работает:
Полученный тензор имеет целочисленный тип данных, а входящий слой модели ожидает дробный, поэтому пропускаем тензор через функцию preprocess, которую мы создали двумя шагами ранее, а также добавляем размерность тензора с помощью квадратных скобок:
Теперь подготовленные данные передаем в модель и получаем предсказание (нас интересует только первая часть предсказания):
Нам нужны метки классов и рамки объектов, причем метки должны иметь значение 1 (person), остальные классы должны быть пропущены:
Чтобы получить рамки только людей мы используем маску, а чтобы нарисовать эти рамки мы объявили функцию draw_bounding_boxes (4 сверху). Рамки будут красными, шириной два пикселя. Но результат мы получим в виде тензора, который конвертируем в PIL-формат с помощью функции to_pil_image (5):
Выведем изображение с помощью функции imshow библиотеки matplotlib (последнее объявление):
На первом изображении модель обнаружила два человека:
На втором только одного, хотя глазами мы видим двух дюдей:
Предобученная модель без дополнительной настройки не очень хорошо распознает перекрывающиеся объекты, что мы видим и на третьем рисунке:
Для подсчета обнаруженных людей нам достаточно выполнить:
Задача нами решена, правда есть некоторые "шероховатости", но в целом предобученная модель вполне себе справляется с обнаружение людей, если они не перекрыты друг другом.