detrsmpl/apis/inference.py

import cv2
import mmcv
import numpy as np
import torch
from mmcv.parallel import collate
from mmcv.runner import load_checkpoint

from detrsmpl.data.datasets.pipelines import Compose
from detrsmpl.models.architectures.builder import build_architecture
from detrsmpl.models.backbones.builder import build_backbone
from detrsmpl.utils.demo_utils import box2cs, xywh2xyxy, xyxy2xywh


def init_model(config, checkpoint=None, device='cuda:0'):
    """Initialize a model from config file.

    Args:
        config (str or :obj:`mmcv.Config`): Config file path or the config
            object.
        checkpoint (str, optional): Checkpoint path. If left as None, the model
            will not load any weights.

    Returns:
        nn.Module: The constructed model.
        (nn.Module, None): The constructed extractor model
    """
    if isinstance(config, str):
        config = mmcv.Config.fromfile(config)
    elif not isinstance(config, mmcv.Config):
        raise TypeError('config must be a filename or Config object, '
                        f'but got {type(config)}')
    config.data.test.test_mode = True

    model = build_architecture(config.model)
    if checkpoint is not None:
        # load model checkpoint
        load_checkpoint(model, checkpoint, map_location=device)
    # save the config in the model for convenience
    model.cfg = config
    model.to(device)
    model.eval()

    extractor = None
    if config.model.type == 'VideoBodyModelEstimator':
        extractor = build_backbone(config.extractor.backbone)
        if config.extractor.checkpoint is not None:
            # load model checkpoint
            load_checkpoint(extractor, config.extractor.checkpoint)
        extractor.cfg = config
        extractor.to(device)
        extractor.eval()
    return model, extractor


class LoadImage:
    """A simple pipeline to load image."""
    def __init__(self, color_type='color', channel_order='bgr'):
        self.color_type = color_type
        self.channel_order = channel_order

    def __call__(self, results):
        """Call function to load images into results.

        Args:
            results (dict): A result dict contains the image_path.

        Returns:
            dict: ``results`` will be returned containing loaded image.
        """
        if isinstance(results['image_path'], str):
            results['image_file'] = results['image_path']
            img = mmcv.imread(results['image_path'], self.color_type,
                              self.channel_order)
        elif isinstance(results['image_path'], np.ndarray):
            results['image_file'] = ''
            if self.color_type == 'color' and self.channel_order == 'rgb':
                img = cv2.cvtColor(results['image_path'], cv2.COLOR_BGR2RGB)
            else:
                img = results['image_path']
        else:
            raise TypeError('"image_path" must be a numpy array or a str or '
                            'a pathlib.Path object')

        results['img'] = img
        return results


