# Copyright (c) Alibaba, Inc. and its affiliates. import copy import json import mmcv import numpy as np import torch from mmcv.image import imwrite from mmcv.utils.path import is_filepath from mmcv.visualization.image import imshow from easycv.core.visualization import imshow_bboxes, imshow_keypoints from easycv.datasets.pose.data_sources.top_down import DatasetInfo from easycv.datasets.pose.pipelines.transforms import bbox_cs2xyxy from easycv.file import io from easycv.predictors.builder import PREDICTORS, build_predictor from easycv.utils.checkpoint import load_checkpoint from easycv.utils.config_tools import mmcv_config_fromfile from easycv.utils.misc import deprecated from .base import InputProcessor, OutputProcessor, PredictorV2 np.set_printoptions(suppress=True) def _box2cs(image_size, box): """This encodes bbox(x,y,w,h) into (center, scale) Args: x, y, w, h Returns: tuple: A tuple containing center and scale. - np.ndarray[float32](2,): Center of the bbox (x, y). - np.ndarray[float32](2,): Scale of the bbox w & h. """ x, y, w, h = box[:4] aspect_ratio = image_size[0] / image_size[1] center = np.array([x + w * 0.5, y + h * 0.5], dtype=np.float32) if w > aspect_ratio * h: h = w * 1.0 / aspect_ratio elif w < aspect_ratio * h: w = h * aspect_ratio # pixel std is 200.0 scale = np.array([w / 200.0, h / 200.0], dtype=np.float32) scale = scale * 1.25 return center, scale def vis_pose_result( model, img, result, radius=4, thickness=1, kpt_score_thr=0.3, bbox_color='green', dataset_info=None, out_file=None, pose_kpt_color=None, pose_link_color=None, text_color='white', font_scale=0.5, bbox_thickness=1, win_name='', show=False, wait_time=0, ): """Visualize the detection results on the image. Args: model (nn.Module): The loaded detector. img (str | np.ndarray): Image filename or loaded image. result (list[dict]): The results to draw over `img` (bbox_result, pose_result). radius (int): Radius of circles. thickness (int): Thickness of lines. kpt_score_thr (float): The threshold to visualize the keypoints. skeleton (list[tuple()]): Default None. out_file (str|None): The filename of the output visualization image. show (bool): Whether to show the image. Default: False. wait_time (int): Value of waitKey param. Default: 0. out_file (str or None): The filename to write the image. Default: None. """ # get dataset info if (dataset_info is None and hasattr(model, 'cfg') and 'dataset_info' in model.cfg): dataset_info = DatasetInfo(model.cfg.dataset_info) if not dataset_info: raise ValueError('Please provide `dataset_info`!') skeleton = dataset_info.skeleton pose_kpt_color = dataset_info.pose_kpt_color pose_link_color = dataset_info.pose_link_color if hasattr(model, 'module'): model = model.module img = mmcv.imread(img) img = img.copy() bbox_result = result.get('bbox', []) pose_result = result['keypoints'] if len(bbox_result) > 0: bboxes = np.vstack(bbox_result) labels = None if 'label' in result: labels = result['label'] # draw bounding boxes imshow_bboxes( img, bboxes, labels=labels, colors=bbox_color, text_color=text_color, thickness=bbox_thickness, font_scale=font_scale, show=False) imshow_keypoints(img, pose_result, skeleton, kpt_score_thr, pose_kpt_color, pose_link_color, radius, thickness) if show: imshow(img, win_name, wait_time) if out_file is not None: imwrite(img, out_file) return img class PoseTopDownInputProcessor(InputProcessor): def __init__(self, cfg, dataset_info, detection_predictor_config, bbox_thr=None, pipelines=None, batch_size=1, cat_id=None, mode='BGR'): self.detection_predictor = build_predictor(detection_predictor_config) self.dataset_info = dataset_info self.bbox_thr = bbox_thr self.cat_id = cat_id super().__init__( cfg, pipelines=pipelines, batch_size=batch_size, threads=1, mode=mode) def get_detection_outputs(self, input, cat_id=None): det_results = self.detection_predictor(input['img'], keep_inputs=False) person_results = self._process_detection_results( det_results, cat_id=cat_id) return person_results def _process_detection_results(self, det_results, cat_id=None): """Process det results, and return a list of bboxes. Args: det_results (list|tuple): det results. cat_id (int | str): category id or name to reserve, if None, reserve all detection results. Returns: person_results (list): a list of detected bounding boxes """ # Only support one sample/image if isinstance(det_results, tuple): det_results = det_results[0] elif isinstance(det_results, list): det_results = det_results[0] else: det_results = det_results bboxes = det_results['detection_boxes'] scores = det_results['detection_scores'] classes = det_results['detection_classes'] if cat_id is not None: if isinstance(cat_id, str): assert cat_id in self.detection_predictor.cfg.CLASSES, f'cat_id "{cat_id}" not in detection classes list: {self.detection_predictor.cfg.CLASSES}' assert det_results.get('detection_class_names', None) is not None detection_class_names = det_results['detection_class_names'] keeped_ids = [ i for i in range(len(detection_class_names)) if str(detection_class_names[i]) == str(cat_id) ] else: keeped_ids = classes == cat_id bboxes = bboxes[keeped_ids] scores = scores[keeped_ids] classes = classes[keeped_ids] person_results = [] for idx, bbox in enumerate(bboxes): person = {} bbox = np.append(bbox, scores[idx]) person['bbox'] = bbox person_results.append(person) return person_results def process_single(self, input): output = super()._load_input(input) person_results = self.get_detection_outputs(output, cat_id=self.cat_id) box_id = 0 # Select bboxes by score threshold bboxes = np.array([res['bbox'] for res in person_results]) if self.bbox_thr is not None: assert bboxes.shape[1] == 5 valid_idx = np.where(bboxes[:, 4] > self.bbox_thr)[0] bboxes = bboxes[valid_idx] person_results = [person_results[i] for i in valid_idx] results = [] for person_result in person_results: box = person_result['bbox'] # x,y,x,y,s boxc = [box[0], box[1], box[2] - box[0], box[3] - box[1]] # x,y,w,h center, scale = _box2cs(self.cfg.data_cfg['image_size'], boxc) data = { 'image_id': 0, 'center': center, 'scale': scale, 'bbox': box, 'bbox_score': box[4] if len(box) == 5 else 1, 'bbox_id': box_id, # need to be assigned if batch_size > 1 'joints_3d': np.zeros((self.cfg.data_cfg.num_joints, 3), dtype=np.float32), 'joints_3d_visible': np.zeros((self.cfg.data_cfg.num_joints, 3), dtype=np.float32), 'rotation': 0, 'flip_pairs': self.dataset_info.flip_pairs, 'ann_info': { 'image_size': np.array(self.cfg.data_cfg['image_size']), 'num_joints': self.cfg.data_cfg['num_joints'], }, 'image_file': output['filename'], 'img': output['img'], 'img_shape': output['img_shape'], 'ori_shape': output['ori_shape'], 'img_fields': output['img_fields'], } box_id += 1 data_processor = self.processor(data) data_processor['bbox'] = box results.append(data_processor) return results def __call__(self, inputs): """Process all inputs list. And collate to batch and put to target device. If you need custom ops to load or process a batch samples, you need to reimplement it. """ batch_outputs = [] for inp in inputs: for res in self.process_single(inp): batch_outputs.append(res) if len(batch_outputs) < 1: return batch_outputs batch_outputs = self._collate_fn(batch_outputs) batch_outputs['img_metas']._data = [[ img_meta[i] for img_meta in batch_outputs['img_metas']._data for i in range(len(img_meta)) ]] return batch_outputs class PoseTopDownOutputProcessor(OutputProcessor): def __call__(self, inputs): output = {} output['keypoints'] = inputs['preds'] output['bbox'] = np.array(inputs['boxes']) # x1, y1, x2, y2 score return output # TODO: Fix when multi people are detected in each sample, # all the people results will be passed to the pose model, # resulting in a dynamic batch_size, which is not supported by jit script model. @PREDICTORS.register_module() class PoseTopDownPredictor(PredictorV2): """Pose topdown predictor. Args: model_path (str): Path of model path. config_file (Optinal[str]): Config file path for model and processor to init. Defaults to None. detection_model_config: Dict of person detection model predictor config, example like ``dict(type="", model_path="", config_file="", ......)`` batch_size (int): Batch size for forward. bbox_thr (float): Bounding box threshold to filter output results of detection model cat_id (int | str): Category id or name to filter target objects. device (str | torch.device): Support str('cuda' or 'cpu') or torch.device, if is None, detect device automatically. save_results (bool): Whether to save predict results. save_path (str): File path for saving results, only valid when `save_results` is True. pipelines (list[dict]): Data pipeline configs. mode (str): The image mode into the model. """ def __init__(self, model_path, config_file=None, detection_predictor_config=None, batch_size=1, bbox_thr=None, cat_id=None, device=None, pipelines=None, save_results=False, save_path=None, mode='BGR', model_type=None, *args, **kwargs): assert batch_size == 1, 'Only support batch_size=1 now!' self.cat_id = cat_id self.bbox_thr = bbox_thr self.detection_predictor_config = detection_predictor_config self.model_type = model_type if self.model_type is None: if model_path.endswith('jit'): assert config_file is not None self.model_type = 'jit' elif model_path.endswith('blade'): import torch_blade assert config_file is not None self.model_type = 'blade' else: self.model_type = 'raw' assert self.model_type in ['raw', 'jit', 'blade'] super(PoseTopDownPredictor, self).__init__( model_path, config_file=config_file, batch_size=batch_size, device=device, save_results=save_results, save_path=save_path, pipelines=pipelines, input_processor_threads=1, mode=mode, *args, **kwargs) if hasattr(self.cfg, 'dataset_info'): dataset_info = self.cfg.dataset_info if is_filepath(dataset_info): cfg = mmcv_config_fromfile(dataset_info) dataset_info = cfg._cfg_dict['dataset_info'] else: from easycv.datasets.pose.data_sources.coco import COCO_DATASET_INFO dataset_info = COCO_DATASET_INFO self.dataset_info = DatasetInfo(dataset_info) def _build_model(self): if self.model_type != 'raw': with io.open(self.model_path, 'rb') as infile: model = torch.jit.load(infile, self.device) else: model = super()._build_model() return model def prepare_model(self): """Build model from config file by default. If the model is not loaded from a configuration file, e.g. torch jit model, you need to reimplement it. """ model = self._build_model() model.to(self.device) model.eval() if self.model_type == 'raw': load_checkpoint(model, self.model_path, map_location='cpu') return model def model_forward(self, inputs, return_heatmap=False): boxes = inputs['bbox'].cpu().numpy() if self.model_type == 'raw': with torch.no_grad(): result = self.model( **inputs, mode='test', return_heatmap=return_heatmap) else: img_metas = inputs['img_metas'] with torch.no_grad(): img = inputs['img'].to(self.device) tensor_img_metas = copy.deepcopy(img_metas) for meta in tensor_img_metas: meta.pop('image_file') for k, v in meta.items(): meta[k] = torch.tensor(v) output_heatmap = self.model(img, tensor_img_metas) from easycv.models.pose.heads.topdown_heatmap_base_head import decode_heatmap output_heatmap = output_heatmap.cpu().numpy() result = decode_heatmap(output_heatmap, img_metas, self.cfg.model.test_cfg) result['boxes'] = np.array(boxes) return result def get_input_processor(self): return PoseTopDownInputProcessor( cfg=self.cfg, dataset_info=self.dataset_info, detection_predictor_config=self.detection_predictor_config, bbox_thr=self.bbox_thr, pipelines=self.pipelines, batch_size=self.batch_size, cat_id=self.cat_id, mode=self.mode) def get_output_processor(self): return PoseTopDownOutputProcessor() def show_result(self, image, keypoints, radius=4, thickness=3, kpt_score_thr=0.3, bbox_color='green', show=False, save_path=None): vis_result = vis_pose_result( self.model, image, keypoints, dataset_info=self.dataset_info, radius=radius, thickness=thickness, kpt_score_thr=kpt_score_thr, bbox_color=bbox_color, show=show, out_file=save_path) return vis_result class _TorchPoseTopDownOutputProcessor(PoseTopDownOutputProcessor): def __call__(self, inputs): output = super(_TorchPoseTopDownOutputProcessor, self).__call__(inputs) bbox = output['bbox'] keypoints = output['keypoints'] results = [] for i in range(len(keypoints)): results.append({'bbox': bbox[i], 'keypoints': keypoints[i]}) return {'pose_results': results} @deprecated(reason='Please use PoseTopDownPredictor.') @PREDICTORS.register_module() class TorchPoseTopDownPredictorWithDetector(PoseTopDownPredictor): def __init__( self, model_path, model_config={ 'pose': { 'bbox_thr': 0.3, 'format': 'xywh' }, 'detection': { 'model_type': None, 'reserved_classes': [], 'score_thresh': 0.0, } }, ): """ init model Args: model_path: pose and detection model file path, split with `,`, make sure the first is pose model, second is detection model model_config: config string for model to init, in json format """ if isinstance(model_config, str): model_config = json.loads(model_config) reserved_classes = model_config['detection'].pop( 'reserved_classes', []) if len(reserved_classes) == 0: reserved_classes = None else: assert len(reserved_classes) == 1 reserved_classes = reserved_classes[0] model_list = model_path.split(',') assert len(model_list) == 2 # first is pose model, second is detection model pose_model_path, detection_model_path = model_list detection_model_type = model_config['detection'].pop('model_type') if detection_model_type == 'TorchYoloXPredictor': detection_predictor_config = dict( type=detection_model_type, model_path=detection_model_path, model_config=model_config['detection']) else: detection_predictor_config = dict( model_path=detection_model_path, **model_config['detection']) pose_kwargs = model_config['pose'] pose_kwargs.pop('format', None) super().__init__( model_path=pose_model_path, detection_predictor_config=detection_predictor_config, cat_id=reserved_classes, **pose_kwargs, ) def get_output_processor(self): return _TorchPoseTopDownOutputProcessor() def show_result(self, image_path, keypoints, radius=4, thickness=1, kpt_score_thr=0.3, bbox_color='green', show=False, save_path=None): dataset_info = self.dataset_info # get dataset info if (dataset_info is None and hasattr(self.model, 'cfg') and 'dataset_info' in self.model.cfg): dataset_info = DatasetInfo(self.model.cfg.dataset_info) if not dataset_info: raise ValueError('Please provide `dataset_info`!') skeleton = dataset_info.skeleton pose_kpt_color = dataset_info.pose_kpt_color pose_link_color = dataset_info.pose_link_color if hasattr(self.model, 'module'): self.model = self.model.module img = self.model.show_result( image_path, keypoints, skeleton, radius=radius, thickness=thickness, pose_kpt_color=pose_kpt_color, pose_link_color=pose_link_color, kpt_score_thr=kpt_score_thr, bbox_color=bbox_color, show=show, out_file=save_path) return img