runner_builder.py

Go to the documentation of this file.

from smpose.runner.mmpose_trt_runner import MMPoseTRTRunner
 
 
__runner_builders__ = {
    "MMPoseTRTRunner": MMPoseTRTRunner,
}
 
def build_pose_runner(args):
    """
    pose/runner 의 모델을 빌드하는 함수
    현재 pose/runner의 모델은 mmpose기반으로 구현되어 있음
    TODO: pytorch 모델 빌드하는 기능 추가 필요
 
    args : 
        type -> 모델의 타입  __runner_builders__ 의 키값이 들어가야함
        type을 제외한 나머지 파라미터는 모델의 파라미터값이 있어야함
            세부적으로 모듈의 타입과 모듈의 파라미터가 정의되어 있어야함
    
    returns:
        파이토이 모델 or TensorRT 모델
    """
    assert args.type in __runner_builders__, \
            f'not found pose model type : {args.type}'
    
    runner_builder = __runner_builders__[args.type]
    runner_args = args.copy()
    runner_args.pop('type')
    return runner_builder(**runner_args)