def inference_image_based_model(
    model,
    img_or_path,
    det_results,
    bbox_thr=None,
    format='xywh',
):
    """Inference a single image with a list of person bounding boxes.

    Args:
        model (nn.Module): The loaded pose model.
        img_or_path (Union[str, np.ndarray]): Image filename or loaded image.
        det_results (List(dict)): the item in the dict may contain
            'bbox' and/or 'track_id'.
            'bbox' (4, ) or (5, ): The person bounding box, which contains
            4 box coordinates (and score).
            'track_id' (int): The unique id for each human instance.
        bbox_thr (float, optional): Threshold for bounding boxes.
            Only bboxes with higher scores will be fed into the pose detector.
            If bbox_thr is None, ignore it. Defaults to None.
        format (str, optional): bbox format ('xyxy' | 'xywh'). Default: 'xywh'.
            'xyxy' means (left, top, right, bottom),
            'xywh' means (left, top, width, height).

    Returns:
        list[dict]: Each item in the list is a dictionary,
            containing the bbox: (left, top, right, bottom, [score]),
            SMPL parameters, vertices, kp3d, and camera.
    """
    # only two kinds of bbox format is supported.
    assert format in ['xyxy', 'xywh']
    mesh_results = []
    if len(det_results) == 0:
        return []

    # Change for-loop preprocess each bbox to preprocess all bboxes at once.
    bboxes = np.array([box['bbox'] for box in det_results])

    # Select bboxes by score threshold
    if bbox_thr is not None:
        assert bboxes.shape[1] == 5
        valid_idx = np.where(bboxes[:, 4] > bbox_thr)[0]
        bboxes = bboxes[valid_idx]
        det_results = [det_results[i] for i in valid_idx]

    if format == 'xyxy':
        bboxes_xyxy = bboxes
        bboxes_xywh = xyxy2xywh(bboxes)
    else:
        # format is already 'xywh'
        bboxes_xywh = bboxes
        bboxes_xyxy = xywh2xyxy(bboxes)

    # if bbox_thr remove all bounding box
    if len(bboxes_xywh) == 0:
        return []

    cfg = model.cfg
    device = next(model.parameters()).device

    # build the data pipeline
    inference_pipeline = [LoadImage()] + cfg.inference_pipeline
    inference_pipeline = Compose(inference_pipeline)

    assert len(bboxes[0]) in [4, 5]

    batch_data = []
    input_size = cfg['img_res']
    aspect_ratio = 1 if isinstance(input_size,
                                   int) else input_size[0] / input_size[1]

    for i, bbox in enumerate(bboxes_xywh):
        center, scale = box2cs(bbox, aspect_ratio, bbox_scale_factor=1.25)
        # prepare data
        data = {
            'image_path': img_or_path,
            'center': center,
            'scale': scale,
            'rotation': 0,
            'bbox_score': bbox[4] if len(bbox) == 5 else 1,
            'sample_idx': i,
        }
        data = inference_pipeline(data)
        batch_data.append(data)

    batch_data = collate(batch_data, samples_per_gpu=1)

    if next(model.parameters()).is_cuda:
        # scatter not work so just move image to cuda device
        batch_data['img'] = batch_data['img'].to(device)

    # get all img_metas of each bounding box
    batch_data['img_metas'] = [
        img_metas[0] for img_metas in batch_data['img_metas'].data
    ]

    # forward the model
    with torch.no_grad():
        results = model(
            img=batch_data['img'],
            img_metas=batch_data['img_metas'],
            sample_idx=batch_data['sample_idx'],
        )

    for idx in range(len(det_results)):
        mesh_result = det_results[idx].copy()
        mesh_result['bbox'] = bboxes_xyxy[idx]
        mesh_result['camera'] = results['camera'][idx]
        mesh_result['smpl_pose'] = results['smpl_pose'][idx]
        mesh_result['smpl_beta'] = results['smpl_beta'][idx]
        mesh_result['vertices'] = results['vertices'][idx]
        mesh_result['keypoints_3d'] = results['keypoints_3d'][idx]
        mesh_results.append(mesh_result)
    return mesh_results


def inference_video_based_model(model,
                                extracted_results,
                                with_track_id=True,
                                causal=True):
    """Inference SMPL parameters from extracted featutres using a video-based
    model.

    Args:
        model (nn.Module): The loaded mesh estimation model.
        extracted_results (List[List[Dict]]): Multi-frame feature extraction
            results stored in a nested list. Each element of the outer list
            is the feature extraction results of a single frame, and each
            element of the inner list is the feature information of one person,
            which contains:
                features (ndarray): extracted features
                track_id (int): unique id of each person, required when
                    ``with_track_id==True```
                bbox ((4, ) or (5, )): left, right, top, bottom, [score]
        with_track_id: If True, the element in extracted_results is expected to
            contain "track_id", which will be used to gather the feature
            sequence of a person from multiple frames. Otherwise, the extracted
            results in each frame are expected to have a consistent number and
            order of identities. Default is True.
        causal (bool): If True, the target frame is the first frame in
            a sequence. Otherwise, the target frame is in the middle of a
            sequence.

    Returns:
        list[dict]: Each item in the list is a dictionary, which contains:
            SMPL parameters, vertices, kp3d, and camera.
    """
    cfg = model.cfg
    device = next(model.parameters()).device
    seq_len = cfg.data.test.seq_len
    mesh_results = []
    # build the data pipeline
    inference_pipeline = Compose(cfg.inference_pipeline)
    target_idx = 0 if causal else len(extracted_results) // 2

    input_features = _gather_input_features(extracted_results)
    feature_sequences = _collate_feature_sequence(input_features,
                                                  with_track_id, target_idx)
    if not feature_sequences:
        return mesh_results

    batch_data = []

    for i, seq in enumerate(feature_sequences):

        data = {
            'features': seq['features'],
            'sample_idx': i,
        }

        data = inference_pipeline(data)
        batch_data.append(data)

    batch_data = collate(batch_data, samples_per_gpu=len(batch_data))

    if next(model.parameters()).is_cuda:
        # scatter not work so just move image to cuda device
        batch_data['features'] = batch_data['features'].to(device)

    with torch.no_grad():
        results = model(features=batch_data['features'],
                        img_metas=batch_data['img_metas'],
                        sample_idx=batch_data['sample_idx'])

    results['camera'] = results['camera'].reshape(-1, seq_len, 3)
    results['smpl_pose'] = results['smpl_pose'].reshape(-1, seq_len, 24, 3, 3)
    results['smpl_beta'] = results['smpl_beta'].reshape(-1, seq_len, 10)
    results['vertices'] = results['vertices'].reshape(-1, seq_len, 6890, 3)
    results['keypoints_3d'] = results['keypoints_3d'].reshape(
        -1, seq_len, 17, 3)

    for idx in range(len(feature_sequences)):
        mesh_result = dict()
        mesh_result['camera'] = results['camera'][idx, target_idx]
        mesh_result['smpl_pose'] = results['smpl_pose'][idx, target_idx]
        mesh_result['smpl_beta'] = results['smpl_beta'][idx, target_idx]
        mesh_result['vertices'] = results['vertices'][idx, target_idx]
        mesh_result['keypoints_3d'] = results['keypoints_3d'][idx, target_idx]
        mesh_result['bbox'] = extracted_results[target_idx][idx]['bbox']
        # 'track_id' is not included in results generated by mmdet
        if 'track_id' in extracted_results[target_idx][idx].keys():
            mesh_result['track_id'] = extracted_results[target_idx][idx][
                'track_id']
        mesh_results.append(mesh_result)
    return mesh_results


