Приветствую всех читателей Хабра! Сегодня я хочу поделиться с вами своим опытом запуска YOLACT на edge-устройстве RockChip. Несмотря на то, что процесс запуска занял больше времени, чем я ожидал, я решил поделиться с вами своими наработками, чтобы помочь другим разработчикам, которые могут столкнуться с той же задачей. В конце концов я нашёл способ запуска yolact, который позволил достичь высокой производительности и качества модели. Надеюсь, что мой опыт будет полезен для вас и поможет вам избежать ошибок, которые я совершил. Приятного чтения!

Так уж вышло, что задачам классификации изображений, детекции и сегментации объектов на изображениях посвящено множество статей, обзоров и мануалов. Вообще, тема алгоритмов компьютерного зрения является интересной. При этом, когда речь заходит об instance-segmentation, авторы ограничиваются обзором Mask-RCNN и Fast R-CNN, сетей которые вышли более семи лет назад. Запуск в реал‑тайме этих сетей в принципе невозможен даже на мощном железе, что уж говорить об «одноплатниках». Пытаясь восполнить этот пробел и имея под рукой мощный и дешевый «firefly» я натолкнулся на yolact.

Немного о самой yolact

Статья «YOLACT. Real-time Instance Segmentation» вышла в 2019 году. И на тот момент сетка была в четыре раза производительней других алгоритмов, решающих подобную задачу (33 fps против 8). При этом точность (mAp) у yolact была примерно на 17 процентов ниже чем у знаменитой Mask-RCNN. Авторы заявляют, что yolact это первая instance-segmentation сеть работающая в "реал-тайме", что весьма и весьма впечатляет.

Нужно сказать, что instance-segmentation является наиболее интересной задачей сегментации, так как ближе всего повторяет процесс, который происходит у нас в мозгу, когда мы смотрим на окружающие объекты. Мы видим не только границы объектов, но и понимаем к какому классу эти объекты принадлежат. Yolact делает это за счёт "пирамиды свёрточных слоёв", выделяющих карты признаков разного масштаба, в основании. И Protonet, использующейся для генерации k-масок (k = 32). Как пишут авторы, архитектура сети была основана на RetinaNet + FPN.

Немного о RockChip

Firefly RK3588 имеет на борту нейропроцессор (NPU) с тремя ядрами, мощностью до 6 TOPS. И для того чтобы использовать весь его потенциал (запустить инференс на NPU) необходимо перевести модель в определенный формат (граф) - rknn, а веса модели сжать (квантизировать) до точности 16 или 8 бит. Сделать это можно с помощью небольшого скрипта, пример которого находится в официальном репозитории firefly. Конвертация в rknn-формат может идти напрямую из model.pt, а также из model.onnx.

Обзор yolact-проектов

dbolya. В интернете существует несколько репозиториев с реализацией yolact. В первую очередь - это репозиторий от авторов оригинальной статьи: dbolya. Проект не отличается дружелюбностью к пользователю (несмотря на объемный "Readme"), но запускается из коробки. Большим минусом данного проекта является невозможность сконвертировать модель в onnx формат. Автор признается, что не собирается переписывать модель, чтобы добавить возможность конвертации в onnx, так как это снижает её производительность. Конвертация напрямую в rknn также невозможна. Поэтому, если вы собираетесь запускать сегментацию на компьютере, сервере и т. п. можно использовать этот проект, а дальше можно не читать.

Ma-Dan. К счастью, существует "форк" yolact от Ma-Dan, в котором из класса модели выкинуты "ненужные" декораторы и всякие прочие jit-компиляторы, что позволяет сконвертировать её в onnx формат. Но у Ma-Dan модели есть странный выход (Рис. 1), из-за чего конвертация в rknn-формат, необходимый для запуска на firefly, не проходит. Наверное, этот "output" можно удалить и всё заработает, но я этого не проверял.

PINTO. Ещё одну версию yolact можно найти среди сотен других моделей у автора PINTO0309 (Честно, я не знаю чем занимается этот человек в свободное время, но подозреваю, что он живет на Венере, так как в сутках там ~5832 часа, иначе как объяснить его количество «contributions»). Всё что нужно - скачать репозиторий, скачать подготовленные автором модели и пост-процессы (в формате onnx) с помощью download.sh — скрипта. Поздравляю, вы на 90% приблизились к запуску yolact на edge-девайсе.

Тут есть несколько моментов. Во‑первых, нужно сконвертировать модель (например «yolact_base_54_800 000_550×550») в rknn формат. Цифры в конце названия модели означают размер входного изображения. Во‑вторых, нужно сконвертировать «postprocess550×550» в rknn формат. А это сделать невозможно не так просто — привет слои «reshape 1×30 963». Проще запустить постпроцесс с помощью onnxruntime. В‑третьих, то что у автора называется постпроцесс (postprocess550×550.onnx), на самом деле это лишь часть постпроцесса. Так, если сравнить его с оригинальным репозиторием (dbolya/yolact), то окажется что это часть класса Detect(), который выдаёт четыре аутпута: классы, боксы, скоры и маски (classes, boxes, scores, masks). Поэтому, для получения результата как на картинке (Рис.2), нужно добавить следующие строчки:

