Распознавание изображений, обнаружение объектов и оценка позы с использованием Tensorflow Lite на Raspberry Pi

Вступление

Что такое Edge (или Fog) C вычисление?

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

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

Периферийные вычисления были разработаны в связи с экспоненциальным ростом количества устройств IoT, подключенных к Интернету для получения информации из облака или доставки данных обратно в облако. И многие устройства Интернета вещей (IoT) генерируют огромные объемы данных во время своей работы.

Граничные вычисления предоставляют новые возможности в приложениях IoT, особенно для тех, кто полагается на машинное обучение (ML) для таких задач, как обнаружение объектов и позы, распознавание изображений (и лиц), языковая обработка и предотвращение препятствий. Данные изображений являются отличным дополнением к IoT, но также являются значительным потребителем ресурсов (таких как мощность, память и обработка). Обработка изображений «на грани» с использованием классических моделей AI / ML - это большой скачок!

Tensorflow Lite - машинное обучение (ML) на передовой !!

Машинное обучение можно разделить на два отдельных процесса: обучение и вывод, как описано в Блоге Gartner:

  • Обучение. Обучение - это процесс создания алгоритма машинного обучения. Обучение включает использование структуры глубокого обучения (например, TensorFlow) и набора обучающих данных (см. Левую часть рисунка выше). Данные Интернета вещей предоставляют источник обучающих данных, которые специалисты по обработке данных и инженеры могут использовать для обучения моделей машинного обучения для различных случаев, от обнаружения сбоев до потребительского интеллекта.
  • Вывод. Под выводом понимается процесс использования обученного алгоритма машинного обучения для прогнозирования. Данные IoT могут использоваться в качестве входных данных для обученной модели машинного обучения, позволяя делать прогнозы, которые могут направлять логику принятия решений на устройстве, на пограничном шлюзе или в другом месте системы IoT (см. Правую часть рисунка выше).

TensorFlow Lite - это платформа глубокого обучения с открытым исходным кодом, которая позволяет вывод машинного обучения на устройстве с низкой задержкой и небольшим двоичным размером. Он разработан, чтобы упростить машинное обучение на устройствах на границе сети вместо того, чтобы отправлять данные с сервера туда и обратно.

Машинное обучение на устройстве может помочь улучшить:

  • Задержка: нет возврата к серверу и обратно.
  • Конфиденциальность: данные не должны покидать устройство
  • Подключение: подключение к Интернету не требуется.
  • Энергопотребление: сетевые подключения требуют большого количества энергии.

TensorFlow Lite (TFLite) состоит из двух основных компонентов:

  • Конвертер TFLite, который преобразует модели TensorFlow в эффективную форму для использования интерпретатором и может вносить оптимизацию для улучшения размера двоичного файла и производительности.
  • Интерпретатор TFLite работает со специально оптимизированными моделями на многих различных типах оборудования, включая мобильные телефоны, встроенные устройства Linux и микроконтроллеры.

Таким образом, обученная и сохраненная модель TensorFlow (например, model.h5) может быть преобразована с помощью TFLite Converter в TFLite FlatBuffer (например, model.tflite ), который будет использоваться TF Lite Interpreter внутри устройства Edge (как Raspberry Pi) для выполнения логического вывода новых данных.

Например, я обучил с нуля простую модель классификации изображений CNN на моем Mac (Сервер на рисунке выше). Окончательная модель имела 225 610 параметров для обучения, используя в качестве входных данных набор данных CIFAR10: 60 000 изображений (форма: 32, 32, 3). Обученная модель (cifar10_model.h5) имела размер 2,7 МБ. Используя TFLite Converter, модель, используемая на Raspberry Pi (model_cifar10.tflite), стала с 905 КБ (примерно 1/3 от исходного размера). Вывод с обеими моделями (.h5 на Mac и .tflite на RPi) дает одинаковые результаты. Оба блокнота можно найти на GitHub.

Raspberry Pi - Установка TFLite

Также можно обучать модели с нуля на Raspberry Pi, и для этого необходим полный пакет TensorFlow. Но как только мы займемся только частью вывода, мы установим только интерпретатор TensorFlow Lite.

Пакет только для интерпретатора является частью размера полного пакета TensorFlow и включает в себя минимум кода, необходимый для выполнения выводов с помощью TensorFlow Lite. Он включает только класс Python tf.lite.Interpreter, используемый для выполнения .tflite моделей.

