Дружим YOLACT и RockChip: запуск инстанс-сегментации на китайском одноплатнике

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

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

Статья «YOLACT. Real-time Instance Segmentation» вышла в 2019 году. И на тот момент сетка была в четыре раза производительней других алгоритмов, решающих подобную задачу (33 fps против 8). При этом точность (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 файл с разметкой в формате GitHub. В нём находятся две rknn‑сетки: «yolact_550.rknn» и «yolact_544.rknn», и описаны два постпроцесса: ONNX и RKNN, для первой и второй модели, соответственно. Спасибо за внимание!