pipnet.h

// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once
#include "fastdeploy/fastdeploy_model.h"
#include "fastdeploy/vision/common/processors/transform.h"
#include "fastdeploy/vision/common/result.h"

namespace fastdeploy {

namespace vision {

namespace facealign {
class FASTDEPLOY_DECL PIPNet : public FastDeployModel {
 public:
  PIPNet(const std::string& model_file, const std::string& params_file = "",
         const RuntimeOption& custom_option = RuntimeOption(),
         const ModelFormat& model_format = ModelFormat::ONNX);

  std::string ModelName() const { return "PIPNet"; }
  virtual bool Predict(cv::Mat* im, FaceAlignmentResult* result);

  int GetNumLandmarks() {return num_landmarks_; }
  std::vector<float> GetMeanVals() { return mean_vals_; }
  std::vector<float> GetStdVals() { return std_vals_; }
  std::vector<int> GetSize() { return size_; }
  void SetNumLandmarks(const int& num_landmarks);
  void SetMeanVals(const std::vector<float>& mean_vals) {
       mean_vals_ = mean_vals;
  }
  void SetStdVals(const std::vector<float>& std_vals) { std_vals_ = std_vals; }
  void SetSize(const std::vector<int>& size) { size_ = size; }

 private:
  bool Initialize();

  bool Preprocess(Mat* mat, FDTensor* outputs,
                  std::map<std::string, std::array<int, 2>>* im_info);

  bool Postprocess(std::vector<FDTensor>& infer_result,
                   FaceAlignmentResult* result,
                   const std::map<std::string, std::array<int, 2>>& im_info);
  void GenerateLandmarks(std::vector<FDTensor>& infer_result,
                         FaceAlignmentResult* result,
                         float img_height, float img_width);
  std::map<int, int> num_lms_map_;
  std::map<int, int> max_len_map_;
  std::map<int, std::vector<int>> reverse_index1_map_;
  std::map<int, std::vector<int>> reverse_index2_map_;
  int num_nb_;
  int net_stride_;
  // Now PIPNet support num_landmarks in {19, 29, 68, 98}
  std::vector<int> supported_num_landmarks_;
  // tuple of (width, height), default (256, 256)
  std::vector<int> size_;

  // Mean parameters for normalize, size should be the the same as channels,
  // default mean_vals = {0.485f, 0.456f, 0.406f}
  std::vector<float> mean_vals_;
  // Std parameters for normalize, size should be the the same as channels,
  // default std_vals = {0.229f, 0.224f, 0.225f}
  std::vector<float> std_vals_;
  // number of landmarks
  int num_landmarks_;
};

}  // namespace facealign
}  // namespace vision
}  // namespace fastdeploy