Давайте откроем терминал на Raspberry Pi и установим колесо Python, необходимое для вашей конкретной конфигурации системы. Параметры можно найти по этой ссылке: Python Quickstart. Например, в моем случае я использую Linux ARM32 (Raspbian Buster - Python 3.7), поэтому командная строка выглядит так:

$ sudo pip3 install https://dl.google.com/coral/python/tflite_runtime-2.1.0.post1-cp37-cp37m-linux_armv7l.whl

Если вы хотите дважды проверить, какая версия ОС у вас установлена ​​на Raspberry Pi, выполните команду:

$ uname -

Как показано на изображении ниже, если вы получите… arm7l…, операционная система будет 32-битной Linux.

Установка колеса Python - единственное требование для работы интерпретатора TFLite в Raspberry Pi. Можно дважды проверить, в порядке ли установка, вызвав интерпретатор TFLite на терминале, как показано ниже. Если ошибок не появилось, у нас все в порядке.

Классификация изображений

Вступление

Одна из наиболее классических задач IA, применяемых к компьютерному зрению (CV), - это классификация изображений. Начиная с 2012 года, IA и Deep Learning (DL) навсегда изменились, когда сверточная нейронная сеть (CNN) под названием AlexNet (в честь своего ведущего разработчика Алекса Крижевского) достигла ошибки в топ-5 в 15,3%. Конкурс ImageNet 2012. Согласно The Economist: «Внезапно люди начали обращать внимание (в DL) не только в сообществе ИИ, но и во всей технологической отрасли в целом.

Этот проект, почти через восемь лет после Алекса Крижевска, более современная архитектура (MobileNet) также был предварительно обучен на миллионах изображений с использованием того же набора данных ImageNet, в результате чего было получено 1000 различных классов. Эта предварительно обученная и квантованная модель была преобразована в .tflite и использована здесь.

Сначала давайте на Raspberry Pi перейдем в рабочий каталог (например, Image_Recognition). Затем необходимо создать два подкаталога, один для моделей, а другой для изображений:

$ mkdir images
$ mkdir models

Оказавшись в каталоге модели, давайте загрузим предварительно обученную модель (по этой ссылке можно загрузить несколько разных моделей). Мы будем использовать квантованную модель Mobilenet V1, предварительно обученную с изображениями 224x224 пикселей. ZIP-файл, который можно загрузить из Классификации изображений TensorFlow Lite с помощью wget:

$ cd models
$ wget https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_1.0_224_quant_and_labels.zip

Далее распакуйте файл:

$ unzip mobilenet_v1_1.0_224_quant_and_labels

Скачиваются два файла:

  • mobilenet_v1_1.0_224_quant.tflite: преобразованная модель TensorFlow-Lite
  • labels_mobilenet_quant_v1_224.txt: набор данных ImageNet 1000 ярлыков классов.

Теперь получите несколько изображений (например, .png, .jpg) и сохраните их в подкаталоге created images.

На GitHub можно найти изображения, использованные в этом руководстве.

Установка Raspberry Pi OpenCV и Jupyter Notebook

OpenCV (Open Source Computer Vision Library) - это библиотека программного обеспечения для компьютерного зрения и машинного обучения с открытым исходным кодом. Это полезно в качестве поддержки при работе с изображениями. Если очень просто установить его на Mac или ПК, это будет немного «уловкой», чтобы сделать это на Raspberry Pi, но я рекомендую его использовать.

Пожалуйста, следуйте этому замечательному руководству от Q-Engineering, чтобы установить OpenCV на Raspberry Pi: Установите OpenCV 4.4.0 на Raspberry Pi 4. Хотя оно написано для Raspberry Pi 4, руководство также можно использовать без каких-либо изменений для Raspberry 3 или 2. .

Затем установите Jupyter Notebook. Это будет наша платформа для разработки.

$ sudo pip3 install jupyter
$ jupyter notebook

Кроме того, во время установки OpenCV должен был быть установлен NumPy, если не сделать это сейчас, то же самое с MatPlotLib.

$ sudo pip3 install numpy
$ sudo apt-get install python3-matplotlib

И готово! У нас есть все необходимое для того, чтобы начать наш путь ИИ к краю!

Вывод классификации изображений

Создайте новую записную книжку Jupyter и следуйте инструкциям ниже или загрузите полную записную книжку с GitHub.

