model_runner_builder.py

Go to the documentation of this file.

 
from smtrack.builder.runner_builder import build_track_runner
from smtrack.builder.runner_builder import __runner_builders__ as track_runner_builders
from smdetect.builder.runner_builder import build_detect_runner
from smdetect.builder.runner_builder import __runner_builders__ as det_runner_builders
from smpose.builder.runner_builder import build_pose_runner
from smpose.builder.runner_builder import __runner_builders__ as pose_runner_builders
from smaction.builder.runner_builder import build_action_runner
from smaction.builder.runner_builder import __runner_builders__ as action_runner_builders
from smrunner.utils import load_model
 
from mmengine.config import Config
 
#TODO: 현재 사용 안함. 사용 안하는게 확실해 지면 제거
__model_builder__ = {
    "ByteTrackerRunner"     : build_track_runner,
    "MMDetPyTorchRunner"    : build_detect_runner,
    "MMDetTRTRunner"        : build_detect_runner,
    "YoloV8Runner"          : build_detect_runner,
    "MMPoseTRTRunner"       : build_pose_runner,
    "MMActionPyTorchRunner" : build_action_runner,
    "STGCNRunner"           : build_action_runner,
    "MCSTGCNRunner"         : build_action_runner,
}
 
 
def build_model(cfg, device='cpu'):    
    """
    세이프모션 라이브러리의 모델들을 빌드하는 기능
    action, detect, pose, track 모델들을 빌드할 수 있음
    args:
        cfg (str or Config): 모델의 config 파일 경로 또는 mmengine.config.Config.fromfile()로 생성한 객체
        device (str): 모델이 구동될 디바이스
    return : 세이프모션 라이브러리의 모델
    """
 
    #문자열로 입력 받을 경우 mmengine.config.Config.fromfile() cfg 생성
    if isinstance(cfg, str):
        cfg = Config.fromfile(cfg)
    
    model = cfg.model #모델 파라미터
    model_args = model.copy() #복사
 
    #모델 빌드
    if model.type in det_runner_builders: #객체 검출 모델  빌드
        model = build_detect_runner(model_args)
        
    elif model.type in pose_runner_builders: #포즈 추정 모델  빌드
        model = build_pose_runner(model_args)
 
    elif model.type in action_runner_builders: #행동 인식 모델  빌드
        model = build_action_runner(model_args)
        if cfg.test.model_path is not None:
            load_model(model, cfg.test.model_path, device=device)
 
    elif model.type in track_runner_builders: #추적 모델 빌드
        model = build_track_runner(model_args)
    else:
        return None
    
    return model
    
    # assert model.type in __model_builder__, \
    #         f'not found model type : {model.type}'
    
    # model_builder = __model_builder__[model.type]
    # model_args = model.copy()
    # return model_builder(model_args)