Function inference_video_based_model

detrsmpl/apis/inference.py:204–292

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 oute

(model,
                                extracted_results,
                                with_track_id=True,
                                causal=True)

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))