runner_builder.py

Go to the documentation of this file.

from smdetect.runner.mmdet_pytorch_runner import MMDetPyTorchRunner
from smdetect.runner.mmdet_trt_runner import MMDetTRTRunner
from smdetect.runner.yolov8_runner import YoloV8Runner
 
__runner_builders__ = {
    "MMDetPyTorchRunner": MMDetPyTorchRunner,
    "MMDetTRTRunner"    : MMDetTRTRunner,
    "YoloV8Runner"      : YoloV8Runner,
}
 
def build_detect_runner(args):
    """
    detect/runner 의 모델을 빌드하는 함수
    현재 detect/runner의 모델은 오픈소스기반으로 구현되어 있음
 
    args : 
        type -> 모델의 타입  __runner_builders__ 의 키값이 들어가야함
        type을 제외한 나머지 파라미터는 모델의 파라미터값이 있어야함
            세부적으로 모듈의 타입과 모듈의 파라미터가 정의되어 있어야함
    
    returns:
        파이토이 모델 or TensorRT 모델
    """
    assert args.type in __runner_builders__, \
            f'not found detection model type : {args.type}'
    
    runner_builder = __runner_builders__[args.type]
    runner_args = args.copy()
    runner_args.pop('type')
    return runner_builder(**runner_args)