mmdet_trt_runner.MMDetTRTRunner Class Reference

Inheritance diagram for mmdet_trt_runner.MMDetTRTRunner:

Public Member Functions
	__init__ (self, deploy_cfg, model_cfg, backend_model, device='cuda')
	run_detector (self, image)

Public Attributes
	deploy_cfg
	model_cfg
	task_processor
	model
	input_shape

Detailed Description

mmdetection의 Deploy(TensorRT) 모델을 구동시키기 위한 클래스 
args:
    deploy_cfg (str): mmtection의 deploy 내용이 정의된 config 파일 경로
    model_cfg (str): mmtection의 모델이 정의되 config 파일 경로
    backend_model (str): TenosrRT 모델이 저장된 파일 경로
    device (str): 모델이 구동될 디바이스

Definition at line 5 of file mmdet_trt_runner.py.

Constructor & Destructor Documentation

__init__()

mmdet_trt_runner.MMDetTRTRunner.__init__	(		self,
			deploy_cfg,
			model_cfg,
			backend_model,
			device = 'cuda' )

Definition at line 14 of file mmdet_trt_runner.py.

    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)
 

Member Function Documentation

run_detector()

mmdet_trt_runner.MMDetTRTRunner.run_detector	(		self,
			image )

객체 검출 모델을 구동시키는 기능
args:
    image (np.array): RGB 이미지
return (dict): 객체 검출 결과

Definition at line 21 of file mmdet_trt_runner.py.

    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]
        

Member Data Documentation

deploy_cfg

mmdet_trt_runner.MMDetTRTRunner.deploy_cfg

Definition at line 16 of file mmdet_trt_runner.py.

input_shape

mmdet_trt_runner.MMDetTRTRunner.input_shape

Definition at line 19 of file mmdet_trt_runner.py.

model

mmdet_trt_runner.MMDetTRTRunner.model

Definition at line 18 of file mmdet_trt_runner.py.

model_cfg

mmdet_trt_runner.MMDetTRTRunner.model_cfg

Definition at line 16 of file mmdet_trt_runner.py.

task_processor

mmdet_trt_runner.MMDetTRTRunner.task_processor

Definition at line 17 of file mmdet_trt_runner.py.

---

The documentation for this class was generated from the following file:

• smdetect/runner/mmdet_trt_runner.py