mmdet_pytorch_runner.py

Go to the documentation of this file.

from mmdet.apis import init_detector, inference_detector
 
class MMDetPyTorchRunner(object):
    """
    mmdetection의 모델을 동작시키기 위한 클래스
    args:
        model_cfg (str): mmdetection의 모델이 정의된 config파일 경로
        model_checkpoint (str): 모델의 파라미터가 저장된 파일 경로
        device (str): 모델이 구동될 디바이스
    """
    def __init__(self, model_cfg, model_checkpoint, device='cuda'):
        
        self.model = init_detector(model_cfg, model_checkpoint, device=device)
 
    def run_detector(self, image):
        """
        객체 검출기를 동작시키는 기능
        args:
            image (np.ndarray): RGB 이미지
        """
 
        result = inference_detector(self.model, image)
 
        return result