Библиотеки импорта:

import numpy as np
import matplotlib.pyplot as plt
import cv2
import tflite_runtime.interpreter as tflite

Загрузите модель TFLite и разместите тензоры:

interpreter = tflite.Interpreter(model_path=’./models/mobilenet_v1_1.0_224_quant.tflite’)
interpreter.allocate_tensors()

Получите входные и выходные тензоры:

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

Сведения о вводе предоставят вам необходимую информацию о том, как модель должна быть заполнена изображением:

Форма (1, 224x224x3) сообщает, что изображение с размерами: (224x224x3) должно вводиться одно за другим (размер пакета: 1). Dtype uint8 сообщает, что значения являются 8-битными целыми числами.

Детали вывода показывают, что в результате логического вывода будет получен массив из 1001 целочисленных значений (8 бит). Эти значения являются результатом классификации изображений, где каждое значение представляет собой вероятность того, что конкретная метка связана с изображением.

Например, предположим, что мы хотим классифицировать изображение, имеющее форму (1220, 1200, 3). Во-первых, нам нужно изменить его форму на (224, 224, 3) и добавить размер партии 1, как определено во входных данных: (1, 224, 224, 3). Результатом вывода будет массив размером 1001, как показано ниже:

Шаги для кодирования этих операций:

  1. Введите изображение и преобразуйте его в RGB (OpenCV читает изображение как BGR):
image_path = './images/cat_2.jpg'
image = cv2.imread(image_path)
img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

2. Предварительно обработайте изображение, изменив форму и добавив размер партии:

img = cv2.resize(img, (224, 224))
input_data = np.expand_dims(img, axis=0)

3. Укажите данные, которые будут использоваться для тестирования, и запустите интерпретатор:

interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()

4. Получите результаты и сопоставьте их с классами:

predictions = interpreter.get_tensor(output_details[0][‘index’])[0]

Выходные значения (прогнозы) варьируются от 0 до 255 (максимальное значение 8-битного целого числа). Чтобы получить прогноз в диапазоне от 0 до 1, выходное значение следует разделить на 255. Индекс массива, связанный с наивысшим значением, является наиболее вероятной классификацией такого изображения.

Имея индекс, мы должны определить, к какому классу он отнесен (например, автомобиль, кошка или собака). Текстовый файл, загруженный с моделью, имеет метку, связанную с каждым индексом, от 0 до 1000.

Давайте сначала создадим функцию для загрузки файла .txt в качестве словаря:

def load_labels(path):
    with open(path, 'r') as f:
        return {i: line.strip() for i, line in enumerate(f.readlines())}

И создайте словарь с именем label и просмотрите некоторые из них:

labels = load_labels('./models/labels_mobilenet_quant_v1_224.txt')

Возвращаясь к нашему примеру, давайте получим 3 лучших результата (самые высокие вероятности):

top_k_indices = 3
top_k_indices = np.argsort(predictions)[::-1][:top_k_results]

Мы видим, что 3 ведущих индекса относятся к кошкам. Содержание прогноза - это вероятность, связанная с каждой из меток. Как объяснялось ранее, разделив на 255., мы можем получить значение от 0 до 1. Давайте создадим цикл, чтобы просмотреть верхние результаты, напечатать этикетку и вероятности:

for i in range(top_k_results):
    print("\t{:20}: {}%".format(
        labels[top_k_indices[i]],
        int((predictions[top_k_indices[i]] / 255.0) * 100)))

Давайте создадим функцию для плавного вывода различных изображений:

def image_classification(image_path, labels, top_k_results=3):
    image = cv2.imread(image_path)
    img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    plt.imshow(img)
img = cv2.resize(img, (w, h))
    input_data = np.expand_dims(img, axis=0)
interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    predictions = interpreter.get_tensor(output_details[0]['index'])[0]
top_k_indices = np.argsort(predictions)[::-1][:top_k_results]
print("\n\t[PREDICTION]        [Prob]\n")
    for i in range(top_k_results):
        print("\t{:20}: {}%".format(
            labels[top_k_indices[i]],
            int((predictions[top_k_indices[i]] / 255.0) * 100)))

На рисунке ниже показаны некоторые тесты с использованием функции:

Общая производительность потрясающая! С момента ввода пути к изображению на карте памяти до момента распечатки этого результата весь процесс с высокой точностью занял менее полсекунды!

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

