summaryrefslogtreecommitdiffhomepage
diff options
context:
space:
mode:
-rw-r--r--tracker-neuralnet/ftnoir_tracker_neuralnet.cpp37
-rw-r--r--tracker-neuralnet/ftnoir_tracker_neuralnet.h2
2 files changed, 26 insertions, 13 deletions
diff --git a/tracker-neuralnet/ftnoir_tracker_neuralnet.cpp b/tracker-neuralnet/ftnoir_tracker_neuralnet.cpp
index 50d1e7b4..320c8f23 100644
--- a/tracker-neuralnet/ftnoir_tracker_neuralnet.cpp
+++ b/tracker-neuralnet/ftnoir_tracker_neuralnet.cpp
@@ -105,10 +105,9 @@ mat33 rotation_from_two_vectors(const vec3 &a, const vec3 &b)
}
-/* Computes correction due to head being off screen center.
- x, y: In screen space, i.e. in [-1,1]
- focal_length_x: In screen space
-*/
+// Computes correction due to head being off screen center.
+// x, y: In screen space, i.e. in [-1,1]
+// focal_length_x: In screen space
mat33 compute_rotation_correction(const cv::Point2f &p, float focal_length_x)
{
return rotation_from_two_vectors(
@@ -144,6 +143,20 @@ T iou(const cv::Rect_<T> &a, const cv::Rect_<T> &b)
return double{i.area()} / (a.area()+b.area()-i.area());
}
+// Returns width and height of the input tensor, or throws.
+// Expects the model to take one tensor as input that must
+// have the shape B x C x H x W, where B=C=1.
+cv::Size get_input_image_shape(const Ort::Session &session)
+{
+ if (session.GetInputCount() != 1)
+ throw std::invalid_argument("Model must take exactly one input tensor");
+ const std::vector<std::int64_t> shape =
+ session.GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
+ if (shape.size() != 4)
+ throw std::invalid_argument("Model takes the input tensor in the wrong shape");
+ return { static_cast<int>(shape[3]), static_cast<int>(shape[2]) };
+}
+
} // namespace
@@ -222,13 +235,16 @@ double Localizer::last_inference_time_millis() const
}
-PoseEstimator::PoseEstimator(Ort::MemoryInfo &allocator_info, Ort::Session &&session) :
- session{std::move(session)},
- scaled_frame(input_img_height, input_img_width, CV_8U),
- input_mat(input_img_height, input_img_width, CV_32F)
+PoseEstimator::PoseEstimator(Ort::MemoryInfo &allocator_info, Ort::Session &&_session) :
+ session{std::move(_session)}
{
+ const cv::Size input_image_shape = get_input_image_shape(session);
+
+ scaled_frame = cv::Mat(input_image_shape, CV_8U);
+ input_mat = cv::Mat(input_image_shape, CV_32F);
+
{
- const std::int64_t input_shape[4] = { 1, 1, input_img_height, input_img_width };
+ const std::int64_t input_shape[4] = { 1, 1, input_image_shape.height, input_image_shape.width };
input_val = Ort::Value::CreateTensor<float>(allocator_info, input_mat.ptr<float>(0), input_mat.total(), input_shape, 4);
}
@@ -295,7 +311,7 @@ std::optional<PoseEstimator::Face> PoseEstimator::run(
auto p = input_mat.ptr(0);
- cv::resize(cropped, scaled_frame, { input_img_width, input_img_height }, 0, 0, cv::INTER_AREA);
+ cv::resize(cropped, scaled_frame, scaled_frame.size(), 0, 0, cv::INTER_AREA);
// Automatic brightness amplification.
const int brightness = find_input_intensity_90_pct_quantile();
@@ -306,7 +322,6 @@ std::optional<PoseEstimator::Face> PoseEstimator::run(
assert (input_mat.ptr(0) == p);
assert (!input_mat.empty() && input_mat.isContinuous());
- assert (input_mat.cols == input_img_width && input_mat.rows == input_img_height);
const char* input_names[] = {"x"};
const char* output_names[] = {"pos_size", "quat", "box"};
diff --git a/tracker-neuralnet/ftnoir_tracker_neuralnet.h b/tracker-neuralnet/ftnoir_tracker_neuralnet.h
index 5f9c6fbe..c74a7b1d 100644
--- a/tracker-neuralnet/ftnoir_tracker_neuralnet.h
+++ b/tracker-neuralnet/ftnoir_tracker_neuralnet.h
@@ -118,8 +118,6 @@ class PoseEstimator
private:
// Operates on the private image data members
int find_input_intensity_90_pct_quantile() const;
- inline static constexpr int input_img_width = 129;
- inline static constexpr int input_img_height = 129;
Ort::Session session{nullptr};
// Inputs
cv::Mat scaled_frame{}, input_mat{};