| import os |
| from rtmlib import RTMO |
|
|
| class RTMO_GPU(RTMO): |
|
|
| def __init__(self, |
| onnx_model: str = None, |
| model_input_size: tuple = (640, 640), |
| mean: tuple = None, |
| std: tuple = None, |
| to_openpose: bool = False, |
| backend: str = 'onnxruntime', |
| device: str = 'cuda'): |
| |
| if backend == 'onnxruntime': |
|
|
| if not os.path.exists(onnx_model): |
| from rtmlib.tools.file import download_checkpoint |
| onnx_model = download_checkpoint(onnx_model) |
|
|
| import onnxruntime as ort |
|
|
| providers = {'cpu': 'CPUExecutionProvider', |
| 'cuda': [ |
| ('CUDAExecutionProvider', { |
| 'device_id': 0, |
| 'arena_extend_strategy': 'kNextPowerOfTwo', |
| 'gpu_mem_limit': 2 * 1024 * 1024 * 1024, |
| 'cudnn_conv_algo_search': 'DEFAULT', |
| 'do_copy_in_default_stream': True, |
| 'enable_cuda_graph': False |
| }), |
| 'CPUExecutionProvider']} |
| |
| self.session = ort.InferenceSession(path_or_bytes=onnx_model, |
| providers=providers[device]) |
|
|
| print(f'load {onnx_model} with {backend} backend') |
|
|
| self.onnx_model = onnx_model |
| self.model_input_size = model_input_size |
| self.mean = mean |
| self.std = std |
| self.backend = backend |
| self.device = device |
| self.to_openpose = to_openpose |
|
|
| else: |
| super().__init__(onnx_model, |
| model_input_size, |
| mean, |
| std, |
| to_openpose, |
| backend, |
| device) |
|
|
