Привет, Хабр! Меня зовут Кирилл Колодяжный, я разрабатываю системы хранения данных в YADRO и изучаю нестандартные подходы к машинному обучению: создаю ML-проекты на С++. 

Это вторая часть цикла о разработке приложения для обнаружения предметов на С++ . В прошлом материале мы выяснили, как создать проект в IDE Android Studio, реализовать сессию непрерывного захвата и преобразовать изображение в матрицу OpenCV. Ссылку вы найдете в конце статьи.

В этой статье продолжим реализацию проекта и обсудим следующие шаги:

• Как подключить к проекту библиотеки машинного обучения PyTorch и NCNN.

• Как получить модели YOLOv5 и YOLOv4 для использования на мобильном устройстве.

• Как реализовать инференс моделей для обнаружения объектов.

• Как обработать результаты работы моделей YOLO, реализовав алгоритмы Non-Maximum-Suppression и Intersection-Over-Union.

В конце сравним производительность PyTorch и NCNN и решим, какой фреймворк подойдет для задачи лучше.

Как подключить к проекту PyTorch

Использование готового дистрибутива

Готовая сборка — самый простой способ использовать PyTorch в приложении для Android. В репозитории Maven доступен бинарный дистрибутив PyTorch для мобильных устройств org.pytorch: pytorch_android_lite. Правда, у дистрибутива устаревшая версия. Если вам подходит, достаточно добавить в файл build.gradle.kts следующие строки:

dependencies {
  	...
 	implementation 'org.pytorch:pytorch_android:1.6.0-SNAPSHOT'
 	extractForNativeBuild 'org.pytorch:pytorch_android:1.6.0-SNAPSHOT'
 }

В иных случаях последнюю версию нужно собирать из исходного кода.

Самостоятельная сборка из исходного кода

Собрать мобильную версию PyTorch можно как обычную, но с дополнительными параметрами CMake. Допустим, у нас уже установлен Native Development Kit для Android (NDK), который включает в себя соответствующую версию компилятора C/C++ и нативные библиотеки Android, нужные для сборки приложения. Следующий фрагмент кода показывает, как использовать командную строку для скачивания нужной версии PyTorch и создания мобильной сборки для Android:

cd /home/USER
 git clone https://github.com/pytorch/pytorch.git
 cd pytorch/
 git checkout v2.3.1
 git submodule update --init --recursive
 
 export ANDROID_NDK=$START_DIR/android/ndk/26.1.10909125
 export ANDROID_ABI='arm64-v8a'
 export ANDROID_STL_SHARED=1
 
 scripts/build_android.sh \
 -DBUILD_CAFFE2_MOBILE=OFF \
 -DBUILD_SHARED_LIBS=ON \
 -DUSE_VULKAN=OFF \
 -DCMAKE_PREFIX_PATH=$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())') \
 -DPYTHON_EXECUTABLE=$(python -c 'import sys; print(sys.executable)') \

Здесь /home/USER — это домашний каталог пользователя. Когда дело доходит до создания мобильной версии PyTorch, основное требование — переменная среды ANDROID_NDK, которая указывает на каталог установки Android NDK. 

Самый простой способ установить инструменты разработки Android — загрузить IDE Android Studio и использовать из нее инструмент SDK Manager. Вы можете найти SDK Manager в меню Tools | SDK Manager. Менеджер можно использовать для установки соответствующих версий Android SDK. 

Также можно установить соответствующий NDK и утилиту CMake, используя SDK Tools во вкладке в окне менеджера. Переменную среды ANDROID_ABI можем использовать для указания нужной архитектуры процессора, чтобы компилятор генерировал код, зависящий от конкретной архитектуры. В этом примере я использую архитектуру arm64-v8a. 

Далее я использовал скрипт build_android.sh из исходного кода для сборки мобильной версии PyTorch. Этот скрипт использует CMake, поэтому и принимает определения параметров CMake в качестве аргументов. Обратите внимание, что я передаю параметр BUILD_CAFFE2_MOBILE=OFF, чтобы отключить сборку мобильной версии Caffe2. Ее сложно использовать в текущей версии, поскольку библиотека устарела. 

Второй важный параметр — BUILD_SHARED_LIBS=ON, он нужен для создания разделяемых библиотек. Честно говоря, у меня были проблемы со сборкой статических библиотек PyTorch предыдущих версий, поэтому я не пробовал их собирать для последней. Возможно, уже все работает исправно. 

Кроме того, я отключил поддержку Vulkan API, используя DUSE_VULKAN = OFF, потому что она все еще экспериментальная, в ней есть проблемы с компиляцией. Я настроил и другие параметры, например, пути установки Python для промежуточной генерации кода сборки. Теперь, когда я собрал мобильные библиотеки PyTorch, их нужно скопировать в папку JniLibs в проекте так, чтобы получилась следующая структура:

