mmpose_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
import numpy as np
 
from smutils.utils_image import crop_image
 
class MMPoseTRTRunner(object):
    def __init__(self, deploy_cfg, model_cfg, backend_model, device='cuda', batch_size=1):
        """
        mmpose의 Deploy(TensorRT) 모델을 구동시키기 위한 클래스 
        args:
            deploy_cfg (str): deploy config 파일 경로
            model_cfg (str): 모델이 정의된 config 파일 경로
            backend_model (list[str]): TensorRT 모델 저장 경로
            device (str): 모델을 구동할 디바이스
            batch_size (int): 모델의 입력 배치 사이즈
        """
        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)
        self.batch_size = batch_size
        self.test_pipeline = self.build_test_pipeline(self.model_cfg)
 
    def build_test_pipeline(self, cfg):
        """
        모델의 입력 생성을 위한 pipeline을 빌드하는 기능
        registry의 TRANSFORMS을 사용함, 미리 빌드를 해두어야함
        mmlab의 다른 클래스(mmdetect)을 함께 사용하면 mmpose를 사용하기 전에 registry의 설정이 추가적으로 필요함
        args:
            cfg : 모델 config 파일을 로드한 객체
        return:
            모델의 입력 생성을 위한 pipeline
        """
        from mmpose.registry import TRANSFORMS
        from mmcv.transforms import Compose
        test_pipeline = [
            TRANSFORMS.build(c) for c in cfg.test_dataloader.dataset.pipeline
        ]
        test_pipeline = Compose(test_pipeline)
        return test_pipeline
 
    def make_src(self, image, image_id, bboxes=None):
        """
        박스영역을 잘라내고 잘라낸 이미지에 대한 정보를 생성하는 기능
        args:
            image (np.array): RGB 이미지
            image_id (int): 이미지 아이디
            bboxes (np.array): 객체 영역의 박스
        return:
            crop_images (list[np.array]): 박스 영역을 자른 이미지
            bboxes (np.array): crop_images의 박스
            track_ids (list[float]): 추적 아이디
            image_ids (list[int]): 이미지 아이디 리스트
        """
        if bboxes is None: #입력 박스가 없을 경우, 영상하나가 입력 데이터임
            h, w = image.shape[:2]
            return [image], np.array([[0, 0, w, h]]), np.array([0]), [image_id]
        
        num_person, dim = bboxes.shape
        box_index = 1 if dim == 6 else 0 #박스 부분의 시작 인덱스
        track_ids = bboxes[:, 0] if dim == 6 else np.arange(num_person) #추적 아이디, dim이 6이 아니면 0부터 각기 다르게 설정
        image_ids = [image_id]*num_person #이미지 아이디는 모두 동일하게 부여됨
        crop_images = crop_image(image, bboxes[:, box_index:box_index+4]) #영상 자르기
 
        return crop_images, bboxes[:, box_index:box_index+4], track_ids, image_ids
 
    def divide_into_batches(self, data, batch_size):
        """
        데이터를 배치 단위로 분할하는 기능
        args:
            data (list): 데이터
            batch_size (int)): 배치 크기
        return (list): 배치 단위로 분할된 데이터
        """
 
        #배치 사이즈가 1이면 그대로 반환
        if batch_size == 1:
            return data
 
        #배치 단위로 분할
        batches = [data[i:i+batch_size] for i in range(0, len(data), batch_size)]
        return batches
 
    def run_detector(self, images, bboxes=None):
        """
        포즈 추정기를 동작시키는 기능
        args:
            image (np.array): RGB 이미지
            bboxes (np.array): 객체 영역의 박스
        """
 
        #입력 타입 체크
        is_batch = isinstance(images, (list, tuple))
 
        #입력 통일 시키기
        images = [images] if not is_batch else images
        if not is_batch:
            bboxes = [bboxes]
        elif bboxes is None:
            bboxes = [None]*len(images)
 
        #데이터 생성, 영역 자르기
        crop_images, bbox_list, track_ids, image_ids = [], [], [], []
        for image_id, (image, bbox) in enumerate(zip(images, bboxes)):
            crops, bbs, tids, imids = self.make_src(image, image_id, bbox)
            crop_images.extend(crops)
            bbox_list.append(bbs)
            track_ids.append(tids)
            image_ids.extend(imids)
 
        bbox_list = np.vstack(bbox_list)
        track_ids = np.hstack(track_ids)
 
        #배치 단위로 분할
        crop_images = self.divide_into_batches(crop_images, self.batch_size)
        with torch.no_grad():
            results = []
 
            for image in crop_images:
                #모델 입력 생성
                processed_input = self.task_processor.create_input(
                    image, 
                    self.input_shape, 
                    test_pipeline=self.test_pipeline
                )[0]
                
                #모델 구동
                inference_result = self.model.test_step(processed_input)
 
                #모델 출력 모으기
                for item in inference_result:
                    results.append(item)
                    
            # results = [item for image in crop_images for item in self.model.test_step(self.task_processor.create_input(image, self.input_shape, test_pipeline=self.test_pipeline)[0])]
 
        #출력 데이터 생성, 추적아이디, 박스, 스켈레톤 데이터가 모두 포함됨
        pose_results = []
        pre_img_id = -1
        for res, bbox, track_id, img_id in zip(results, bbox_list, track_ids, image_ids):
            keypoints = res.pred_instances.keypoints + bbox[:2] #스켈레톤, 원본 이미지 기준으로 좌표 변환
            keypoint_scores = np.expand_dims(res.pred_instances.keypoint_scores, axis=-1) #스켈레톤 스코어
            keypoints = np.dstack((keypoints, keypoint_scores)) #스켈레톤 좌표, 스코어 합치기
            
            # [1, 17, 3] -> [17, 3]
            if keypoints.ndim == 3:
                keypoints = keypoints[0]
            
            #원본 이미지가 다를 경우
            if pre_img_id != img_id:
                pre_img_id = img_id
                pose_results.append([])
 
            #출력 데이터 생성
            data = {'track_id': track_id, 'bbox': bbox, 'keypoints': keypoints}
            pose_results[-1].append(data)
 
        #입력이 원본 이미지 한장일 경우
        if is_batch:
            return pose_results
        else:
            if len(pose_results) == 0:
                return pose_results
            return pose_results[0]