cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

mediapipe/calculators/pytorch/pytorch_converter_calculator.cc

Browse Source 
Signed-off-by: Pierre Fenoll <pierrefenoll@gmail.com>
This commit is contained in:
Pierre Fenoll 2020-10-20 00:23:57 +02:00
parent 302ce764f3
commit 0b1987add9
1 changed files with 6 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
mediapipe/calculators/pytorch/pytorch_converter_calculator.cc
View File

@ -45,7 +45,8 @@ constexpr char kImageTag[] = "IMAGE";
constexpr char kImageGpuTag[] = "IMAGE_GPU"; constexpr char kImageGpuTag[] = "IMAGE_GPU";
constexpr char kTensorsTag[] = "TENSORS"; constexpr char kTensorsTag[] = "TENSORS";
using Outputs = std::vector<torch::jit::IValue>; using Output = torch::jit::IValue;
using Outputs = std::vector<Output>;
} // namespace } // namespace
// Calculator for normalizing and converting an ImageFrame // Calculator for normalizing and converting an ImageFrame
@ -183,7 +184,6 @@ REGISTER_CALCULATOR(PyTorchConverterCalculator);
cv::Mat img_float; cv::Mat img_float;
// Normalize to [0;1] // Normalize to [0;1]
image.convertTo(img_float, CV_32F, 1.0 / 255); image.convertTo(img_float, CV_32F, 1.0 / 255);
// FIXME: try NCWH directly here
auto img_tensor = torch::from_blob(img_float.data, {1, width, height, 3}); auto img_tensor = torch::from_blob(img_float.data, {1, width, height, 3});
// Permute from NWHC to NCWH // Permute from NWHC to NCWH
img_tensor = img_tensor.permute({0, 3, 1, 2}); img_tensor = img_tensor.permute({0, 3, 1, 2});
@ -199,10 +199,10 @@ REGISTER_CALCULATOR(PyTorchConverterCalculator);
} }
if (has_tensors_tag_) { if (has_tensors_tag_) {
// FIXME: move to ctor auto output_tensors =
auto output_tensors = absl::make_unique<Outputs>(); absl::make_unique<Outputs, std::initializer_list<Output>>({
output_tensors->reserve(1); img_tensor.cpu(),
output_tensors->emplace_back(img_tensor.cpu()); });
cc->Outputs() cc->Outputs()
.Tag(kTensorsTag) .Tag(kTensorsTag)
.Add(output_tensors.release(), cc->InputTimestamp()); .Add(output_tensors.release(), cc->InputTimestamp());