基于C++的TNN分割代码
时间: 2024-03-28 13:41:45 浏览: 145
以下是一个基于C++的TNN分割代码的示例:
```
#include <tnn/core/macro.h>
#include <tnn/utils/blob_converter.h>
#include <tnn/utils/dims_utils.h>
#include <tnn/utils/omp_utils.h>
#include <tnn/utils/cpu_utils.h>
#include <tnn/device/cpu/cpu_device.h>
#include <tnn/interpreter/default_model_interpreter.h>
#include <tnn/interpreter/layer_resource_generator.h>
#include <tnn/network/tensorrt/tensorrt_network.h>
#include <tnn/utils/string_utils_inner.h>
#include <tnn/utils/string_utils.h>
#include <tnn/utils/mat_utils.h>
#include <string>
#include <iostream>
#include <opencv2/opencv.hpp>
using namespace TNN_NS;
using namespace cv;
int main(int argc, char** argv) {
if (argc != 3) {
std::cerr << "Usage: " << argv[0] << " proto_file model_file" << std::endl;
return 1;
}
// create network
auto proto_content = fdLoadFile(argv[1]);
auto model_content = fdLoadFile(argv[2]);
auto interpreter = std::make_shared<DefaultModelInterpreter>();
auto network = interpreter->CreateNetwork(proto_content, model_content);
CHECK_TNN_OK(network->SetCpuNumThreads(4));
CHECK_TNN_OK(network->SetDeviceType(DeviceType::DEVICE_ARM));
// get input info
auto input_info_map = network->GetInputInfo();
auto input_name = input_info_map.begin()->first;
auto input_info = input_info_map.begin()->second;
auto input_dims = input_info->shape;
// prepare input blob
auto input_data_type = input_info->data_type;
auto input_blob = std::make_shared<Blob>(input_data_type, input_dims);
auto input_mat = cv::imread("input.jpg");
auto input_mat_resize = Mat(input_dims[2], input_dims[3], CV_32FC3);
cv::resize(input_mat, input_mat_resize, input_mat_resize.size(), 0, 0, cv::INTER_LINEAR);
auto input_data = input_blob->GetHandle().base;
ConvertFromMat(input_mat_resize, input_data, input_dims);
// forward
std::vector<std::shared_ptr<Blob>> input_blobs = {input_blob};
auto status = network->Forward(input_blobs);
if (status != TNN_OK) {
std::cerr << "Forward Failed: " << TNN_NS::StatusUtils::StatusDescription(status) << std::endl;
return 2;
}
// get output info
auto output_info_map = network->GetOutputInfo();
auto output_name = output_info_map.begin()->first;
auto output_info = output_info_map.begin()->second;
auto output_dims = output_info->shape;
// prepare output blob
auto output_data_type = output_info->data_type;
auto output_blob = std::make_shared<Blob>(output_data_type, output_dims);
auto output_data = output_blob->GetHandle().base;
// get output blob data
auto output_size = DimsVectorUtils::Count(output_dims);
auto output_data_size = DataTypeUtils::GetBytesSize(output_data_type) * output_size;
std::memcpy(output_data, output_info->resource->data, output_data_size);
// postprocess
std::vector<float> output_data_vec(output_size);
ConvertFromRawData(output_data, output_data_vec.data(), output_data_type, output_size);
auto output_mat = Mat(output_dims[2], output_dims[3], CV_32FC1, output_data_vec.data());
cv::threshold(output_mat, output_mat, 0.5, 255, cv::THRESH_BINARY);
// save output
imwrite("output.jpg", output_mat);
return 0;
}
```
该示例中,使用一个给定的TNN模型进行图像分割。需要提供模型的 `.tnnproto` 和 `.tnnmodel` 文件。然后使用 OpenCV 加载输入图像,将其转换为输入 Blob,并在网络上执行正向传递。最后,将输出 Blob 转换为 OpenCV Mat,进行二值化并保存为输出图像。
阅读全文