Обнаружение объекта

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

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

Для этой задачи мы загрузим модель Mobilenet V1, предварительно обученную с использованием набора данных COCO (Общие объекты в контексте). Этот набор данных содержит более 200 000 изображений с пометками в 91 категории.

Скачивание модели и этикеток

На терминале Raspberry выполните команды:

$ cd ./models 
$ curl -O http://storage.googleapis.com/download.tensorflow.org/models/tflite/coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.zip
$ unzip coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.zip
$ curl -O  https://dl.google.com/coral/canned_models/coco_labels.txt
$ rm coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.zip
$ rm labelmap.txt

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

coco_labels.txt  
detect.tflite

Шаги по выполнению логического вывода для нового изображения очень похожи на шаги, выполняемые с помощью классификации изображений, за исключением того, что:

  • ввод: изображение должно иметь форму 300x300 пикселей
  • вывод: включает не только метку и вероятность («оценка»), но также относительное положение окна («ограничивающая рамка») относительно того, где находится объект на изображении.

Теперь мы должны загрузить метки и модель, распределяя тензоры.

labels = load_labels('./models/coco_labels.txt')
interpreter = Interpreter('./models/detect.tflite')
interpreter.allocate_tensors()

Предварительная обработка ввода такая же, как и раньше, но над выводом нужно работать, чтобы получить более читаемый вывод. Приведенные ниже функции помогут в этом:

def set_input_tensor(interpreter, image):
    """Sets the input tensor."""
    tensor_index = interpreter.get_input_details()[0]['index']
    input_tensor = interpreter.tensor(tensor_index)()[0]
    input_tensor[:, :] = image
def get_output_tensor(interpreter, index):
    """Returns the output tensor at the given index."""
    output_details = interpreter.get_output_details()[index]
    tensor = np.squeeze(interpreter.get_tensor(output_details['index']))
    return tensor

С помощью вышеуказанных функций detect_objects () вернет результаты вывода:

  • идентификатор метки объекта
  • счет
  • ограничивающая рамка, которая покажет, где находится объект.

Мы включили «порог», чтобы избегать объектов с низкой вероятностью правильности. Обычно мы должны рассматривать балл выше 50%.

def detect_objects(interpreter, image, threshold):
    set_input_tensor(interpreter, image)
    interpreter.invoke()
    
    # Get all output details
    boxes = get_output_tensor(interpreter, 0)
    classes = get_output_tensor(interpreter, 1)
    scores = get_output_tensor(interpreter, 2)
    count = int(get_output_tensor(interpreter, 3))
    results = []
    for i in range(count):
        if scores[i] >= threshold:
            result = {
                'bounding_box': boxes[i],
                'class_id': classes[i],
                'score': scores[i]
            }
            results.append(result)
    return results

Если мы применим указанную выше функцию к измененному изображению (так же, как в примере классификации), мы должны получить:

Большой! Менее чем за 200 мс с вероятностью 77% объект с идентификатором 16 был обнаружен в области, ограниченной «ограничивающей рамкой»: (0,028011084, 0,020121813, 0,9886069, 0,802299). Эти четыре числа соответственно связаны с ymin, xmin, ymax и xmax.

Учтите, что y идет сверху (ymin) вниз (ymax), а x идет слева (xmin) вправо (xmax), как показано на рисунке ниже:

Имея четыре значения ограничивающего прямоугольника, мы фактически имеем координаты верхнего / левого угла и нижнего / правого угла. Имея оба края и зная форму изображения, можно нарисовать прямоугольник вокруг объекта.

Затем мы должны выяснить, что означает class_id, равный 16. Открывая файл coco_labels.txt, как словарь, каждый из его элементов имеет связанный индекс, и, проверяя индекс 16, мы получаем, как и ожидалось, 'cat'. Вероятность - это значение, возвращаемое из оценки .