def prep_display(results):
    def crop(bbox, shape):
        x1 = max(int(bbox[0] * shape[1]), 0)
        y1 = max(int(bbox[1] * shape[0]), 0)
        x2 = max(int(bbox[2] * shape[1]), 0)
        y2 = max(int(bbox[3] * shape[0]), 0)
        return (slice(y1, y2), slice(x1, x2))
    bboxes, scores, class_ids, masks = [], [], [], []
    for result, mask in zip(results[0][0], results[1]):
        bbox = result[:4].tolist()
        score = result[4]
        class_id = int(result[5])
        if threshold <= score:
            mask = np.where(mask > 0.5, class_id + 1, 0).astype(np.uint8)
            region = crop(bbox, mask.shape)
            cropped = np.zeros(mask.shape, dtype=np.uint8)
            cropped[region] = mask[region]
            bboxes.append(bbox)
            class_ids.append(class_id)
            scores.append(score)
            masks.append(cropped)
    return bboxes, scores, class_ids, masks

и, наконец, отрисовываем результат:

def onnx_draw(frame, bboxes, scores, class_ids, masks):
      colors = get_colors(len(COCO_CLASSES))
      frame_height, frame_width = frame.shape[0], frame.shape[1]
      # Draw masks
      if len(masks) > 0:
          mask_image = np.zeros(MASK_SHAPE, dtype=np.uint8)
          for mask in masks:
              color_mask = np.array(colors, dtype=np.uint8)[mask]
              filled = np.nonzero(mask)
              mask_image[filled] = color_mask[filled]
          mask_image = cv2.resize(mask_image, (frame_width, frame_height), cv2.INTER_NEAREST)
          cv2.addWeighted(frame, 0.5, mask_image, 0.5, 0.0, frame)
      # Draw boxes
      for bbox, score, class_id in zip(bboxes, scores, class_ids):
          x1, y1 = int(bbox[0] * frame_width), int(bbox[1] * frame_height)
          x2, y2 = int(bbox[2] * frame_width), int(bbox[3] * frame_height)
          color = colors[class_id + 1]
          frame = draw_box(frame, (x1, y1, x2, y2), color, class_id, score)
      return frame

Но есть и четвертый, самый неприятный момент — yolact_base_54_800000_550x550.onnx от PINTO нельзя обучить. А если у вас получится это сделать, то в случае если количество классов будет отличаться от исходных 80, то postprocess550x550.onnx перестанет работать. Так как размер входного тензора у postprocess550x550 фиксированный. PINTO пишет, что размер входных тензоров у onnx-моделей можно менять с помощью определённых программ, но заморачиваться с этим каждый раз когда вам нужно обучить модель на новые классы, вместо того, чтобы иметь нормально работающий постпроцесс, как-то не хочется.
Подводя итоги, если вы не собираетесь переучивать yolact на свои классы, а хотите быстро запустить модель для демонстрации, то можно обойтись репозиторием PINTO.

feiyuhuahuo. Последним и окончательным вариантом оказался Yolact_minimal. Сам проект как и модель является упрощенным вариантом оригинального проекта dbolya/yolact. Обучение модели, оценка точности и конвертация её в onnx-формат делаются в несколько строчек и изменением конфигов. Модель прекрасно конвертируется в fp16, а если использовать resnet101 в качестве backbone, то и квантуется до int8, правда точность после квантизации оставляет желать лучшего. Один из вариантов увеличения точности квантованной модели — это использование Quantization Aware Training.
Единственный минус — для запуска инференса на устройстве, необходимо перенести туда весь пост процесс, причем использовать можно только те функции, которые работают с numpy-массивами. После этого можно получить кое-какие результаты.

Скачиваем случайную картинку с интернета, обрезаем до нужного размера и скармливаем её нейронке. Смотрим на результат:

Видим, что результат пока не удовлетворительный. Куча ббоксов, маски перекрывают друг‑друга, странный «score». Как бы там ни было, уже что‑то появилось и с этим можно работать. «Score» больше единицы наводит на некоторые мысли. Сравнивая rknn модель и модель Yolact_minimal, например с помощью netron.app (как оказалось он прекрасно «кушает» rknn‑модели), можно увидеть, что выход у исходной сетки прежде чем быть отфильтрован по «nms_score_thre» проходит через функцию softmax (class_pred = F.softmax(class_pred, -1)). Понятно, значит при конвертации потерялся softmax, давайте добавлять:

def np_softmax(x):
    np_max = np.max(x, axis=1)
    sft_max = []
    for idx, pred in enumerate(x):
        e_x = np.exp(pred - np_max[idx])
        sft_max.append(e_x / e_x.sum())
    sft_max = np.array(sft_max)
    return sft_max
  
  # Обратите внимание, что softmax рассчитывается в цикле, 
  # для каждого вектора класса отдельно (что логично). 

После чего получаем симпатичный результат. И это на устройстве размером с пластиковую карту!

Всё работает как надо, и теперь можно выставлять уровень уверенности, порог который плохо распознанные объекты не должны пересекать.

Осталось обучить модель под ваши задачи. Как это сделать написано в репозитории Yolact_minimal. Если кратко, то главное получить JSON файл с разметкой в формате COCO (custom_ann.json). И вставить название нужных классов в конфиг.

Тренировка на 50-ти картинках занимает около 1,5 часов на одной карточке V100. Сеть работает только с изображениями размеры которых кратны 32. Мной были выбраны изображения 544 на 544px.

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

Если у вас есть какие-либо вопросы или предложения, не стесняйтесь писать их в комментариях. Я всегда готовы поделиться своим опытом и помочь в решении сложных задач.

Также хочу порекомендовать свой репозиторий для запуска yolact: GitHub. В нём находятся две rknn‑сетки: «yolact_550.rknn» и «yolact_544.rknn», и описаны два постпроцесса: ONNX и RKNN, для первой и второй модели, соответственно. Спасибо за внимание!