Узнайте, как использовать модель MobileNet SSD, предварительно обученную на наборе данных COCO в Python, для работы на периферийных устройствах с полным кодом и не максимальным подавлением.
Tensorflow недавно выпустил свой API обнаружения объектов для Tensorflow 2, который имеет очень большой зоопарк моделей. Однако они предоставили только одну модель SSD MobileNet v1 с Tensorflow lite, которая описана здесь. В этом сообщении в блоге они предоставили коды для его запуска на устройствах Android и IOS, но не на периферийных устройствах. С ростом популярности периферийных устройств и выпуском OpenCV наборов пространственного ИИ я хотел восполнить недостающий пробел. Так что именно это мы и увидим здесь.
Требования
Для запуска этого не требуется установка API обнаружения объектов Tensorflow. Для его запуска достаточно простой установки Tensorflow вместе с OpenCV для обработки изображений.
pip install tensorflow pip install opencv
Я использовал некоторый код, предоставленный в API обнаружения объектов, чтобы упростить работу, но нет необходимости беспокоиться об этом, так как не требуется его явная установка, и вы можете найти весь код здесь.
Предварительно обученную модель можно скачать из блога Tensorflow здесь или она также предоставляется вместе с кодом.
Этикетки
Наш первый шаг к тому, чтобы сделать метки в формате, который потребуется позже, - это вложенный словарь с идентификатором и меткой внутри. Мы воспользуемся текстовым файлом карты меток, поставляемым с моделью, и преобразуем его в требуемый формат.
Примечание. - Ярлыки с «???» необходимо игнорировать, и эти индексы пропускаются, кроме первого. Таким образом, каждая метка сдвигается на одну позицию назад. Например, метка «человек» находится в первой строке, но ей будет присвоена метка 0 и так далее.
Это можно найти в комментариях в его реализации на Java.
// Модель SSD Mobilenet V1 предполагает, что класс 0 является фоновым классом
// в файле меток и класса метки начинаются с 1 до number_of_classes + 1,
// пока outputClasses соответствует индексу класса от 0 до number_of_classes
Это можно сделать с помощью следующего кода:
def create_category_index(label_path='path_to/labelmap.txt'):
f = open(label_path)
category_index = {}
for i, val in enumerate(f):
if i != 0:
val = val[:-1]
if val != '???':
category_index.update({(i-1): {'id': (i-1), 'name': val}})
f.close()
return category_index
Здесь, чтобы игнорировать первую строку, я использую оператор if и сохраняю i-1. Это создает словарь, как показано ниже.
В нем будет 80 строк с ключами до 89.
Переводчик TfLite
Когда ярлыки готовы, давайте разберемся с интерпретатором TfLite и тем, как мы можем получить результаты.
Инициализировать интерпретатор
import tensorflow as tf interpreter = tf.lite.Interpreter(model_path="path/detect.tflite") interpreter.allocate_tensors()
Просто загрузите правильный путь к вашей модели tflite и назначьте тензоры.
Детали ввода и вывода
Чтобы получить информацию о вводе и выводе, напишите:
input_details = interpreter.get_input_details() output_details = interpreter.get_output_details()
Теперь давайте изучим их, чтобы увидеть, какие типы входов дать и какие выходы мы получим.
Детали ввода представляют собой список только из 1 элемента, который является словарем, как показано ниже.
Здесь мы видим фигуру ввода [1, 300, 300, 3]. Кроме этого, требуется, чтобы входное изображение имело тип данных np.uint8.
Выходные данные представляют собой список из 4 элементов, каждый из которых содержит словарь, как и входные данные. Каждый вызов возвращает 10 результатов, причем первый элемент хранит ограничивающие прямоугольники, второй элемент - классы обнаружения, третий элемент - оценки обнаружения и, наконец, последний элемент - количество возвращенных обнаружений. Возвращаемые ограничивающие рамки нормализованы, поэтому их можно масштабировать до любого входного измерения.
Чтобы получить выходные данные, нам нужно прочитать изображение, преобразовать его в RGB, если используется OpenCV, соответствующим образом изменить его размер и вызвать интерпретатор после установки тензора с входным кадром. Затем требуемые значения могут быть достигнуты с помощью get_tensor функции интерпретатора.
import cv2
img = cv2.imread('image.jpg')
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_rgb = cv2.resize(img_rgb, (300, 300), cv2.INTER_AREA)
img_rgb = img_rgb.reshape([1, 300, 300, 3])
interpreter.set_tensor(input_details[0]['index'], img_rgb)
interpreter.invoke()
de_boxes = interpreter.get_tensor(output_details[0]['index'])[0]
det_classes = interpreter.get_tensor(output_details[1]['index'])[0]
det_scores = interpreter.get_tensor(output_details[2]['index'])[0]
num_det = interpreter.get_tensor(output_details[3]['index'])[0]
Использование кода обнаружения объектов для рисования
Для визуализации я использовал код Python, доступный здесь, который можно использовать не только для рисования ограничивающих рамок, но также при необходимости ключевых точек и масок экземпляров. Нам нужно передать изображение для рисования, ограничивающие рамки, обнаруженные классы, оценки обнаружения и словарь меток. Помимо этого, мы также устанавливаем нормализованные координаты в значение true, поскольку мы получаем нормализованные координаты ограничивающего прямоугольника от интерпретатора.
from object_detection.utils import visualization_utils as vis_util
vis_util.visualize_boxes_and_labels_on_image_array(
img,
output_dict['detection_boxes'],
output_dict['detection_classes'],
output_dict['detection_scores'],
category_index,
use_normalized_coordinates=True,
min_score_thresh=0.6,
line_thickness=3)
Полученные результаты
Результат будет примерно таким, как показано ниже.
Однако есть еще одна проблема, которую необходимо решить, как показано ниже.
Это можно сделать с помощью не максимального подавления.
Не максимальное подавление
Я не собираюсь объяснять это, поскольку это уже было подробно освещено в различных статьях в Интернете. Один из таких примеров - эта статья. Чтобы реализовать это, я собираюсь использовать combined_non_max_suppression Tensorflow Image для выполнения этой задачи, поскольку он позволяет нам работать с несколькими классами одновременно. Он принимает выходные данные и возвращает прогнозы, оставшиеся после порога.
def apply_nms(output_dict, iou_thresh=0.5, score_thresh=0.6):
q = 90 # no of classes
num = int(output_dict['num_detections'])
boxes = np.zeros([1, num, q, 4])
scores = np.zeros([1, num, q])
# val = [0]*q
for i in range(num):
# indices = np.where(classes == output_dict['detection_classes'][i])[0][0]
boxes[0, i, output_dict['detection_classes'][i], :] = output_dict['detection_boxes'][i]
scores[0, i, output_dict['detection_classes'][i]] = output_dict['detection_scores'][i]
nmsd = tf.image.combined_non_max_suppression(
boxes=boxes,
scores=scores,
max_output_size_per_class=num,
max_total_size=num,
iou_threshold=iou_thresh,
score_threshold=score_thresh,
pad_per_class=False,
clip_boxes=False)
valid = nmsd.valid_detections[0].numpy()
output_dict = {
'detection_boxes' : nmsd.nmsed_boxes[0].numpy()[:valid],
'detection_classes' : nmsd.nmsed_classes[0].numpy().astype(np.int64)[:valid],
'detection_scores' : nmsd.nmsed_scores[0].numpy()[:valid],
}
return output_dict
Полный код приведен ниже, или вы посетите мое репо на Github, которое также содержит visualization_utils.py и модели.
Прежде чем закончить, я хотел бы прояснить одну вещь: если вы попытаетесь запустить его в Windows с процессором Intel, вы получите ужасный fps. У меня на i5 было ~ 2, и для сравнения та же модель Tensorflow без tflite дала мне ~ 8 кадров в секунду. Это объясняется здесь. Однако на периферийных устройствах это не будет проблемой, а значительно меньший объем памяти принесет пользу их ограничениям в памяти.
Хотя эта модель не очень точна, я надеюсь, что я предоставил бы шаблон, чтобы упростить вашу задачу при использовании детектора объектов, если Tflite.