app
 |--src
 |  |--main
 |  |--…
 |  |--JniLibs
 | 	`--arm64-v8a
 |    	|--libc10.so
 |    	|--libtorch_cpu.so
 |    	|--libtorch_global_deps.so
 |    	`--libtorch.so
 `...

Если вы собираете проект для другой архитектуры, например для armeabi-v7a, создайте в JniLibs соответствующую папку и скопируйте туда собранные библиотеки. Также, чтобы использовать PyTorch в проекте, необходимо настроить компилятор так, чтобы он находил заголовочные файлы и библиотеки. Для этого передадим соответствующие пути через переменные CMake — изменить gradle.build.kts следующим образом:

...
 val pytorchDir = "/home/USER/development/android/pytorch/build_android/install/"
 val opencvDir = "/home/USER/development/android/OpenCV-android-sdk/sdk/native/jni/"
 val ncnnDir = "/home/USER/development/android/ncnn-20240820-android-vulkan/arm64-v8a/lib/cmake/ncnn/"
 ...
 externalNativeBuild {
 	cmake {
     	arguments += "-DANDROID_ARM_NEON=TRUE"
     	arguments += "-DANDROID_STL=c++_shared"
     	arguments += "-DTorch_DIR=${pytorchDir}"
     	arguments += "-DOpenCV_DIR=${opencvDir}"
     	arguments += "-Dncnn_DIR=${ncnnDir}"
 	}
 }
 ndk {
 	abiFilters.add("arm64-v8a")
 }

После этого надо изменить файл CMakeLists.txt для подключения соответствующих пакетов:

set(Torch_LIBS_DIR ${Torch_DIR}/lib/)
 set(Torch_LIBS ${Torch_DIR}/lib/libtorch_cpu.so ${Torch_DIR}/lib/libc10.so)
 set(Torch_INCLUDE_DIRS ${Torch_DIR}/include)
 
 find_package(OpenCV REQUIRED)
 
 find_package(ncnn REQUIRED)
 
 # ...
 
 add_library(${CMAKE_PROJECT_NAME} SHARED ${SOURCES})
 
 target_link_libraries(${CMAKE_PROJECT_NAME}
 	android
 	m
 	log
 	app_glue
 	camera2ndk
 	mediandk
 	ncnn
 	opencv_java
 	${Torch_LIBS}
 )

Для PyTorch тоже можно использовать find_package, но я не смог заставить его работать нормально в самостоятельной сборке. Напишите в комментариях, если у вас получилось сделать это.

Экспорт модели YOLOv5 для дальнейшего использования в нативной части

Мы собрали мобильную версию PyTorch, и у нас остался единственный вариант загрузки и использования готовых моделей — TorchScript. Я знаю про фреймворк для встроенных устройств ExecuTorch с отдельным подходом к экспорту и импорту моделей, но на момент написания статьи он был сырым, и у меня не вышло экспортировать YOLOv5. 

TorchScript использует трассировку модели в реальном времени, чтобы получить определение модели особого типа, которое может выполняться движком PyTorch независимо от API. В PyTorch экспортировать такие определения можно только из Python API, но мы можем использовать C++ API для загрузки модели и ее выполнения. К сожалению, мобильная версия PyTorch не позволяет нам создавать модели с помощью полнофункционального C++ API. 

Итак, чтобы использовать модель YOLOv5 в TorchScript, надо выполнить следующие шаги:

1. Скачать репозиторий модели и установить зависимости:

 > git clone https://github.com/ultralytics/yolov5
 > cd yolov5
 > pip install -r requirements.txt

2. Выполнить скрипт для экспорта PyTroch-модели, оптимизированной для мобильных устройств:

> python export.py --weights yolov5s.torchscript --include torchscript --optimize

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

В моем проекте уже есть такой файл, только я его преобразовал в TXT-формат, чтобы было проще парсить. После этого полученные файлы надо скопировать в Android-проект в папку assets, чтобы они были упакованы в установочный APK-пакет. Должна получиться следующая структура:

app
 |--src
 |  `--main
 | 	|--...
 | 	|--cpp
 | 	|--JniLibs
 | 	|--assests
 | 	|  |--yolov5.torchscript
 | 	|  `--classes.txt
 | 	`—...
 `...

Когда подключили PyTorch, можно переходить к фреймворку NCNN. 

Как подключить к проекту NCNN

NCNN — фреймворк для логического вывода (инференс) от компании Tencent. Он специально оптимизирован для работы на мобильных и встраиваемых платформах. Поддерживает различные наборы SIMD-инструкций и работу на мобильных GPU с использованием Vulkan. Я решил сравнить его с PyTorch, потому что не смог заставить мобильную версию PyTorch использовать Vulkan.

NCNN можно собрать из исходников, на сайте есть хорошая инструкция. Если не хотите собирать самостоятельно, используйте актуальные сборки из репозитория проекта. На момент написания статьи файл сборки назывался так ncnn-20240820-android-vulkan.zip — там сразу поставляются бинарники разных архитектур. Архив достаточно распаковать в удобное место, прописать пути для CMake и подключить пакет. 

Пример, как это сделать, уже есть в листингах для PyTorch. При пользовании OpenCV и NCNN в одном проекте у меня возникли проблемы из-за несовместимых ключей сборки для этих библиотек. Ключи экспортируют, используя CMake. Чтобы исправить ситуацию, достаточно изменить файл .../ncnn-XXXXXX-android-vulkan/cmake/ncnn/ncnn.cmake — в нем надо убрать строку -fno-RTTI; -fno-Exceptions.

Загрузка моделей YOLOv5 и YOLOv4 для дальнейшего использования в нативной части

Модели для NCNN подготавливаются через формат PyTorch Neural Network eXchange (PNNX). Подробную инструкцию для подготовки оставлю в дополнительных материалах в конце статьи. А готовые популярные модели можно сказать отсюда. В репозитории NCNN есть примеры, как использовать эти модели, в том числе как разбирать и анализировать выходные тензоры. Для проекта я буду использовать модель YOLOV5_V60. Для этого я скачал файлы yolov5s_6.0.param и yolov5s_6.0.bin. 

Кроме модели YOLOV5_V60, я буду использовать оптимизированную модель YOLOV4, которую нашел в проекте YOLOv5_NCNN. Для нее я использовал файлы yolo-fastest-opt.param и yolo-fastest-opt.bin. Загруженные файлы моделей также следует поместить в папку assets для того, чтобы они были упакованы в APK и были доступны во время исполнения приложения.

Рекомендую ознакомится с проектом YOLOv5_NCNN. Он во многом повторяет то, что мы уже разобрали, к тому же, там есть примеры работы с другими моделями на Android. 

Чтение упакованных файлов из assets

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

std::vector<char> read_asset(AAssetManager *asset_manager, const std::string &name) {
 	std::vector<char> buf;
 	AAsset *asset = AAssetManager_open(asset_manager, name.c_str(), AASSET_MODE_UNKNOWN);
 	if (asset != nullptr) {
     	LOGI("Open asset %s OK", name.c_str());
     	off_t buf_size = AAsset_getLength(asset);
     	buf.resize(buf_size + 1, 0);
     	auto num_read = AAsset_read(asset, buf.data(), buf_size);
     	LOGI("Read asset %s OK", name.c_str());
     	if (num_read == 0)
     	buf.clear();
     	AAsset_close(asset);
     	LOGI("Close asset %s OK", name.c_str());
 	}
 	return buf;
 }

Для чтения данных из пакета приложения я использовал четыре функции из Android NDK. Функция AAssetManager_open открывает ресурс и возвращает указатель на объект AAsset. Путь к ресурсу указывается в формате пути к файлу, где корнем этого пути является папка assets. После открытия ресурсов я использовал функцию AAsset_getLength, чтобы получить размер файла, и выделил память в контейнере std::vector<char> с помощью метода std::vector::resize. Затем я использовал функцию AAsset_read() для чтения всего файла в объект buf. У функция AAsset_read() следующие аргументы:

• Указатель на объект AAsset, из которого будет происходить чтение.

• Указатель void* на буфер памяти для чтения.

• Размер считываемых байт.

Мы видим, что API для работы с файлами в assets практически ничем не отличается от стандартного API библиотеки C для файловых операций. Когда заканчиваем работу с объектом AAsset, нужно использовать функцию AAsset_close, чтобы уведомить систему о том, что нам больше не нужен доступ к этому ресурсу. Представление считанных данных в виде объекта std::vector<char> удобно использовать для чтения сырых байтов. 

Но для других целей это не подходит, поэтому я написал два класса адаптера для более удобной работы с текстовыми данными, специализированных для моделей PyTorch.

Адаптер для чтения моделей PyTorch

Для загрузки TorchScript-моделей из памяти есть функция torch::jit::_load_for_mobile. Однако она не работает со стандартными потоками и типами C++. Вместо этого функция принимает указатель на объект класса caffe2::serialize::ReadAdapterInterface. Код ниже показывает, как сделать конкретную реализацию класса caffe2::serialize::ReadAdapterInterface, которая позволяет работать с контейнером std::vector<char>:

class ReadAdapter : public caffe2::serialize::ReadAdapterInterface {
 public:
 	explicit ReadAdapter(const std::vector<char> &buf) : buf_(&buf) {}
 	[[nodiscard]] size_t size() const override {
     	return buf_->size();
 	}
 	[[nodiscard]] size_t read(uint64_t pos,
                           	void *buf, size_t n,
                           	[[maybe_unused]]const char *what) const override {
     	std::copy_n(buf_->begin() + static_cast<ptrdiff_t>(pos),
                 	n,
                 	reinterpret_cast<char *>(buf));
     	return n;
 	}
 private:
 	const std::vector<char> *buf_;
 };

Класс ReaderAdapter переопределяет методы size и read из базового класса caffe2::serialize::ReadAdapterInterface. Их реализации довольно очевидны: метод size возвращает размер базового контейнера, а метод read копирует n байт из вектора в целевой буфер с помощью стандартного алгоритма std::copy_n.

Адаптер для чтения моделей текстовых файлов

Несмотря на то, что std::vector<char> и так представляет собой набор байт, использовать его напрямую для парсинга текста не очень удобно, поэтому я написал следующий класс:

template<typename CharT, typename TraitsT = std::char_traits<CharT> >
 struct VectorStreamBuf : public std::basic_streambuf<CharT, TraitsT> {
 	explicit VectorStreamBuf(std::vector<CharT> &vec) {
     	this->setg(vec.data(), vec.data(), vec.data() + vec.size());
 	}
 };

Класс VectorStreamBuf я унаследовал от std::basic_streambuf. В конструкторе инициализировал внутренние данные streambuf значениями char из входного вектора. В дальнейшем объект этого класса можно будет использовать в качестве обычного входного потока C++. Тут, конечно, копируются данные, и в некоторых случаях это может влиять на производительность. Однако в нашем проекте я использую его один раз: для чтения небольшого файла с данными об идентификаторах классов объектов и их строковых именах.

Чтение имен классов

Как я писал выше, файл с именами классов classes.txt, который находится в папке assets, — это преобразованный из формата YAML словарь соответствий числового идентификатора класса объекта и его строкового представления (метки). Каждая строка в файле представляется так:

[Идентификатор]:[имя класса]

Для чтения этого файла я написал следующую функцию:

using Classes = std::map<size_t, std::string>;
 ...
 Classes load_classes(std::istream &stream) {
 	LOGI("Init classes start OK");
 	Classes classes;
 	if (stream) {
     	std::string line;
     	std::string id;
     	std::string label;
     	size_t idx = 0;
     	while (std::getline(stream, line)) {
         	auto pos = line.find_first_of(':');
         	id = line.substr(0, pos);
         	label = line.substr(pos + 1);
         	classes.insert({idx, label});
         	++idx;
     	}
 	}
 	LOGI("Init classes finish OK");
 	return classes;
 }

Я читаю этот файл построчно и использую функцию std::getline для входного потока. Далее я разбиваю каждую строку на две части, используя символ : как разделитель. Первая часть каждой строки — идентификатор класса, вторая — имя класса. С помощью метода строки find_first_of находится позиция разделителя, и части выделяются с помощью метода substr. Чтобы было удобно сопоставлять идентификатор класса с его именем, я помещаю данные в словарь типа std::map.

Реализация инференса моделей

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

Общей функциональность тут будет применение алгоритмов Non-Maximum-Suppression (NMS) и Intersection-Over-Union (IOU) к полученным результатам обнаружения объектов. В прошлой части я показывал, как конкретные классы, реализующие этот интерфейс, создаются в конструкторе класса CameraCapture:

CameraCapture::CameraCapture(android_app *app) : android_app_(app) {
 	yolo_v5_torch_ = std::make_shared<TorchYolo>(app->activity->assetManager, &stats_);
 	yolo_v5_ncnn_ = std::make_shared<NcnnYolo>(app->activity->assetManager, &stats_, false);
 	yolo_v4_ncnn_ = std::make_shared<NcnnYolo>(app->activity->assetManager, &stats_, true);
 	yolo_ = yolo_v4_ncnn_;
 }

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

Сам объект детектора изображений используется в методе CameraCapture::process_image следующим образом:

auto results = yolo_->detect(rgba_img_);
 for (auto &result: results) {
 	...
 }

Как визуализировать полученные результаты, я рассказывал в первой части.

Общий интерфейс

Класс интерфейса детектора объектов очень простой:

struct DetectionResult {
 	int class_index;
 	std::string class_name;
 	float score;
 	cv::Rect rect;
 };
 
 class ObjectDetector {
 	public:
 	virtual ~ObjectDetector() = default;
 	virtual std::vector<DetectionResult> detect(const cv::Mat &image) = 0;
 	virtual const std::string& name() const = 0;
 };

У интерфейса два метода: detect и name. Метод name используется для вывода статистики и просто возвращает имя реализации, то есть используемой модели. А метод detect принимает на вход изображение в виде матрицы OpenCV и возвращает контейнер с объектами типа DetectionResult. 

Я использовал только значения class_name и rect для визуализации результатов. Это соответственно строка с именем класса обнаруженного объекта и ограничивающий его прямоугольник. Поля class_index и score — это числовой идентификатор класса объекта и значение, показывающее, насколько модель уверена в предсказанном результате (confidence score).

Реализация inference для PyTorch

Рассмотрим класс TorchYolo, в котором реализуется обнаружение объектов с помощью  PyTorch и модели YOLOv5. Начну с описания конструктора, в котором происходит загрузка модели и создание TorchScript-модуля.

TorchYolo::TorchYolo(AAssetManager *asset_manager, Stats *stats) : stats_(stats) {
 	auto model_buf = read_asset(asset_manager, model_file_name);
 	model_ = torch::jit::_load_for_mobile(std::make_unique<ReadAdapter>(model_buf));
 	const std::string classes_file_name = "classes.txt";
 	auto classes_buf = read_asset(asset_manager, classes_file_name);
 	VectorStreamBuf<char> stream_buf(classes_buf);
 	std::istream is(&stream_buf);
 	classes_ = load_classes(is);
 }

Помните, я добавил файлы yolov5s.torchscript и classes.txt в папку assets нашего проекта? В конструкторе я загружаю бинарный файл модели и файл словаря классов, используя описанную выше функцию read_asset. Затем полученные данные модели я использую для загрузки и инициализации модуля TorchScript с помощью функции torch:: jit::_load_for_mobile. Также для совместимости типов входных аргументов тут применяется описанный выше класс ReadAdapter. 

Обратите внимание, что модель TorchScript нужно сохранить с оптимизацией для мобильных устройств и загрузить с соответствующей функцией. На самом деле при компиляции PyTorch для мобильных устройств обычная функциональность torch::jit::load автоматически отключается и становится недоступной. После загрузки модели я считываю файл словаря с классами и сохраняю его в переменной класса, чтобы использовать потом в методе output2results. Тут также можно обратить внимание на использование типа VectorStreamBuf для передачи текстового потока в функцию load_classes. 

Основной метод, в котором происходит обнаружение объектов, — detect, и он реализуется так:

std::vector<DetectionResult> TorchYolo::detect(const cv::Mat &image) {
 	// yolov5 input size
 	constexpr int input_width = 640;
 	constexpr int input_height = 640;
 	cv::cvtColor(image, rgb_img_, cv::COLOR_RGBA2RGB);
 	cv::resize(rgb_img_, rgb_img_, cv::Size(input_width, input_height));
 	auto img_scale_x = static_cast<float>( image.cols) / input_width;
 	auto img_scale_y = static_cast<float>( image.rows) / input_height;
 
 	auto input_tensor = mat2tensor(rgb_img_);
 	std::vector<torch::jit::IValue> inputs;
 	inputs.emplace_back(input_tensor);
 
 	auto output = model_.forward(inputs).toTuple()->elements()[0].toTensor().squeeze(0);
 
 	output2results(output, img_scale_x, img_scale_y);
 	return non_max_suppression(results_, iou_threshold, nms_limit);
 }

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

В качестве аргумента метод принимает объект матрицы OpenCV, который представляет собой RGBA-изображение. В начале я определил константы, которые представляют ширину и высоту входных данных модели. Использовал значения 640x640, потому что модель YOLO обучена на изображениях такого размера. Я изменяю размер входного изображения, подгоняя его под константы. Также я удалил альфа-канал, создав RGB-изображение. 

Далее я рассчитал коэффициенты масштабирования с помощью исходных размеров изображения. Они будут использоваться для масштабирования границ обнаруженных объектов до размера исходного изображения. 

Изменив размер изображения, я преобразовывал матрицу OpenCV в тензор PyTorch, используя функцию mat2tensor, реализацию которой мы обсудим позже. Тензор PyTorch добавляем в контейнер значений torch::jit::IValue, приведение типов тут выполняется автоматически. В контейнере inputs у нас единственный элемент, поскольку модель YOLO принимает на вход только изображение RGB. Конечно, несколько изображений можно организовать в батч. 

Затем я использовал функцию forward объекта model_ для выполнения вывода. Для модуля TorchScript API эта функция возвращает тип torch:: jit::IValue. Он является кортежем (tuple), поэтому я явно преобразовал возвращенный объект к типу torch:: jit:: Tuple и использовал его первый элемент. Этот элемент я привел к типу torch::Tensor и изменил его размерность, удаляя из него размерность батча с помощью метода squeeze. 

И в результате получил объект output типа torch::Tensor размером Nx85. Где N — количество обнаруженных объектов. А 85 означает 80 значений оценок принадлежности к классу, 4 координаты ограничивающей рамки (x, y, ширина, высота) и 1 балл достоверности предсказания модели. Тензор output, содержащий результат, преобразуется в контейнер объектов типа DetectionResult в методе output2results. Обнаруженных объектов может быть достаточно много, так как результаты могут перекрываться на плоскости. Далее я использовал метод non_max_suppression для выбора наилучших результатов обнаружения.

Как создать PyTorch-тензор

Для преобразования матрицы изображения в тензор я использовал функцию mat2tensor. Она реализуется следующим образом:

torch::Tensor mat2tensor(const cv::Mat &image) {
 	torch::Tensor tensor_image = torch::from_blob(image.data,
                                               	{1, image.rows, image.cols, image.channels()},
                                               	at::kByte);
 
 	tensor_image = tensor_image.to(at::kFloat) / 255.;
 
 	tensor_image = torch::transpose(tensor_image, 1, 2);
 	tensor_image = torch::transpose(tensor_image, 1, 3);
 	return tensor_image;
 }

Сначала я использовал функцию torch::from_blob для создания объекта типа torch::Tensor на основе сырых данных. Указатель на данные взял из объекта OpenCV, воспользовавшись полем data. Формат использованной размерности [высота, ширина, каналы] соответствует формату, используемому в OpenCV, где последнее измерение — это количество каналов. Для тензора я также добавил и указал размерность батча. Затем я перевел тензор в тип float и нормализовал его к [0,1]. Так как PyTorch и модель YOLO используют другую компоновку размерностей [батч, каналы, высота, ширина], то я соответствующим образом транспонировал тензорные каналы.

Преобразование результирующего тензора в контейнер результатов

Следующая функция — output2results, она преобразует выходной тензорный объект в контейнер структур DetectionResult. Он реализуется так:

void TorchYolo::output2results(const torch::Tensor &output,
                            	float img_scale_x,
                            	float img_scale_y) {
 	auto outputs = output.accessor<float, 2>();
 	auto output_row = output.size(0);
 	auto output_column = output.size(1);
 	results_.clear();
 	for (int64_t i = 0; i < output_row; i++) {
     	auto score = outputs[i][4];
     	if (score > confidence_threshold) {
         	// вычислить ограничивающий прямоугольник
         	// вычислить идентификатор класса
         	results_.push_back(DetectionResult {
                                     	.class_index=cls,
                                     	.class_name=classes_[cls],
                                     	.score = score,
                                     	.rect=cv::Rect(left, top, bw, bh),
                                 	});
     	}
 	}
 }

Вначале я использовал метод accessor<float, 2> объекта torch::Tensor, чтобы получить объект accessor для тензора. Этот объект позволяет использовать оператор квадратных скобок для доступа к элементам в многомерном тензоре. Число 2 в параметрах шаблона означает, что тензор двумерный. Затем я сделал цикл, который проходит по строкам тензора, потому что каждая строка соответствует одному результату обнаружения. Внутри цикла я выполняю следующие шаги:

1. Получаю оценку достоверности из элемента с индексом 4.

2. Продолжаю обработку строки, если оценка достоверности превышает пороговое значение.

3. Вычисляю ограничивающий прямоугольник.

4. Вычисляю идентификатор класса.

5. Создаю структуру DetectionResult с полученными значениями и добавляю ее в контейнер results_.

Расчет ограничивающего прямоугольника выполняется следующим образом:

float cx = outputs[i][0];
 float cy = outputs[i][1];
 float w = outputs[i][2];
 float h = outputs[i][3];
 int left = static_cast<int>(img_scale_x * (cx - w / 2));
 int top = static_cast<int>(img_scale_y * (cy - h / 2));
 int bw = static_cast<int>(img_scale_x * w);
 int bh = static_cast<int>(img_scale_y * h);

Получил элементы 0, 1, 2, 3, которые являются координатами [x, y, ширина, высота]. И, используя ранее рассчитанные масштабные коэффициенты, преобразовывал эти координаты в формат [слева, сверху, ширина, высота]. Масштабирование нужно для возврата в систему координат исходного изображения, размер которого отличается от поданного в модель.

Выбор идентификатора класса реализован следующим образом:

float max = outputs[i][5];
 int cls = 0;
 for (int64_t j = 0; j < output_column - 5; j++) {
 	if (outputs[i][5 + j] > max) {
     	max = outputs[i][5 + j];
     	cls = static_cast<int>(j);
 	}
 }

В цикле проходил по элементам 5–84 — вероятностям классов — и выбирал класс с максимальным значением. Индекс элемента с максимальной вероятностью использовал в качестве идентификатора класса.

Далее рассмотрим функции этого метода, нужные для уточнения результатов обнаружения.

Реализация inference для NCNN

Инференс для фреймфорка NCNN реализован в классе NcnnYolo. В репозитории данный класс поддерживает работу с двумя моделями: YOLOv5 и оптимизированной, но менее точной YOLOv4. 

В этом материале я показываю особенности работы с фреймворком на примере YOLOv4. Для YOLOv5 нужен более сложный код обработки выходного тензора для преобразования его в контейнер результатов, но это особенности работы с конкретной моделью, а не работы с NCNN. 

Итак, начнем также с реализации конструктора:

NcnnYolo::NcnnYolo(AAssetManager *asset_manager, Stats *stats, bool v4)
     	: stats_(stats)
     	, is_yolo4_(v4) {
 	model_.opt.use_vulkan_compute = true;
 	model_.opt.use_fp16_arithmetic = true;
 	model_.opt.num_threads = 8;
 
 	if (model_.load_param(asset_manager, "yolo-fastest-opt.param")) {
     	throw std::runtime_error("Can't load NCNN params");
 	}
 	if (model_.load_model(asset_manager,  "yolo-fastest-opt.bin")) {
     	throw std::runtime_error("Can't load NCNN model");
 	}
 	const std::string classes_file_name = "classes.txt";
 	auto classes_buf = read_asset(asset_manager, classes_file_name);
 	VectorStreamBuf<char> stream_buf(classes_buf);
 	std::istream is(&stream_buf);
 	classes_ = load_classes(is);
 }

Перед загрузкой модели я установил параметры оптимизации — это важно сделать вначале. У NCNN есть много разных настроек оптимизации, но подбирать оптимальный набор для конкретного устройства придется самостоятельно. Часть настроек уже выставлена, и из коробки тоже будет работать неплохо. Но, например, если на вашем устройстве есть GPU, то можно попробовать включить поддержку вычислений используя Vulkan. 

Важную роль играет и количество CPU-потоков для вычислений, поэтому стоит ознакомится со списком доступных настроек и поэкспериментировать с их комбинациями. Например, на телефоне POCO F3, который я использовал, есть GPU Adreno 650, а в NCNN есть набор настроек конкретно под эту GPU. Тем не менее, мне не удалось получить увеличение производительности при их использовании. 

А вот использование Vulkan действительно немного повысило производительность. После установки настроек производительности я загружаю файл *.param с параметрами модели и описанием архитектуры. А потом и файл *.bin с конкретными весами модели. Для загрузки использую методы load_param и load_model объекта model_ с типом \ncnn::Net. Удобно, что работа с загрузкой файлов из папки assets уже реализована в NCNN, и достаточно передать имя файла и указатель на объект AAssetManager. После я загружаю словарь классов уже знакомым способом — с помощью функции load_classes, так же, как делал для PyTorch.

Detect из интерфейса ObjectDetector — основной метод, в котором происходит, он реализуется следующим образом:

...
 // in header
 ncnn::Mat in_;
 ncnn::Mat output_;
 ...
 std::vector<DetectionResult> NcnnYolo::detect(const cv::Mat &image) {
 	img_width_ = image.cols;
 	img_height_ = image.rows;
 
 	mat2tensor_nopad(image);
 
 	auto extractor = model_.create_extractor();
 	extractor.input(0, in_);
 	extractor.extract("output", output_);
 
 	output2results(output8_);
 	return non_max_suppression(results_, iou_threshold, nms_limit);
 }

Реализация этого метода точно такая же, как в классе TorchYolo. Отличий лишь несколько: 

• я сохраняю размеры входного изображения, 

• для вывода модели используется объект extractor с методом extract вместо метода forward. 

Перед заполнением входных параметров модели я преобразовал входное изображение в объект типа ncnn::Mat, используя функцию mat2tensor_nopad. Она сохраняет значение в поле in_. 

Для передачи входных параметров используется метод input, объекта extractor, который принимает индекс входного параметра и его значение по ссылке. Вместо индекса можно использовать строковое имя. Если входных параметров несколько, то надо установить каждый. Для получения выходного значения вывода модели метод extract также принимает индекс или имя соответствующего тензора. Также эти функции принимают объекты тензоров по ссылке, что позволяет не выделять память каждый раз. 

После получения выходного тензора я его преобразую в контейнер объектов DetectionResult и передаю на уточнение в функцию non_max_suppression. Ее рассмотрим далее.

Создание NCNN-тензора

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

void NcnnYolo::mat2tensor_nopad(const cv::Mat &image) {
 	constexpr int target_width = 320;
 	constexpr int target_height = 320;
 	in_ = ncnn::Mat::from_pixels_resize(image.data,
 	ncnn::Mat::PIXEL_RGBA2RGB,
 	image.cols, image.rows,
 	target_width, target_height);
 	const float norm_vals[3] = {1 / 255.f, 1 / 255.f, 1 / 255.f};
 	in_.substract_mean_normalize(nullptr, norm_vals);
 }

Как я уже писал, модель оптимизирована и настроена на работу с разрешением 320х320, поэтому вначале я задавал соответствующие константы. Фреймворк NCNN предоставляет набор функций для создания тензоров из сырых данных. Поэтому далее я использовал функцию ncnn::Mat::from_pixels_resize для преобразования OpenCV-матрицы в объект типа ncnn::Mat. 

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

В NCNN есть много функций для преобразования изображений, поэтому стоит изучить документацию и исходный код — возможно, что-то подойдет конкретно для вашей задачи. После получения изображения в виде объекта ncnn::Mat я привел значения к диапазону [0, 1]. На этом преобразование изображения закончено.

Преобразование результирующего тензора в контейнер результатов

Следующая функция, которую я использовал — output2results. Она преобразует выходной тензорный объект в контейнер структур DetectionResult. Реализуется так:

void NcnnYolo::output2results(const ncnn::Mat &output) {
 	for (int i = 0; i < output.h; i++) {
     	const float *row = output.row(i);
     	auto score = row[1];
     	if (score > confidence_threshold) {
         	auto x1 = row[2] * static_cast<float>(img_width_);
         	auto y1 = row[3] * static_cast<float>(img_height_);
         	auto x2 = row[4] * static_cast<float>(img_width_);
         	auto y2 = row[5] * static_cast<float>(img_height_);
         	int cls = static_cast<int>(row[0] - 1);
         	results_.push_back(
         	DetectionResult{
             	.class_index=cls,
             	.class_name=classes_[cls],
             	.score = score,
             	.rect=cv::Rect(static_cast<int>(x1), static_cast<int>(y1),
             	static_cast<int>(x2 - x1), static_cast<int>(y2 - y1)),
         	});
     	}
 	}
 }

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

Сначала я получаю оценку достоверности по индексу 1 и сравниваю с пороговым значением. Если достоверность результата оценена моделью как достаточно большая, я продолжаю разбор результата. Значение по индексу 0 — это индекс класса, но в диапазоне 1–80, поэтому я уменьшил его на единицу. А значения с 2 по 5 — это коэффициенты для восстановления координаты для верхнего левого и правого нижнего углов ограничивающего прямоугольника. 

Изображения надо умножить на соответствующие размеры исходного изображения, чтобы привести их в систему координат. Эти размеры я специально сохранил в методе detect. Такой метод представления координат дает менее точные ограничивающие прямоугольники. Далее я просто заполняю объект структуры DetectionResult и добавляю его к результирующему контейнеру.

Уточнение результатов обнаружения объектов

Non-Maximum-Suppression (NMS) и Intersection-Over-Union (IOU) — ключевые алгоритмы YOLO для уточнения и фильтрации полученных результатов обнаружения. NМS используется для удаления повторяющихся результатов, которые накладываются друг на друга. Это работает путем сравнения предсказанных ограничивающих прямоугольников и удаления тех, которые сильное совпадают с другими. Например, если для одного и того же объекта предсказаны два ограничивающих прямоугольника, то NMS сохранит только тот, что имеет наибольшую оценку достоверности, а остальные удалит.

IOU — еще один алгоритм, используемый совместно с NMS для измерения перекрытия между ограничивающими прямоугольниками. IOU вычисляет отношение площади пересечения к площади объединения между двумя прямоугольниками. Получаемое значение находится в диапазоне от 0 до 1, где 0 означает отсутствие перекрытия, а 1 указывает на идеальное перекрытие. В коде я использовал два пороговых значения: одно — для первичной фильтрации результатов по значению оценки достоверности в методах output2results, а второе — для фильтра значений IOU. 

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

Реализация алгоритма Non-Maximum-Suppression (NMS)

Я реализовал NMS в методе non_max_suppression:

std::vector<DetectionResult> non_max_suppression(std::vector<DetectionResult> &detections,
                                              	float threshold,
                                              	size_t nms_limit) {
 	std::sort(detections.begin(), detections.end(), [](const auto &r1, const auto &r2) {
     	return r1.score > r2.score;
 	});
 	std::vector<DetectionResult> selected;
 	std::vector<bool> active(detections.size(), true);
 	int num_active = static_cast<int>(active.size());
 	bool done = false;
 	for (size_t i = 0; i < detections.size() && !done; i++) {
     	if (active[i]) {
         	const auto &box_a = detections[i];
         	selected.push_back(box_a);
         	if (selected.size() >= nms_limit)
             	break;
         	for (size_t j = i + 1; j < detections.size(); j++) {
             	if (active[j]) {
                 	const auto &box_b = detections[j];
                 	auto iou = IOU(box_a.rect, box_b.rect);
                 	if (iou > threshold) {
                     	active[j] = false;
                     	num_active -= 1;
                     	if (num_active <= 0) {
                         	done = true;
                         	break;
                     	}
                 	}
             	}
         	}
     	}
 	}
 	return selected;
 }

Сначала я отсортировал все результаты обнаружения по значению оценки достоверности в порядке убывания. Отметил все результаты как активные. Если результат обнаружения активен, то его можно сравнивать с другим. В противном случае результат уже подавлен (suppressed). 

Затем каждый активный результат обнаружения я последовательно сравнил с остальными активными результатами. Помните, что контейнер отсортирован. Сравнение выполняется путем вычисления значения IOU для ограничивающих прямоугольников и сравнения значения IOU с пороговым значением. 

Если значение IOU превышает пороговое значение, то результат с меньшим значением достоверности помечается как неактивный, то есть мы подавляем его (suppress). Во внешнем цикле подавленные результаты тоже игнорируются. Во вложенных циклах также есть проверка на максимально допустимое количество результатов, обратите внимание на использование значения nms_limit.

Реализация алгоритма Intersection-Over-Union(IOU)

Алгоритм IOU для двух ограничивающих прямоугольников я реализовал так:

float IOU(const cv::Rect &a, const cv::Rect &b) {
 	if (a.empty()) return 0.0f;
 	if (b.empty()) return 0.0f;
 	auto min_x = std::max(a.x, b.x);
 	auto min_y = std::max(a.y, b.y);
 	auto max_x = std::min(a.x + a.width, b.x + b.width);
 	auto max_y = std::min(a.y + a.height, b.y + b.height);
 	auto area = std::max(max_y - min_y, 0) * std::max(max_x - min_x, 0);
 	return static_cast<float>(area) / static_cast<float>(a.area() + b.area() - area);
 }

Сначала проверил, что ограничивающие прямоугольники не пустые. Если хотя бы один из них пуст, то значение IOU равно нулю. 

Затем вычислил площадь пересечения путем нахождения минимального и максимального значений для X и Y с учетом размеров обеих ограничивающих прямоугольников, а также нашел произведение разницы между этими значениями. 

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

NMS и IoU помогают повысить точность обнаружения для моделей YOLO, устраняя ложные срабатывания и гарантируя, что в конечный результат будут включены максимально корректные обнаружения.

Какой фреймворк стоит выбрать

По итогам работы над приложением я собрал таблицу с замерами производительности фреймворков на устройстве POCO 5G c процессором Snapdragon 870SM8250 и GPU Adreno 650.

Надо сказать, что модель NCNN YOLOv4, несмотря на отличную производительность, обладает значительно худшей точностью. Она использует уменьшенное разрешение для входных изображение, и это одна из существенных причин ее хорошей производительности. 

Также использование Vulkan немного повышает производительность — это видно по результатам модели NCNN YOLOv5. Сравнить производительность преобразований OpenCV и NCNN не получилось, в таблице видны разные значения для mat2tensor.

Отмечу, что мне фреймворк NCNN показался более перспективным для использования на Android по нескольким параметрам:

• возможность тонкой настройки,

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

• интеграция с системой Android, например загрузка данных из asset. 

К тому же, использование OpenCV для работы с изображениями существенно облегчает прототипирование и разработку.

Что еще изучить

• Компьютерное зрение на C++: пишем приложение для поиска объектов под Android

• PyTorch C++ API

• Документация для разработчиков приложений

• Android NDK

• Руководство PyTroch для мобильной разработки

• Руководство PyTroch для экспорта оптимизированных мобильных скриптов

• ExcuTorch – новый фреймворк для запуска PyTorch на встроенных устройствах

• Сборка NCNN для Android

• Подготовка моделей для NCNN