Давайте создадим общую функцию для обнаружения нескольких объектов на одном изображении. Первая функция, начиная с пути к изображению, выполнит вывод, возвращая изображение с измененным размером и результаты (несколько идентификаторов, каждый со своими оценками и ограничивающими рамками:

def detectObjImg_2(image_path, threshold = 0.51):
img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    image = cv2.resize(img, (width, height),
                       fx=0.5,
                       fy=0.5,
                       interpolation=cv2.INTER_AREA)
    results = detect_objects(interpreter, image, threshold)
return img, results

Имея измененное изображение и результаты вывода, приведенную ниже функцию можно использовать для рисования прямоугольника вокруг объектов, указав для каждого из них его метку и вероятность:

def detect_mult_object_picture(img, results):
    HEIGHT, WIDTH, _ = img.shape
    aspect = WIDTH / HEIGHT
    WIDTH = 640
    HEIGHT = int(640 / aspect)
    dim = (WIDTH, HEIGHT)
   img = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)
   for i in range(len(results)):
        id = int(results[i]['class_id'])
        prob = int(round(results[i]['score'], 2) * 100)
        
        ymin, xmin, ymax, xmax = results[i]['bounding_box']
        xmin = int(xmin * WIDTH)
        xmax = int(xmax * WIDTH)
        ymin = int(ymin * HEIGHT)
        ymax = int(ymax * HEIGHT)
        text = "{}: {}%".format(labels[id], prob)
        if ymin > 10: ytxt = ymin - 10
        else: ytxt = ymin + 15
        img = cv2.rectangle(img, (xmin, ymin), (xmax, ymax),
                            COLORS[id],
                            thickness=2)
        img = cv2.putText(img, text, (xmin + 3, ytxt), FONT, 0.5, COLORS[id],
                          2)
   return img

Ниже приведены некоторые результаты:

Полный код можно найти на GitHub.

Обнаружение объекта с помощью камеры

Если у вас есть PiCam, подключенный к Raspberry Pi, вы можете снимать видео и выполнять распознавание объектов покадрово, используя те же функции, которые были определены ранее. Пожалуйста, следуйте этому руководству, если у вас нет работающей камеры в Pi: Начало работы с модулем камеры.

Во-первых, важно определить размер кадра, который будет захвачен камерой. Мы будем использовать 640x480.

WIDTH = 640
HEIGHT = 480

Далее вы должны активировать камеру:

cap = cv2.VideoCapture(0)
cap.set(3, WIDTH)
cap.set(4, HEIGHT)

И запустите приведенный ниже код в цикле. Пока не будет нажата клавиша «q», камера будет снимать видео кадр за кадром, рисуя ограничивающую рамку с соответствующими метками и вероятностями.

