mmdet_trt_runner.py

Go to the documentation of this file.

from mmdeploy.apis.utils import build_task_processor
from mmdeploy.utils import get_input_shape, load_config
import torch
 
class MMDetTRTRunner(object):
    """
    mmdetection의 Deploy(TensorRT) 모델을 구동시키기 위한 클래스 
    args:
        deploy_cfg (str): mmtection의 deploy 내용이 정의된 config 파일 경로
        model_cfg (str): mmtection의 모델이 정의되 config 파일 경로
        backend_model (str): TenosrRT 모델이 저장된 파일 경로
        device (str): 모델이 구동될 디바이스
    """
    def __init__(self, deploy_cfg, model_cfg, backend_model, device='cuda'):
        
        self.deploy_cfg, self.model_cfg = load_config(deploy_cfg, model_cfg)
        self.task_processor = build_task_processor(self.model_cfg, self.deploy_cfg, device)
        self.model = self.task_processor.build_backend_model(backend_model)
        self.input_shape = get_input_shape(deploy_cfg)
 
    def run_detector(self, image):
        """
        객체 검출 모델을 구동시키는 기능
        args:
            image (np.array): RGB 이미지
        return (dict): 객체 검출 결과
        """
 
        #배치 단위인지 채크
        is_batch = False
        if isinstance(image, (list, tuple)):
            is_batch = True
 
        #입력 생성
        input_data, _ = self.task_processor.create_input(image, self.input_shape)
        
        #모델 구동
        with torch.no_grad():
            result = self.model.test_step(input_data)
 
        #출력 포멧 설정, mmaction의 이전 버전과 포맷 매칭
        #TODO: 
        ret_data = []
        for d in result:
            masks = d.pred_instances.masks #TODO : Instance 모델일 경우에만 작동되도록 바꿔야함
            bboxes = d.pred_instances.bboxes.cpu() #박스
            scores = d.pred_instances.scores.cpu() #스코어
            labels = d.pred_instances.labels.cpu() #라벨
            
            #기존 출력 포맷으로 변환
            scores = scores.unsqueeze(1)
            bboxes = torch.cat((bboxes, scores), dim=1)
            data = {}
            data['det_bboxes'] = bboxes
            data['det_labels'] = labels
            data['masks'] = masks
            ret_data.append(data)
 
        if is_batch: #배치 단위
            return ret_data
        else: #이미지 한장
            return ret_data[0]