def feature_extract(
    model,
    img_or_path,
    det_results,
    bbox_thr=None,
    format='xywh',
):
    """Extract image features with a list of person bounding boxes.

    Args:
        model (nn.Module): The loaded feature extraction model.
        img_or_path (Union[str, np.ndarray]): Image filename or loaded image.
        det_results (List(dict)): the item in the dict may contain
            'bbox' and/or 'track_id'.
            'bbox' (4, ) or (5, ): The person bounding box, which contains
            4 box coordinates (and score).
            'track_id' (int): The unique id for each human instance.
        bbox_thr (float, optional): Threshold for bounding boxes.
            If bbox_thr is None, ignore it. Defaults to None.
        format (str, optional): bbox format. Default: 'xywh'.
            'xyxy' means (left, top, right, bottom),
            'xywh' means (left, top, width, height).

    Returns:
        list[dict]: The bbox & pose info,
            containing the bbox: (left, top, right, bottom, [score])
            and the features.
    """
    # only two kinds of bbox format is supported.
    assert format in ['xyxy', 'xywh']

    cfg = model.cfg
    device = next(model.parameters()).device

    feature_results = []
    if len(det_results) == 0:
        return feature_results

    # Change for-loop preprocess each bbox to preprocess all bboxes at once.
    bboxes = np.array([box['bbox'] for box in det_results])
    assert len(bboxes[0]) in [4, 5]

    # Select bboxes by score threshold
    if bbox_thr is not None:
        assert bboxes.shape[1] == 5
        valid_idx = np.where(bboxes[:, 4] > bbox_thr)[0]
        bboxes = bboxes[valid_idx]
        det_results = [det_results[i] for i in valid_idx]

    # if bbox_thr remove all bounding box
    if len(bboxes) == 0:
        return feature_results

    if format == 'xyxy':
        bboxes_xyxy = bboxes
        bboxes_xywh = xyxy2xywh(bboxes)
    else:
        # format is already 'xywh'
        bboxes_xywh = bboxes
        bboxes_xyxy = xywh2xyxy(bboxes)

    # build the data pipeline
    extractor_pipeline = [LoadImage()] + cfg.extractor_pipeline
    extractor_pipeline = Compose(extractor_pipeline)
    batch_data = []
    input_size = cfg['img_res']
    aspect_ratio = 1 if isinstance(input_size,
                                   int) else input_size[0] / input_size[1]

    for i, bbox in enumerate(bboxes_xywh):
        center, scale = box2cs(bbox, aspect_ratio, bbox_scale_factor=1.25)
        # prepare data
        data = {
            'image_path': img_or_path,
            'center': center,
            'scale': scale,
            'rotation': 0,
            'bbox_score': bbox[4] if len(bbox) == 5 else 1,
            'sample_idx': i,
        }
        data = extractor_pipeline(data)
        batch_data.append(data)

    batch_data = collate(batch_data, samples_per_gpu=1)

    if next(model.parameters()).is_cuda:
        # scatter not work so just move image to cuda device
        batch_data['img'] = batch_data['img'].to(device)

    # get all img_metas of each bounding box
    batch_data['img_metas'] = [
        img_metas[0] for img_metas in batch_data['img_metas'].data
    ]

    # forward the model
    with torch.no_grad():
        results = model(batch_data['img'])

        if isinstance(results, list) or isinstance(results, tuple):
            results = results[-1].mean(dim=-1).mean(dim=-1)

    for idx in range(len(det_results)):
        feature_result = det_results[idx].copy()
        feature_result['bbox'] = bboxes_xyxy[idx]
        feature_result['features'] = results[idx].cpu().numpy()
        feature_results.append(feature_result)

    return feature_results


