libtorch vision 基础功能介绍

最新推荐文章于 2024-08-24 21:27:59 发布

GondorFu

最新推荐文章于 2024-08-24 21:27:59 发布

阅读量732

点赞数 1

分类专栏：深度学习

本文链接：https://blog.csdn.net/a40850273/article/details/119727306

版权

深度学习专栏收录该内容

35 篇文章 0 订阅

订阅专栏

libtorch Tensor 元素引用

tensor.index({0, "...", 49, 212}) // 取第 1 维上 (0， *， 49， 212) 上的所有元素

Opencv 读取图像常用转换

torch::from_blob：从 cv::Mat 转 Tensor

cv::cvtColor(frame, frame, CV_BGR2RGB); // step1. BGR->RGB
frame.convertTo(frame, CV_32FC3, 1.0f / 255.0f);  // step2. 归一化
// step3. 转tensor
int img_h = frame.rows;
int img_w = frame.cols;
int depth = frame.channels()
auto input_tensor = torch::from_blob(frame.data, {1, img_h, img_w, depth});
input_tensor = input_tensor.permute({0, 3, 1, 2});  // step4. 通道转换

libtorch 模型加载与推理

libtorch 模型推理输入数据类型：std::vector<torch::jit::IValue>

std::vector<torch::jit::IValue> inputs;
inputs.push_back(input_tensor.to(device_type));
 
at::Tensor output = model.forward(inputs).toTensor();
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

libtorch Tensor 转 cv::Mat

libtorch Tensor 取数据操作：out_tensor.data_ptr()

torch::Tensor out_tensor = input_tensor;
//s1:sequeeze去掉多余维度,(1,C,H,W)->(C,H,W)；s2:permute执行通道顺序调整,(C,H,W)->(H,W,C)
out_tensor = out_tensor.squeeze().detach().permute({ 1, 2, 0 });
out_tensor = out_tensor.mul(255).clamp(0, 255).to(torch::kU8); //s3:*255，转uint8
out_tensor = out_tensor.to(torch::kCPU); //迁移至CPU
cv::Mat resultImg(img_h, img_w, CV_8UC3, out_tensor.data_ptr()); // 将Tensor数据拷贝至Mat
// cv::cvtColor(resultImg, resultImg, CV_RGB2BGR); // convert to BGR
cv::imwrite("ret.jpg", resultImg);

libtorch interpolate

namespace F = torch::nn::functional;
at::Tensor tmp = F::interpolate(out[3], F::InterpolateFuncOptions().mode(torch::kBilinear).size(std::vector<int64_t>({image.rows, image.cols})).align_corners(true));

libtorch softmax max argmax

at::Tensor softmax_tmp = torch::softmax(tmp, 1);
std::tuple<torch::Tensor, torch::Tensor> max_classes = torch::max(softmax_tmp, 1);
auto conf= std::get<0>(max_classes).squeeze().to(torch::kFloat).contiguous();
auto pre_seg= std::get<1>(max_classes).squeeze().to(torch::kU8).contiguous();


at::Tensor pre_seg = at::argmax(tmp, 1).squeeze().to(torch::kU8).contiguous();
cv::Mat pre_seg_mat(cv::Size(image.cols, image.rows), CV_8U, pre_seg.data_ptr());