Приветствую всех читателей Хабра! Сегодня я хочу поделиться с вами своим опытом запуска 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где "results" - это четыре тензора, полученных на выходе postprocess550x550.
и, наконец, отрисовываем результат:
функция рисующая маски и боксы на изображении
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.
postprocess550x550
На самом деле вся эта история с postprocess550x550.onnx была связана с предположением, что он будет работать быстрее, чем постпроцесс написанный вручную и использующий циклы на python. Это предположение оказалось неверным.
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, для первой и второй модели, соответственно. Спасибо за внимание!
