[Model] Add FSANet model (#448)

* add yolov5cls

* fixed bugs

* fixed bugs

* fixed preprocess bug

* add yolov5cls readme

* deal with comments

* Add YOLOv5Cls Note

* add yolov5cls test

* add rvm support

* support rvm model

* add rvm demo

* fixed bugs

* add rvm readme

* add TRT support

* add trt support

* add rvm test

* add EXPORT.md

* rename export.md

* rm poros doxyen

* deal with comments

* deal with comments

* add rvm video_mode note

* add fsanet

* fixed bug

* update readme

* fixed for ci

* deal with comments

* deal with comments

* deal with comments

Co-authored-by: Jason <jiangjiajun@baidu.com>
Co-authored-by: DefTruth <31974251+DefTruth@users.noreply.github.com>

docs/api/vision_results/README.md

@@ -14,3 +14,4 @@ FastDeploy根据视觉模型的任务类型，定义了不同的结构体(`fastd
 | MattingResult           | [C++/Python文档](./matting_result.md)           | 图片/视频抠图返回结果      | MODNet、RVM系列模型等         |
 | OCRResult               | [C++/Python文档](./ocr_result.md)               | 文本框检测，分类和文本识别返回结果 | OCR系列模型等                |
 | MOTResult               | [C++/Python文档](./mot_result.md)               | 多目标跟踪返回结果         | pptracking系列模型等         |
+| HeadPoseResult               | [C++/Python文档](./headpose_result.md)               | 头部姿态估计返回结果         | FSANet系列模型等         |

docs/api/vision_results/headpose_result.md

@@ -0,0 +1,25 @@
+# HeadPoseResult 头部姿态结果
+
+HeadPoseResult 代码定义在`fastdeploy/vision/common/result.h`中，用于表明头部姿态结果。
+
+## C++ 定义
+
+`fastdeploy::vision::HeadPoseResult`
+
+```c++
+struct HeadPoseResult {
+  std::vector<float> euler_angles;
+  void Clear();
+  std::string Str();
+};
+```
+
+- **euler_angles**: 成员变量，表示单张人脸图片预测的欧拉角，存放的顺序是(yaw, pitch, roll)， yaw 代表水平转角，pitch 代表垂直角，roll 代表翻滚角，值域都为 [-90,+90]度
+- **Clear()**: 成员函数，用于清除结构体中存储的结果
+- **Str()**: 成员函数，将结构体中的信息以字符串形式输出（用于Debug）
+
+## Python 定义
+
+`fastdeploy.vision.HeadPoseResult`
+
+- **euler_angles**(list of float): 成员变量，表示单张人脸图片预测的欧拉角，存放的顺序是(yaw, pitch, roll)， yaw 代表水平转角，pitch 代表垂直角，roll 代表翻滚角，值域都为 [-90,+90]度

examples/CMakeLists.txt

@@ -49,7 +49,7 @@ function(add_fastdeploy_executable FIELD CC_FILE)
     add_executable(${TEMP_TARGET_NAME} ${TEMP_TARGET_FILE})
     target_link_libraries(${TEMP_TARGET_NAME} PUBLIC fastdeploy)
     if(TARGET gflags)
-      if(NOT ANDROID)
+      if(UNIX)
         target_link_libraries(${TEMP_TARGET_NAME} PRIVATE gflags pthread)
       else()
         target_link_libraries(${TEMP_TARGET_NAME} PRIVATE gflags)

examples/vision/README.md

@@ -2,17 +2,19 @@
 
 本目录下提供了各类视觉模型的部署，主要涵盖以下任务类型
 
-| 任务类型              | 说明                                              | 预测结果结构体                                                                              |
+| 任务类型           | 说明                                  | 预测结果结构体                                                                          |
-|:------------------|:------------------------------------------------|:-------------------------------------------------------------------------------------|
+|:-------------- |:----------------------------------- |:-------------------------------------------------------------------------------- |
-| Detection         | 目标检测，输入图像，检测图像中物体位置，并返回检测框坐标及类别和置信度             | [DetectionResult](../../docs/api/vision_results/detection_result.md)                 |
+| Detection      | 目标检测，输入图像，检测图像中物体位置，并返回检测框坐标及类别和置信度 | [DetectionResult](../../docs/api/vision_results/detection_result.md)       |
-| Segmentation      | 语义分割，输入图像，给出图像中每个像素的分类及置信度                      | [SegmentationResult](../../docs/api/vision_results/segmentation_result.md)           |
+| Segmentation   | 语义分割，输入图像，给出图像中每个像素的分类及置信度          | [SegmentationResult](../../docs/api/vision_results/segmentation_result.md) |
-| Classification    | 图像分类，输入图像，给出图像的分类结果和置信度                         | [ClassifyResult](../../docs/api/vision_results/classification_result.md)             |
+| Classification | 图像分类，输入图像，给出图像的分类结果和置信度             | [ClassifyResult](../../docs/api/vision_results/classification_result.md)   |
-| FaceDetection     | 人脸检测，输入图像，检测图像中人脸位置，并返回检测框坐标及人脸关键点              | [FaceDetectionResult](../../docs/api/vision_results/face_detection_result.md)        |
+| FaceDetection | 人脸检测，输入图像，检测图像中人脸位置，并返回检测框坐标及人脸关键点             | [FaceDetectionResult](../../docs/api/vision_results/face_detection_result.md)   |
-| KeypointDetection | 关键点检测，输入图像，返回图像中人物行为的各个关键点坐标和置信度                | [KeyPointDetectionResult](../../docs/api/vision_results/keypointdetection_result.md) |
+| FaceAlignment |  人脸对齐(人脸关键点检测)，输入图像，返回人脸关键点            | [FaceAlignmentResult](../../docs/api/vision_results/face_alignment_result.md)   |
-| FaceRecognition   | 人脸识别，输入图像，返回可用于相似度计算的人脸特征的embedding             | [FaceRecognitionResult](../../docs/api/vision_results/face_recognition_result.md)    |
+| KeypointDetection   | 关键点检测，输入图像，返回图像中人物行为的各个关键点坐标和置信度         | [KeyPointDetectionResult](../../docs/api/vision_results/keypointdetection_result.md) |
-| Matting           | 抠图，输入图像，返回图片的前景每个像素点的Alpha值                     | [MattingResult](../../docs/api/vision_results/matting_result.md)                     |
+| FaceRecognition | 人脸识别，输入图像，返回可用于相似度计算的人脸特征的embedding            | [FaceRecognitionResult](../../docs/api/vision_results/face_recognition_result.md)   |
-| OCR               | 文本框检测，分类，文本框内容识别，输入图像，返回文本框坐标，文本框的方向类别以及框内的文本内容 | [OCRResult](../../docs/api/vision_results/ocr_result.md)                             |
+| Matting | 抠图，输入图像，返回图片的前景每个像素点的Alpha值            | [MattingResult](../../docs/api/vision_results/matting_result.md)   |
-| MOT               | 多目标跟踪，输入图像，检测图像中物体位置，并返回检测框坐标，对象id及类别置信度        | [MOTResult](../../docs/api/vision_results/mot_result.md)                             |
+| OCR | 文本框检测，分类，文本框内容识别，输入图像，返回文本框坐标，文本框的方向类别以及框内的文本内容            | [OCRResult](../../docs/api/vision_results/ocr_result.md)   |
+| MOT | 多目标跟踪，输入图像，检测图像中物体位置，并返回检测框坐标，对象id及类别置信度        | [MOTResult](../../docs/api/vision_results/mot_result.md)   |
+| HeadPose | 头部姿态估计，返回头部欧拉角            | [HeadPoseResult](../../docs/api/vision_results/headpose_result.md)   |
 
 ## FastDeploy API设计
 

examples/vision/facealign/pfld/cpp/CMakeLists.txt

@@ -11,4 +11,8 @@ include_directories(${FASTDEPLOY_INCS})
 
 add_executable(infer_demo ${PROJECT_SOURCE_DIR}/infer.cc)
 # 添加FastDeploy库依赖
-target_link_libraries(infer_demo ${FASTDEPLOY_LIBS} gflags pthread)
+if(UNIX)
+  target_link_libraries(infer_demo ${FASTDEPLOY_LIBS} gflags pthread)
+else()
+  target_link_libraries(infer_demo ${FASTDEPLOY_LIBS} gflags)
+endif()

examples/vision/facealign/pfld/python/README.md

@@ -16,7 +16,6 @@ cd FastDeploy/examples/vision/facealign/pfld/python
 ## 原版ONNX模型
 wget https://bj.bcebos.com/paddlehub/fastdeploy/pfld-106-lite.onnx
 wget https://bj.bcebos.com/paddlehub/fastdeploy/facealign_input.png
-
 # CPU推理
 python infer.py --model pfld-106-lite.onnx --image facealign_input.png --device cpu
 # GPU推理

examples/vision/facealign/pfld/python/infer.py

@@ -17,11 +17,11 @@ def parse_arguments():
     parser.add_argument(
         "--backend",
         type=str,
-        default="ort",
+        default="default",
-        help="inference backend, ort, ov, trt, paddle, paddle_trt.")
+        help="inference backend, default, ort, ov, trt, paddle, paddle_trt.")
     parser.add_argument(
         "--enable_trt_fp16",
-        type=bool,
+        type=ast.literal_eval,
         default=False,
         help="whether enable fp16 in trt/paddle_trt backend")
     return parser.parse_args()

examples/vision/headpose/README.md

@@ -0,0 +1,7 @@
+# 头部姿态模型
+
+FastDeploy目前支持如下人脸对齐模型部署
+
+| 模型 | 说明 | 模型格式 | 版本 |
+| :--- | :--- | :------- | :--- |
+| [omasaht/headpose-fsanet-pytorch](./fsanet) | FSANet 系列模型 | ONNX | [CommitID:002549c](https://github.com/omasaht/headpose-fsanet-pytorch/commit/002549c) |

examples/vision/headpose/fsanet/README.md

@@ -0,0 +1,25 @@
+# FSANet 模型部署
+
+## 模型版本说明
+
+- [FSANet](https://github.com/omasaht/headpose-fsanet-pytorch/commit/002549c)
+
+## 支持模型列表
+
+目前FastDeploy支持如下模型的部署
+
+- [FSANet 模型](https://github.com/omasaht/headpose-fsanet-pytorch)
+
+## 下载预训练模型
+
+为了方便开发者的测试，下面提供了PFLD导出的各系列模型，开发者可直接下载使用。
+
+| 模型                                                               | 参数大小    | 精度    | 备注 |
+|:---------------------------------------------------------------- |:----- |:----- | :------ |
+| [fsanet-1x1.onnx](https://bj.bcebos.com/paddlehub/fastdeploy/fsanet-1x1.onnx) | 1.2M | - |
+| [fsanet-var.onnx](https://bj.bcebos.com/paddlehub/fastdeploy/fsanet-var.onnx) | 1.2MB | - |
+
+## 详细部署文档
+
+- [Python部署](python)
+- [C++部署](cpp)

examples/vision/headpose/fsanet/cpp/CMakeLists.txt

@@ -0,0 +1,18 @@
+PROJECT(infer_demo C CXX)
+CMAKE_MINIMUM_REQUIRED (VERSION 3.10)
+
+# 指定下载解压后的fastdeploy库路径
+option(FASTDEPLOY_INSTALL_DIR "Path of downloaded fastdeploy sdk.")
+include(${FASTDEPLOY_INSTALL_DIR}/utils/gflags.cmake)
+include(${FASTDEPLOY_INSTALL_DIR}/FastDeploy.cmake)
+
+# 添加FastDeploy依赖头文件
+include_directories(${FASTDEPLOY_INCS})
+
+add_executable(infer_demo ${PROJECT_SOURCE_DIR}/infer.cc)
+# 添加FastDeploy库依赖
+if(UNIX)
+  target_link_libraries(infer_demo ${FASTDEPLOY_LIBS} gflags pthread)
+else()
+  target_link_libraries(infer_demo ${FASTDEPLOY_LIBS} gflags)
+endif()

examples/vision/headpose/fsanet/cpp/README.md

@@ -0,0 +1,74 @@
+# FSANet C++部署示例
+
+本目录下提供`infer.cc`快速完成FSANet在CPU/GPU，以及GPU上通过TensorRT加速部署的示例。
+
+在部署前，需确认以下两个步骤
+
+- 1. 软硬件环境满足要求，参考[FastDeploy环境要求](../../../../../docs/cn/build_and_install/download_prebuilt_libraries.md)  
+- 2. 根据开发环境，下载预编译部署库和samples代码，参考[FastDeploy预编译库](../../../../../docs/cn/build_and_install/download_prebuilt_libraries.md)
+
+以Linux上CPU推理为例，在本目录执行如下命令即可完成编译测试，保证 FastDeploy 版本0.6.0以上(x.x.x >= 0.6.0)支持FSANet模型
+
+```bash
+mkdir build
+cd build
+wget https://bj.bcebos.com/fastdeploy/release/cpp/fastdeploy-linux-x64-x.x.x.tgz
+tar xvf fastdeploy-linux-x64-x.x.x.tgz
+cmake .. -DFASTDEPLOY_INSTALL_DIR=${PWD}/fastdeploy-linux-x64-x.x.x
+make -j
+
+#下载官方转换好的 FSANet 模型文件和测试图片
+wget https://bj.bcebos.com/paddlehub/fastdeploy/fsanet-var.onnx
+wget https://bj.bcebos.com/paddlehub/fastdeploy/headpose_input.png
+# CPU推理
+./infer_demo --model fsanet-var.onnx --image headpose_input.png --device cpu
+# GPU推理
+./infer_demo --model fsanet-var.onnx --image headpose_input.png --device gpu
+# GPU上TensorRT推理
+./infer_demo --model fsanet-var.onnx --image headpose_input.png --device gpu --backend trt
+```
+
+运行完成可视化结果如下图所示
+
+<div width="520">
+<img width="500" height="514" float="left" src="https://user-images.githubusercontent.com/19977378/198279932-3eee424e-98a2-4249-bdeb-0f79127cbc9d.png">
+</div>
+
+以上命令只适用于Linux或MacOS, Windows下SDK的使用方式请参考:  
+- [如何在Windows中使用FastDeploy C++ SDK](../../../../../docs/cn/faq/use_sdk_on_windows.md)
+
+## FSANet C++接口
+
+### FSANet 类
+
+```c++
+fastdeploy::vision::headpose::FSANet(
+        const string& model_file,
+        const string& params_file = "",
+        const RuntimeOption& runtime_option = RuntimeOption(),
+        const ModelFormat& model_format = ModelFormat::ONNX)
+```
+FSANet模型加载和初始化，其中model_file为导出的ONNX模型格式。
+**参数**
+> * **model_file**(str): 模型文件路径
+> * **params_file**(str): 参数文件路径，当模型格式为ONNX时，此参数传入空字符串即可
+> * **runtime_option**(RuntimeOption): 后端推理配置，默认为None，即采用默认配置
+> * **model_format**(ModelFormat): 模型格式，默认为ONNX格式
+#### Predict函数
+> ```c++
+> FSANet::Predict(cv::Mat* im, HeadPoseResult* result)
+> ```
+>
+> 模型预测接口，输入图像直接输出头部姿态预测结果。
+>
+> **参数**
+>
+> > * **im**: 输入图像，注意需为HWC，BGR格式
+> > * **result**: 头部姿态预测结果, HeadPoseResult说明参考[视觉模型预测结果](../../../../../docs/api/vision_results/)
+### 类成员变量
+用户可按照自己的实际需求，修改下列预处理参数，从而影响最终的推理和部署效果
+> > * **size**(vector&lt;int&gt;): 通过此参数修改预处理过程中resize的大小，包含两个整型元素，表示[width, height], 默认值为[112, 112]
+- [模型介绍](../../)
+- [Python部署](../python)
+- [视觉模型预测结果](../../../../../docs/api/vision_results/)
+- [如何切换模型推理后端引擎](../../../../../docs/cn/faq/how_to_change_backend.md)

examples/vision/headpose/fsanet/cpp/infer.cc

@@ -0,0 +1,110 @@
+// 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.
+
+#include "fastdeploy/vision.h"
+#include "gflags/gflags.h"
+
+DEFINE_string(model, "", "Directory of the inference model.");
+DEFINE_string(image, "", "Path of the image file.");
+DEFINE_string(device, "cpu",
+              "Type of inference device, support 'cpu' or 'gpu'.");
+DEFINE_string(backend, "default",
+              "The inference runtime backend, support: ['default', 'ort', "
+              "'paddle', 'ov', 'trt', 'paddle_trt']");
+DEFINE_bool(use_fp16, false, "Whether to use FP16 mode, only support 'trt' and 'paddle_trt' backend");
+
+void PrintUsage() {
+  std::cout << "Usage: infer_demo --model model_path --image img_path --device [cpu|gpu] --backend "
+               "[default|ort|paddle|ov|trt|paddle_trt] "
+               "--use_fp16 false"
+            << std::endl;
+  std::cout << "Default value of device: cpu" << std::endl;
+  std::cout << "Default value of backend: default" << std::endl;
+  std::cout << "Default value of use_fp16: false" << std::endl;
+}
+
+bool CreateRuntimeOption(fastdeploy::RuntimeOption* option) {
+  if (FLAGS_device == "gpu") {
+    option->UseGpu();
+    if (FLAGS_backend == "ort") {
+      option->UseOrtBackend();
+    } else if (FLAGS_backend == "paddle") {
+      option->UsePaddleBackend();
+    } else if (FLAGS_backend == "trt" || 
+               FLAGS_backend == "paddle_trt") {
+      option->UseTrtBackend();
+      option.SetTrtInputShape("images", {1, 3, 64, 64});
+      if (FLAGS_backend == "paddle_trt") {
+        option->EnablePaddleToTrt();
+      }
+      if (FLAGS_use_fp16) {
+        option->EnableTrtFP16();
+      }
+    } else if (FLAGS_backend == "default") {
+      return true;
+    } else {
+      std::cout << "While inference with GPU, only support default/ort/paddle/trt/paddle_trt now, " << FLAG_backend << " is not supported." << std::endl;
+      return false;
+    }
+  } else if (FLAGS_device == "cpu") {
+    if (FLAGS_backend == "ort") {
+      option->UseOrtBackend();
+    } else if (FLAGS_backend == "ov") {
+      option->UseOpenVINOBackend();
+    } else if (FLAGS_backend == "paddle") {
+      option->UsePaddleBackend();
+    } else if (FLAGS_backend == "default") {
+      return true;
+    } else {
+      std::cout << "While inference with CPU, only support default/ort/ov/paddle now, " << FLAG_backend << " is not supported." << std::endl;
+      return false;
+    }
+  } else {
+    std::cerr << "Only support device CPU/GPU now, "  << FLAGS_device << " is not supported." << std::endl;
+    return false;
+  }
+
+  return true;
+}
+
+int main(int argc, char* argv[]) {
+  google::ParseCommandLineFlags(&argc, &argv, true);
+  auto option = fastdeploy::RuntimeOption();
+  if (!CreateRuntimeOption(&option)) {
+    PrintUsage();
+    return -1;
+  }
+
+  auto model = fastdeploy::vision::headpose::FSANet(FLAGS_model, "", option);
+  if (!model.Initialized()) {
+    std::cerr << "Failed to initialize." << std::endl;
+    return -1;
+  }
+
+  auto im = cv::imread(FLAGS_image);
+  auto im_bak = im.clone();
+
+  fastdeploy::vision::HeadPoseResult res;
+  if (!model.Predict(&im, &res)) {
+    std::cerr << "Failed to predict." << std::endl;
+    return -1;
+  }
+  std::cout << res.Str() << std::endl;
+
+  auto vis_im = fastdeploy::vision::VisHeadPose(im_bak, res);
+  cv::imwrite("vis_result.jpg", vis_im);
+  std::cout << "Visualized result saved in ./vis_result.jpg" << std::endl;
+  
+  return 0;
+}

examples/vision/headpose/fsanet/python/README.md

@@ -0,0 +1,67 @@
+# FSANet Python部署示例
+
+在部署前，需确认以下两个步骤
+
+- 1. 软硬件环境满足要求，参考[FastDeploy环境要求](../../../../../docs/cn/build_and_install/download_prebuilt_libraries.md)  
+- 2. FastDeploy Python whl包安装，参考[FastDeploy Python安装](../../../../../docs/cn/build_and_install/download_prebuilt_libraries.md)
+
+本目录下提供`infer.py`快速完成FSANet在CPU/GPU，以及GPU上通过TensorRT加速部署的示例，保证 FastDeploy 版本 >= 0.6.0 支持FSANet模型。执行如下脚本即可完成
+
+```bash
+#下载部署示例代码
+git clone https://github.com/PaddlePaddle/FastDeploy.git
+cd FastDeploy/examples/vision/headpose/fsanet/python
+
+# 下载FSANet模型文件和测试图片
+## 原版ONNX模型
+wget https://bj.bcebos.com/paddlehub/fastdeploy/fsanet-var.onnx
+wget https://bj.bcebos.com/paddlehub/fastdeploy/headpose_input.png
+# CPU推理
+python infer.py --model fsanet-var.onnx --image headpose_input.png --device cpu
+# GPU推理
+python infer.py --model fsanet-var.onnx --image headpose_input.png --device gpu
+# TRT推理
+python infer.py --model fsanet-var.onnx --image headpose_input.png --device gpu --backend trt
+```
+
+运行完成可视化结果如下图所示
+
+<div width="520">
+<img width="500" height="514" float="left" src="https://user-images.githubusercontent.com/19977378/198279932-3eee424e-98a2-4249-bdeb-0f79127cbc9d.png">
+</div>
+
+## FSANet Python接口
+
+```python
+fd.vision.headpose.FSANet(model_file, params_file=None, runtime_option=None, model_format=ModelFormat.ONNX)
+```
+
+FSANet 模型加载和初始化，其中model_file为导出的ONNX模型格式
+
+**参数**
+
+> * **model_file**(str): 模型文件路径
+> * **params_file**(str): 参数文件路径，当模型格式为ONNX格式时，此参数无需设定
+> * **runtime_option**(RuntimeOption): 后端推理配置，默认为None，即采用默认配置
+> * **model_format**(ModelFormat): 模型格式，默认为ONNX
+### predict函数
+
+> ```python
+> FSANet.predict(input_image)
+> ```
+>
+> 模型预测结口，输入图像直接输出头部姿态预测结果。
+>
+> **参数**
+>
+> > * **input_image**(np.ndarray): 输入数据，注意需为HWC，BGR格式
+> **返回**
+>
+> > 返回`fastdeploy.vision.HeadPoseResult`结构体，结构体说明参考文档[视觉模型预测结果](../../../../../docs/api/vision_results/)
+
+## 其它文档
+
+- [FSANet 模型介绍](..)
+- [FSANet C++部署](../cpp)
+- [模型预测结果说明](../../../../../docs/api/vision_results/)
+- [如何切换模型推理后端引擎](../../../../../docs/cn/faq/how_to_change_backend.md)

examples/vision/headpose/fsanet/python/infer.py

@@ -0,0 +1,88 @@
+import fastdeploy as fd
+import cv2
+import os
+
+
+def parse_arguments():
+    import argparse
+    import ast
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", required=True, help="Path of FSANet model.")
+    parser.add_argument("--image", type=str, help="Path of test image file.")
+    parser.add_argument(
+        "--device",
+        type=str,
+        default='cpu',
+        help="Type of inference device, support 'cpu' or 'gpu'.")
+    parser.add_argument(
+        "--backend",
+        type=str,
+        default="default",
+        help="inference backend, default, ort, ov, trt, paddle, paddle_trt.")
+    parser.add_argument(
+        "--enable_trt_fp16",
+        type=ast.literal_eval,
+        default=False,
+        help="whether enable fp16 in trt/paddle_trt backend")
+    return parser.parse_args()
+
+
+def build_option(args):
+    option = fd.RuntimeOption()
+    device = args.device
+    backend = args.backend
+    enable_trt_fp16 = args.enable_trt_fp16
+    if device == "gpu":
+        option.use_gpu()
+        if backend == "ort":
+            option.use_ort_backend()
+        elif backend == "paddle":
+            option.use_paddle_backend()
+        elif backend in ["trt", "paddle_trt"]:
+            option.use_trt_backend()
+            option.set_trt_input_shape("input", [1, 3, 64, 64])
+            if backend == "paddle_trt":
+                option.enable_paddle_to_trt()
+            if enable_trt_fp16:
+                option.enable_trt_fp16()
+        elif backend == "default":
+            return option
+        else:
+            raise Exception(
+                "While inference with GPU, only support default/ort/paddle/trt/paddle_trt now, {} is not supported.".
+                format(backend))
+    elif device == "cpu":
+        if backend == "ort":
+            option.use_ort_backend()
+        elif backend == "ov":
+            option.use_openvino_backend()
+        elif backend == "paddle":
+            option.use_paddle_backend()
+        elif backend == "default":
+            return option
+        else:
+            raise Exception(
+                "While inference with CPU, only support default/ort/ov/paddle now, {} is not supported.".
+                format(backend))
+    else:
+        raise Exception(
+            "Only support device CPU/GPU now, {} is not supported.".format(
+                device))
+
+    return option
+
+
+args = parse_arguments()
+
+# 配置runtime，加载模型
+runtime_option = build_option(args)
+model = fd.vision.headpose.FSANet(args.model, runtime_option=runtime_option)
+
+# for image
+im = cv2.imread(args.image)
+result = model.predict(im.copy())
+print(result)
+# 可视化结果
+vis_im = fd.vision.vis_headpose(im, result)
+cv2.imwrite("visualized_result.jpg", vis_im)
+print("Visualized result save in ./visualized_result.jpg")

fastdeploy/vision.h

@@ -51,6 +51,7 @@
 #include "fastdeploy/vision/ocr/ppocr/recognizer.h"
 #include "fastdeploy/vision/segmentation/ppseg/model.h"
 #include "fastdeploy/vision/tracking/pptracking/model.h"
+#include "fastdeploy/vision/headpose/contrib/fsanet.h"
 #endif
 
 #include "fastdeploy/vision/visualize/visualize.h"

fastdeploy/vision/common/result.cc

@@ -485,5 +485,28 @@ std::string OCRResult::Str() {
   return no_result;
 }
 
+void HeadPoseResult::Clear() {
+  std::vector<float>().swap(euler_angles);
+}
+
+void HeadPoseResult::Reserve(int size) {
+  euler_angles.resize(size);
+}
+
+void HeadPoseResult::Resize(int size) {
+  euler_angles.resize(size);
+}
+
+std::string HeadPoseResult::Str() {
+  std::string out;
+
+  out = "HeadPoseResult: [yaw, pitch, roll]\n";
+  out = out + "yaw: " + std::to_string(euler_angles[0]) + "\n" +
+        "pitch: " + std::to_string(euler_angles[1]) + "\n" +
+        "roll: " + std::to_string(euler_angles[2]) + "\n";
+  return out;
+}
+
+
 }  // namespace vision
 }  // namespace fastdeploy

fastdeploy/vision/common/result.h

@@ -33,7 +33,8 @@ enum FASTDEPLOY_DECL ResultType {
   FACE_RECOGNITION,
   MATTING,
   MASK,
-  KEYPOINT_DETECTION
+  KEYPOINT_DETECTION,
+  HEADPOSE,
 };
 
 struct FASTDEPLOY_DECL BaseResult {
@@ -316,6 +317,25 @@ struct FASTDEPLOY_DECL MattingResult : public BaseResult {
   std::string Str();
 };
 
+/*! @brief HeadPose result structure for all the headpose models
+ */
+struct FASTDEPLOY_DECL HeadPoseResult : public BaseResult {
+  /** \brief EulerAngles for an input image, and the element of `euler_angles` is a vector, contains {yaw, pitch, roll}
+   */
+  std::vector<float> euler_angles;
+
+  ResultType type = ResultType::HEADPOSE;
+  /// Clear headpose result
+  void Clear();
+
+  void Reserve(int size);
+
+  void Resize(int size);
+
+  /// Debug function, convert the result to string to print
+  std::string Str();
+};
+
 
 }  // namespace vision
 }  // namespace fastdeploy

fastdeploy/vision/headpose/contrib/fsanet.cc

@@ -0,0 +1,132 @@
+// 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.
+
+#include "fastdeploy/vision/headpose/contrib/fsanet.h"
+#include "fastdeploy/utils/perf.h"
+#include "fastdeploy/vision/utils/utils.h"
+
+namespace fastdeploy {
+
+namespace vision {
+
+namespace headpose {
+
+FSANet::FSANet(const std::string& model_file,
+               const std::string& params_file,
+               const RuntimeOption& custom_option,
+               const ModelFormat& model_format) {
+  if (model_format == ModelFormat::ONNX) {
+    valid_cpu_backends = {Backend::OPENVINO, Backend::ORT}; 
+    valid_gpu_backends = {Backend::ORT, Backend::TRT}; 
+  } else {
+    valid_cpu_backends = {Backend::PDINFER, Backend::ORT};
+    valid_gpu_backends = {Backend::PDINFER, Backend::ORT, Backend::TRT};
+  }
+  runtime_option = custom_option;
+  runtime_option.model_format = model_format;
+  runtime_option.model_file = model_file;
+  runtime_option.params_file = params_file;
+  initialized = Initialize();
+}
+
+bool FSANet::Initialize() {
+  // parameters for preprocess
+  size = {64, 64};
+
+  if (!InitRuntime()) {
+    FDERROR << "Failed to initialize fastdeploy backend." << std::endl;
+    return false;
+  }
+  return true;
+}
+
+bool FSANet::Preprocess(Mat* mat, FDTensor* output,
+                      std::map<std::string, std::array<int, 2>>* im_info) {
+  // Resize
+  int resize_w = size[0];
+  int resize_h = size[1];
+  if (resize_h != mat->Height() || resize_w != mat->Width()) {
+    Resize::Run(mat, resize_w, resize_h);
+  }
+
+  // Normalize
+  std::vector<float> alpha = {1.0f / 128.0f, 1.0f / 128.0f, 1.0f / 128.0f};
+  std::vector<float> beta = {-127.5f / 128.0f, -127.5f / 128.0f, -127.5f / 128.0f};
+  Convert::Run(mat, alpha, beta);
+
+  // Record output shape of preprocessed image
+  (*im_info)["output_shape"] = {mat->Height(), mat->Width()};
+
+  HWC2CHW::Run(mat);
+  Cast::Run(mat, "float");
+
+  mat->ShareWithTensor(output);
+  output->shape.insert(output->shape.begin(), 1);  // reshape to n, h, w, c
+  return true;
+}
+
+bool FSANet::Postprocess(FDTensor& infer_result, HeadPoseResult* result,
+                       const std::map<std::string, std::array<int, 2>>& im_info) {
+  FDASSERT(infer_result.shape[0] == 1, "Only support batch = 1 now.");
+  if (infer_result.dtype != FDDataType::FP32) {
+    FDERROR << "Only support post process with float32 data." << std::endl;
+    return false;
+  }
+
+  auto iter_in = im_info.find("input_shape");
+  FDASSERT(iter_in != im_info.end(),
+           "Cannot find input_shape from im_info.");
+  int in_h = iter_in->second[0];
+  int in_w = iter_in->second[1];
+
+  result->Clear();
+  float* data = static_cast<float*>(infer_result.Data());
+  for (size_t i = 0; i < 3; ++i) {
+    result->euler_angles.emplace_back(data[i]);
+  }
+
+  return true;
+}
+
+bool FSANet::Predict(cv::Mat* im, HeadPoseResult* result) {
+  Mat mat(*im);
+  std::vector<FDTensor> input_tensors(1);
+
+  std::map<std::string, std::array<int, 2>> im_info;
+
+  // Record the shape of image and the shape of preprocessed image
+  im_info["input_shape"] = {mat.Height(), mat.Width()};
+  im_info["output_shape"] = {mat.Height(), mat.Width()};
+
+  if (!Preprocess(&mat, &input_tensors[0], &im_info)) {
+    FDERROR << "Failed to preprocess input image." << std::endl;
+    return false;
+  }
+  input_tensors[0].name = InputInfoOfRuntime(0).name;
+  std::vector<FDTensor> output_tensors;
+  if (!Infer(input_tensors, &output_tensors)) {
+    FDERROR << "Failed to inference." << std::endl;
+    return false;
+  }
+
+  if (!Postprocess(output_tensors[0], result, im_info)) {
+    FDERROR << "Failed to post process." << std::endl;
+    return false;
+  }
+  return true;
+}
+
+}  // namespace headpose
+}  // namespace vision
+}  // namespace fastdeploy

fastdeploy/vision/headpose/contrib/fsanet.h

@@ -0,0 +1,64 @@
+// 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 headpose {
+/*! @brief FSANet model object used when to load a FSANet model exported by FSANet.
+ */
+class FASTDEPLOY_DECL FSANet : public FastDeployModel {
+ public:
+  /** \brief  Set path of model file and the configuration of runtime.
+   *
+   * \param[in] model_file Path of model file, e.g ./fsanet-var.onnx
+   * \param[in] params_file Path of parameter file, e.g ppyoloe/model.pdiparams, if the model format is ONNX, this parameter will be ignored
+   * \param[in] custom_option RuntimeOption for inference, the default will use cpu, and choose the backend defined in "valid_cpu_backends"
+   * \param[in] model_format Model format of the loaded model, default is ONNX format
+   */
+  FSANet(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 "FSANet"; }
+  /** \brief Predict the face detection result for an input image
+   *
+   * \param[in] im The input image data, comes from cv::imread(), is a 3-D array with layout HWC, BGR format
+   * \param[in] result The output face detection result will be writen to this structure
+   * \return true if the prediction successed, otherwise false
+   */
+  virtual bool Predict(cv::Mat* im, HeadPoseResult* result);
+
+  /// tuple of (width, height), default (64, 64)
+  std::vector<int> size;
+
+ private:
+  bool Initialize();
+
+  bool Preprocess(Mat* mat, FDTensor* outputs,
+                  std::map<std::string, std::array<int, 2>>* im_info);
+
+  bool Postprocess(FDTensor& infer_result, HeadPoseResult* result,
+                   const std::map<std::string, std::array<int, 2>>& im_info);
+};
+
+}  // namespace headpose
+}  // namespace vision
+}  // namespace fastdeploy

fastdeploy/vision/headpose/contrib/fsanet_pybind.cc

@@ -0,0 +1,31 @@
+// 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.
+
+#include "fastdeploy/pybind/main.h"
+
+namespace fastdeploy {
+void BindFSANet(pybind11::module& m) {
+  pybind11::class_<vision::headpose::FSANet, FastDeployModel>(m, "FSANet")
+      .def(pybind11::init<std::string, std::string, RuntimeOption,
+                          ModelFormat>())
+      .def("predict",
+           [](vision::headpose::FSANet& self, pybind11::array& data) {
+             auto mat = PyArrayToCvMat(data);
+             vision::HeadPoseResult res;
+             self.Predict(&mat, &res);
+             return res;
+           })
+      .def_readwrite("size", &vision::headpose::FSANet::size);
+}
+}  // namespace fastdeploy

fastdeploy/vision/headpose/headpose_pybind.cc

@@ -0,0 +1,25 @@
+// 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.
+
+#include "fastdeploy/pybind/main.h"
+
+namespace fastdeploy {
+
+void BindFSANet(pybind11::module& m);
+
+void BindHeadPose(pybind11::module& m) {
+  auto headpose_module = m.def_submodule("headpose", "Headpose models.");
+  BindFSANet(headpose_module);
+}
+}  // namespace fastdeploy

fastdeploy/vision/vision_pybind.cc

@@ -26,6 +26,7 @@ void BindFaceId(pybind11::module& m);
 void BindOcr(pybind11::module& m);
 void BindTracking(pybind11::module& m);
 void BindKeyPointDetection(pybind11::module& m);
+void BindHeadPose(pybind11::module& m);
 #ifdef ENABLE_VISION_VISUALIZE
 void BindVisualize(pybind11::module& m);
 #endif
@@ -113,8 +114,7 @@ void BindVision(pybind11::module& m) {
       .def("__repr__", &vision::MattingResult::Str)
       .def("__str__", &vision::MattingResult::Str);
 
-  pybind11::class_<vision::KeyPointDetectionResult>(m,
+  pybind11::class_<vision::KeyPointDetectionResult>(m, "KeyPointDetectionResult")
-                                                    "KeyPointDetectionResult")
       .def(pybind11::init())
       .def_readwrite("keypoints", &vision::KeyPointDetectionResult::keypoints)
       .def_readwrite("scores", &vision::KeyPointDetectionResult::scores)
@@ -122,6 +122,12 @@ void BindVision(pybind11::module& m) {
       .def("__repr__", &vision::KeyPointDetectionResult::Str)
       .def("__str__", &vision::KeyPointDetectionResult::Str);
 
+  pybind11::class_<vision::HeadPoseResult>(m, "HeadPoseResult")
+      .def(pybind11::init())
+      .def_readwrite("euler_angles", &vision::HeadPoseResult::euler_angles)
+      .def("__repr__", &vision::HeadPoseResult::Str)
+      .def("__str__", &vision::HeadPoseResult::Str);
+
   m.def("enable_flycv", &vision::EnableFlyCV, "Enable image preprocessing by FlyCV.");
   m.def("disable_flycv", &vision::DisableFlyCV, "Disable image preprocessing by FlyCV, change to use OpenCV.");
 
@@ -135,6 +141,7 @@ void BindVision(pybind11::module& m) {
   BindOcr(m);
   BindTracking(m);
   BindKeyPointDetection(m);
+  BindHeadPose(m);
 #ifdef ENABLE_VISION_VISUALIZE
   BindVisualize(m);
 #endif

fastdeploy/vision/visualize/headpose.cc

@@ -0,0 +1,59 @@
+// 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.
+
+#ifdef ENABLE_VISION_VISUALIZE
+
+#include "fastdeploy/vision/visualize/visualize.h"
+#include "opencv2/imgproc/imgproc.hpp"
+
+namespace fastdeploy {
+
+namespace vision {
+
+cv::Mat VisHeadPose(const cv::Mat& im, const HeadPoseResult& result,
+                    int size, int line_size) {
+  const float PI = 3.1415926535;
+  auto vis_im = im.clone();
+  int h = im.rows;
+  int w = im.cols;
+  // vis headpose
+  float pitch = result.euler_angles[0] * PI / 180.f;
+  float yaw = -result.euler_angles[1] * PI / 180.f;
+  float roll = result.euler_angles[2] * PI / 180.f;
+
+  int tdx = w / 2;
+  int tdy = h / 2;
+
+  // X-Axis | drawn in red
+  int x1 = static_cast<int>(size * std::cos(yaw) * std::cos(roll)) + tdx;
+  int y1 = static_cast<int>(size * (std::cos(pitch) * std::sin(roll) + 
+                            std::cos(roll) * std::sin(pitch) * std::sin(yaw))) + tdy;
+  // Y-Axis | drawn in green
+  int x2 = static_cast<int>(-size * std::cos(yaw) * std::sin(roll)) + tdx;
+  int y2 = static_cast<int>(size * (std::cos(pitch) * std::cos(roll) - 
+                            std::sin(pitch) * std::sin(yaw) * std::sin(roll))) + tdy;
+  // Z-Axis | drawn in blue
+  int x3 = static_cast<int>(size * std::sin(yaw)) + tdx;
+  int y3 = static_cast<int>(-size * std::cos(yaw) * std::sin(pitch)) + tdy;
+
+  cv::line(vis_im, cv::Point2i(tdx, tdy), cv::Point2i(x1, y1), cv::Scalar(0, 0, 255), line_size);
+  cv::line(vis_im, cv::Point2i(tdx, tdy), cv::Point2i(x2, y2), cv::Scalar(0, 255, 0), line_size);
+  cv::line(vis_im, cv::Point2i(tdx, tdy), cv::Point2i(x3, y3), cv::Scalar(255, 0, 0), line_size);
+  return vis_im;
+}
+
+}  // namespace vision
+}  // namespace fastdeploy
+
+#endif

fastdeploy/vision/visualize/visualize.h

@@ -94,8 +94,12 @@ FASTDEPLOY_DECL cv::Mat SwapBackground(const cv::Mat& im,
                                        const SegmentationResult& result,
                                        int background_label);
 FASTDEPLOY_DECL cv::Mat VisKeypointDetection(const cv::Mat& im,
-                                      const KeyPointDetectionResult& results,
+                        const KeyPointDetectionResult& results,
-                                      float conf_threshold = 0.5f);
+                        float conf_threshold = 0.5f);
+FASTDEPLOY_DECL cv::Mat VisHeadPose(const cv::Mat& im,
+                                    const HeadPoseResult& result,
+                                    int size = 50,
+                                    int line_size = 1);
 
 }  // namespace vision
 }  // namespace fastdeploy

fastdeploy/vision/visualize/visualize_pybind.cc

@@ -102,6 +102,16 @@ void BindVisualize(pybind11::module& m) {
              FDTensor out;
              vision::Mat(vis_im).ShareWithTensor(&out);
              return TensorToPyArray(out);
+           })
+      .def("vis_headpose",
+           [](pybind11::array& im_data, vision::HeadPoseResult& result,
+              int size, int line_size) {
+             auto im = PyArrayToCvMat(im_data);
+             auto vis_im =
+                 vision::VisHeadPose(im, result, size, line_size);
+             FDTensor out;
+             vision::Mat(vis_im).ShareWithTensor(&out);
+             return TensorToPyArray(out);
            });
 
   pybind11::class_<vision::Visualize>(m, "Visualize")

python/fastdeploy/vision/__init__.py

@@ -23,6 +23,7 @@ from . import facedet
 from . import facealign
 from . import faceid
 from . import ocr
+from . import headpose
 from . import evaluation
 from .utils import fd_result_to_json
 from .visualize import *

python/fastdeploy/vision/headpose/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+
+from __future__ import absolute_import
+from .contrib.fsanet import FSANet

python/fastdeploy/vision/headpose/contrib/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+
+from __future__ import absolute_import

python/fastdeploy/vision/headpose/contrib/fsanet.py

@@ -0,0 +1,68 @@
+# 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.
+
+from __future__ import absolute_import
+import logging
+from .... import FastDeployModel, ModelFormat
+from .... import c_lib_wrap as C
+
+
+class FSANet(FastDeployModel):
+    def __init__(self,
+                 model_file,
+                 params_file="",
+                 runtime_option=None,
+                 model_format=ModelFormat.ONNX):
+        """Load a headpose model exported by FSANet.
+
+        :param model_file: (str)Path of model file, e.g fsanet/fsanet-var.onnx
+        :param params_file: (str)Path of parameters file, if the model_fomat is ModelFormat.ONNX, this param will be ignored, can be set as empty string
+        :param runtime_option: (fastdeploy.RuntimeOption)RuntimeOption for inference this model, if it's None, will use the default backend on CPU
+        :param model_format: (fastdeploy.ModelForamt)Model format of the loaded model, default is ONNX
+        """
+
+        super(FSANet, self).__init__(runtime_option)
+
+        assert model_format == ModelFormat.ONNX, "FSANet only support model format of ModelFormat.ONNX now."
+        self._model = C.vision.headpose.FSANet(
+            model_file, params_file, self._runtime_option, model_format)
+        assert self.initialized, "FSANet initialize failed."
+
+    def predict(self, input_image):
+        """Predict an input image headpose
+
+        :param im: (numpy.ndarray)The input image data, 3-D array with layout HWC, BGR format
+        :return: HeadPoseResult
+        """
+
+        return self._model.predict(input_image)
+
+    @property
+    def size(self):
+        """
+        Returns the preprocess image size, default (64, 64)
+        """
+        return self._model.size
+
+    @size.setter
+    def size(self, wh):
+        """
+        Set the preprocess image size, default (64, 64)
+        """
+        assert isinstance(wh, (list, tuple)),\
+            "The value to set `size` must be type of tuple or list."
+        assert len(wh) == 2,\
+            "The value to set `size` must contatins 2 elements means [width, height], but now it contains {} elements.".format(
+            len(wh))
+        self._model.size = wh

python/fastdeploy/vision/visualize/__init__.py

@@ -109,3 +109,7 @@ def vis_ppocr(im_data, det_result):
 
 def vis_mot(im_data, mot_result, score_threshold=0.0, records=None):
     return C.vision.vis_mot(im_data, mot_result, score_threshold, records)
+
+
+def vis_headpose(im_data, headpose_result, size=50, line_size=1):
+    return C.vision.vis_headpose(im_data, headpose_result, size, line_size)