diff --git a/mediapipe/calculators/tensorflow/BUILD b/mediapipe/calculators/tensorflow/BUILD index 0f8f8706..4aec15dc 100644 --- a/mediapipe/calculators/tensorflow/BUILD +++ b/mediapipe/calculators/tensorflow/BUILD @@ -1054,6 +1054,7 @@ cc_test( "//mediapipe/framework:calculator_framework", "//mediapipe/framework:calculator_runner", "//mediapipe/framework/port:gtest_main", + "//mediapipe/util:packet_test_util", "@org_tensorflow//tensorflow/core:framework", "@org_tensorflow//tensorflow/core:protos_all_cc", ], diff --git a/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator.cc b/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator.cc index cd807b87..ec7cd70f 100644 --- a/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator.cc +++ b/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator.cc @@ -102,7 +102,7 @@ absl::Status TensorToVectorFloatCalculator::Process(CalculatorContext* cc) { } auto output = absl::make_unique>(input_tensor.NumElements()); - const auto& tensor_values = input_tensor.flat(); + const auto& tensor_values = input_tensor.unaligned_flat(); for (int i = 0; i < input_tensor.NumElements(); ++i) { output->at(i) = tensor_values(i); } diff --git a/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator_test.cc b/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator_test.cc index 69d3af60..98ba4f02 100644 --- a/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator_test.cc +++ b/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator_test.cc @@ -16,6 +16,7 @@ #include "mediapipe/framework/calculator_framework.h" #include "mediapipe/framework/calculator_runner.h" #include "mediapipe/framework/port/gtest.h" +#include "mediapipe/util/packet_test_util.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/types.pb.h" @@ -129,5 +130,28 @@ TEST_F(TensorToVectorFloatCalculatorTest, FlattenShouldTakeAllDimensions) { } } +TEST_F(TensorToVectorFloatCalculatorTest, AcceptsUnalignedTensors) { + SetUpRunner(/*tensor_is_2d=*/false, /*flatten_nd=*/false); + + const tf::TensorShape tensor_shape(std::vector{2, 5}); + tf::Tensor tensor(tf::DT_FLOAT, tensor_shape); + auto slice = tensor.Slice(1, 1).flat(); + for (int i = 0; i < 5; ++i) { + slice(i) = i; + } + + auto input_tensor = tensor.SubSlice(1); + // Ensure that the input tensor is unaligned. + ASSERT_FALSE(input_tensor.IsAligned()); + runner_->MutableInputs()->Index(0).packets.push_back( + MakePacket(input_tensor).At(Timestamp(5))); + + ASSERT_TRUE(runner_->Run().ok()); + + EXPECT_THAT(runner_->Outputs().Index(0).packets, + ElementsAre(PacketContainsTimestampAndPayload>( + Timestamp(5), std::vector({0, 1, 2, 3, 4})))); +} + } // namespace } // namespace mediapipe