while True:
    timer = cv2.getTickCount()
    success, img = cap.read()
    img = cv2.flip(img, 0)
    img = cv2.flip(img, 1)
    fps = cv2.getTickFrequency() / (cv2.getTickCount() - timer)
    cv2.putText(img, "FPS: " + str(int(fps)), (10, 470),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
    image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    image = cv2.resize(image, (width, height),
                       fx=0.5,
                       fy=0.5,
                       interpolation=cv2.INTER_AREA)
    start_time = time.time()
    results = detect_objects(interpreter, image, 0.55)
    elapsed_ms = (time.time() - start_time) * 1000
    img = detect_mult_object_picture(img, results)
    cv2.imshow("Image Recognition ==> Press [q] to Exit", img)
if cv2.waitKey(1) & 0xFF == ord('q'):
        break
cap.release()
cv2.destroyAllWindows()

Ниже можно увидеть видео, работающее в реальном времени на экране Raspberry Pi. Обратите внимание, что видео работает со скоростью около 60 FPS (кадров в секунду), что довольно хорошо!

Вот один снимок экрана из приведенного выше видео:

Полный код доступен на GitHub.

Оценка позы

Одна из наиболее интересных и важных областей ИИ - это оценка позы человека в реальном времени, позволяющая машинам понимать, что люди делают на изображениях и видео. Оценка позы была подробно исследована в моей статье Оценка позы для нескольких человек в реальном времени с помощью TensorFlow2.x, но здесь, на Edge, с Raspberry Pi и с помощью TensorFlow Lite, можно легко воспроизвести почти то же самое, что было сделано на Mac.

Модель, которую мы будем использовать в этом проекте, - PoseNet. Мы сделаем вывод так же, как и для классификации изображений и обнаружения объектов, когда изображение подается через предварительно обученную модель. PoseNet поставляется с несколькими различными версиями модели, соответствующими вариантам архитектуры MobileNet v1 и архитектуры ResNet50. В этом проекте предварительно обученной версией является MobileNet V1, которая меньше, быстрее, но менее точна, чем ResNet. Кроме того, существуют отдельные модели для определения позы одного и нескольких человек. Мы рассмотрим модель, обученную для одного человека.

На этом сайте можно исследовать в режиме реального времени и с помощью прямой камеры несколько моделей и конфигураций PoseNet.

Библиотеки для выполнения оценки позы на Raspberry Pi такие же, как и раньше. Интерпретатор NumPy, MatPlotLib, OpenCV и TensorFlow Lite.

Предварительно обученная модель - это posenet_mobilenet_v1_100_257x257_multi_kpt_stripped.tflite, которую можно загрузить по указанной выше ссылке или с сайта TensorFlow Lite - Обзор оценки позы. Модель должна быть сохранена в подкаталоге models.

Начните загрузку модели TFLite и выделение тензоров:

interpreter = tflite.Interpreter(model_path='./models/posenet_mobilenet_v1_100_257x257_multi_kpt_stripped.tflite')
interpreter.allocate_tensors()

Получите входные и выходные тензоры:

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

Как и раньше, глядя на input_details, можно увидеть, что изображение, которое будет использоваться для оценки позы, должно быть (1, 257, 257, 3), что означает, что изображения должны быть изменены на 257x257 пикселей.

Давайте возьмем в качестве исходных данных простую человеческую фигуру, которая поможет нам ее проанализировать:

Первый шаг - это предварительная обработка изображения. Эта конкретная модель не была квантована, что означает, что dtype - это float32. Эта информация необходима для предварительной обработки входного изображения, как показано в приведенном ниже коде.

image = cv2.resize(image, size) 
input_data = np.expand_dims(image, axis=0)
input_data = input_data.astype(np.float32)
input_data = (np.float32(input_data) - 127.5) / 127.5

После предварительной обработки изображения пришло время выполнить логический вывод, снабдив тензор изображением и вызвав интерпретатор:

interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()

Статьей, которая очень помогает понять, как работать с PoseNet, является статья из учебника Иван Кунякин Оценка позы и сопоставление с TensorFlow lite. Там Иван комментирует, что для выходного вектора важны следующие ключевые точки:

  • Тепловые карты 3D-тензор размера (9,9,17), который соответствует вероятности появления каждой из 17 ключевых точек (суставов тела) в определенной части изображения (9,9). . Он используется для определения примерного положения сустава.
  • Векторы смещения: трехмерный тензор размера (9,9,34), который называется векторами смещения. Он используется для более точного расчета положения ключевой точки. Первые 17 третьего измерения соответствуют координатам x, а вторые 17 - координатам y.
output_details = interpreter.get_output_details()[0]
heatmaps = np.squeeze(interpreter.get_tensor(output_details['index']))
output_details = interpreter.get_output_details()[1]
offsets = np.squeeze(interpreter.get_tensor(output_details['index']))

Давайте создадим функцию, которая будет возвращать массив со всеми 17 ключевыми точками (или суставами человека) на основе тепловых карт и смещений.

def get_keypoints(heatmaps, offsets):
    joint_num = heatmaps.shape[-1]
    pose_kps = np.zeros((joint_num, 2), np.uint32)
    max_prob = np.zeros((joint_num, 1))
    for i in range(joint_num):
        joint_heatmap = heatmaps[:,:,i]
        max_val_pos = np.squeeze(
            np.argwhere(joint_heatmap == np.max(joint_heatmap)))
        remap_pos = np.array(max_val_pos / 8 * 257, dtype=np.int32)
        pose_kps[i, 0] = int(remap_pos[0] +
                             offsets[max_val_pos[0], max_val_pos[1], i])
        pose_kps[i, 1] = int(remap_pos[1] +
                             offsets[max_val_pos[0], max_val_pos[1],
                                         i + joint_num])
        max_prob[i] = np.amax(joint_heatmap)
    return pose_kps, max_prob

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

Результирующий массив показывает все 17 координат (y, x) относительно того, где расположены стыки на изображении размером 257 x 257 пикселей. Используя приведенный ниже код. Можно нанести каждый из стыков поверх изображения с измененным размером. Для справки, индекс массива аннотирован, поэтому каждое соединение легко идентифицировать:

y,x = zip(*keypts_array)
plt.figure(figsize=(10,10))
plt.axis([0, image.shape[1], 0, image.shape[0]])  
plt.scatter(x,y, s=300, color='orange', alpha=0.6)
img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plt.imshow(img)
ax=plt.gca() 
ax.set_ylim(ax.get_ylim()[::-1]) 
ax.xaxis.tick_top() 
plt.grid();
for i, txt in enumerate(keypts_array):
    ax.annotate(i, (keypts_array[i][1]-3, keypts_array[i][0]+1))

В результате получаем картинку:

Отлично, теперь пора создать общую функцию для рисования «костей», то есть соединения суставов. Кости будут нарисованы в виде линий, которые являются связями между ключевыми точками с 5 по 16, как показано на рисунке выше. Независимые круги будут использоваться для ключевых точек от 0 до 4, связанных с головой:

def join_point(img, kps, color='white', bone_size=1):
    
    if   color == 'blue'  : color=(255, 0, 0)
    elif color == 'green': color=(0, 255, 0)
    elif color == 'red':  color=(0, 0, 255)
    elif color == 'white': color=(255, 255, 255)
    else:                  color=(0, 0, 0)
    body_parts = [(5, 6), (5, 7), (6, 8), (7, 9), (8, 10), (11, 12), (5, 11),
                  (6, 12), (11, 13), (12, 14), (13, 15), (14, 16)]
    for part in body_parts:
        cv2.line(img, (kps[part[0]][1], kps[part[0]][0]),
                (kps[part[1]][1], kps[part[1]][0]),
                color=color,
                lineType=cv2.LINE_AA,
                thickness=bone_size)
    
    for i in range(0,len(kps)):
        cv2.circle(img,(kps[i,1],kps[i,0]),2,(255,0,0),-1)

Вызывая функцию, мы получаем примерную позу тела на изображении:

join_point(img, keypts_array, bone_size=2)
plt.figure(figsize=(10,10))
plt.imshow(img);

И последнее, но не менее важное: давайте создадим общую функцию для оценки позы, взяв за основу путь изображения:

def plot_pose(img, keypts_array, joint_color='red', bone_color='blue', bone_size=1):
    join_point(img, keypts_array, bone_color, bone_size)
    y,x = zip(*keypts_array)
    plt.figure(figsize=(10,10))
    plt.axis([0, img.shape[1], 0, img.shape[0]])  
    plt.scatter(x,y, s=100, color=joint_color)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    plt.imshow(img)
    ax=plt.gca() 
    ax.set_ylim(ax.get_ylim()[::-1]) 
    ax.xaxis.tick_top() 
    plt.grid();
    return img

def get_plot_pose(image_path, size, joint_color='red', bone_color='blue', bone_size=1):    
    image_original = cv2.imread(image_path)
    image = cv2.resize(image_original, size) 
    input_data = np.expand_dims(image, axis=0)
    input_data = input_data.astype(np.float32)
    input_data = (np.float32(input_data) - 127.5) / 127.5
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    
    output_details = interpreter.get_output_details()[0]
    heatmaps = np.squeeze(interpreter.get_tensor(output_details['index']))
    output_details = interpreter.get_output_details()[1]
    offsets = np.squeeze(interpreter.get_tensor(output_details['index']))
    keypts_array, max_prob = get_keypoints(heatmaps,offsets)
    orig_kps = get_original_pose_keypoints(image_original, keypts_array, size)
    img = plot_pose(image_original, orig_kps, joint_color, bone_color, bone_size)
    
    return orig_kps, max_prob, img

На данный момент с помощью только одной строчки кода можно определить позу на изображениях:

keypts_array, max_prob, img  = get_plot_pose(image_path, size, bone_size=3)

Весь код, разработанный в этом разделе, доступен на GitHub.

Еще один простой шаг - применить эту функцию к кадрам из видео и камеры в реальном времени. Я оставлю это тебе! ;-)

Заключение

TensorFlow Lite - отличный фреймворк для реализации искусственного интеллекта (точнее, машинного обучения) на периферии. Здесь мы исследовали модели машинного обучения, работающие на Raspberry Pi, но сейчас TFLite все больше и больше используется на «краю края», на очень маленьких микроконтроллерах, в так называемом TinyML.

Как всегда, я надеюсь, что эта статья может вдохновить других найти свой путь в фантастическом мире искусственного интеллекта!

Все коды, использованные в этой статье, доступны для скачивания в проекте GitHub: TFLite_IA_at_the_Edge.

Привет с Юга Мира!

Увидимся в моей следующей статье!

Спасибо

Марсело