crosstyan/RapidPoseTriangulation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Using TensorRT backend.

Browse Source
This commit is contained in:
Daniel
2024-12-02 13:50:35 +01:00
parent bf837fbca9
commit c8c48e4bf2
2 changed files with 149 additions and 146 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
scripts/utils_2d_pose_ort.py
View File

@ -13,9 +13,12 @@ class BaseModel(ABC):
def __init__(self, model_path: str, warmup: int):
self.opt = ort.SessionOptions()
providers = ort.get_available_providers()
# ort.set_default_logger_severity(1)
provider = ""
if "CUDAExecutionProvider" in providers:
if "TensorrtExecutionProvider" in providers:
provider = "TensorrtExecutionProvider"
elif "CUDAExecutionProvider" in providers:
provider = "CUDAExecutionProvider"
else:
provider = "CPUExecutionProvider"
@ -328,10 +331,10 @@ def load_model():
print("Loading onnx model ...")
model = TopDown(
"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_with-norm.onnx",
# "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_fp16_with-norm.onnx",
"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_with-norm.onnx",
# "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_fp16_with-norm.onnx",
# "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_with-norm.onnx",
"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_fp16_with-norm.onnx",
# "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_with-norm.onnx",
"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_fp16_with-norm.onnx",
conf_threshold=0.3,
iou_threshold=0.3,
warmup=30,