File size: 2,120 Bytes
7ad7e4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | 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)
|