def _gather_input_features(extracted_results):
    """Gather input features.

    Args:
        extracted_results (List[List[Dict]]):
            Multi-frame feature extraction results

    Returns:
        List[List[dict]]: Multi-frame feature extraction results
            stored in a nested list. Each element of the outer list is the
            feature extraction results of a single frame, and each element of
            the inner list is the extracted results of one person,
            which contains:
                features (ndarray): extracted features
                track_id (int): unique id of each person, required when
                    ``with_track_id==True```
    """
    sequence_inputs = []
    for frame in extracted_results:
        frame_inputs = []
        for res in frame:
            inputs = dict()
            if 'features' in res:
                inputs['features'] = res['features']
            if 'track_id' in res:
                inputs['track_id'] = res['track_id']
            frame_inputs.append(inputs)
        sequence_inputs.append(frame_inputs)
    return sequence_inputs


def _collate_feature_sequence(extracted_features,
                              with_track_id=True,
                              target_frame=0):
    """Reorganize multi-frame feature extraction results into individual
    feature sequences.

    Args:
        extracted_features (List[List[Dict]]): Multi-frame feature extraction
            results stored in a nested list. Each element of the outer list
            is the feature extraction results of a single frame, and each
            element of the inner list is the extracted results of one person,
            which contains:
                features (ndarray): extracted features
                track_id (int): unique id of each person, required when
                    ``with_track_id==True```
        with_track_id (bool): If True, the element in pose_results is expected
            to contain "track_id", which will be used to gather the pose
            sequence of a person from multiple frames. Otherwise, the pose
            results in each frame are expected to have a consistent number and
            order of identities. Default is True.
        target_frame (int): The index of the target frame. Default: 0.
    """
    T = len(extracted_features)
    assert T > 0

    target_frame = (T + target_frame) % T  # convert negative index to positive

    N = len(
        extracted_features[target_frame])  # use identities in the target frame
    if N == 0:
        return []

    C = extracted_features[target_frame][0]['features'].shape[0]

    track_ids = None
    if with_track_id:
        track_ids = [
            res['track_id'] for res in extracted_features[target_frame]
        ]

    feature_sequences = []
    for idx in range(N):
        feature_seq = dict()
        # gather static information
        for k, v in extracted_features[target_frame][idx].items():
            if k != 'features':
                feature_seq[k] = v
        # gather keypoints
        if not with_track_id:
            feature_seq['features'] = np.stack(
                [frame[idx]['features'] for frame in extracted_features])
        else:
            features = np.zeros((T, C), dtype=np.float32)
            features[target_frame] = extracted_features[target_frame][idx][
                'features']
            # find the left most frame containing track_ids[idx]
            for frame_idx in range(target_frame - 1, -1, -1):
                contains_idx = False
                for res in extracted_features[frame_idx]:
                    if res['track_id'] == track_ids[idx]:
                        features[frame_idx] = res['features']
                        contains_idx = True
                        break
                if not contains_idx:
                    # replicate the left most frame
                    features[frame_idx] = features[frame_idx + 1]

            # find the right most frame containing track_idx[idx]
            for frame_idx in range(target_frame + 1, T):
                contains_idx = False
                for res in extracted_features[frame_idx]:
                    if res['track_id'] == track_ids[idx]:
                        features[frame_idx] = res['features']
                        contains_idx = True
                        break
                if not contains_idx:
                    # replicate the right most frame
                    features[frame_idx] = features[frame_idx - 1]
                    # break
            feature_seq['features'] = features
        feature_sequences.append(feature_seq)

    return feature_sequences