def get_end_to_end_preprocess_fn()

in empose/data/transforms.py [0:0]


def get_end_to_end_preprocess_fn(config, smpl_model, randomize_if_configured=False):
    """Factory function to return the preprocessing function depending on the given configuration."""
    normalize_root = NormalizeRoot()
    fk = SMPLFK(smpl_model)

    if config.use_real_offsets:
        noise_level = config.offset_noise_level if randomize_if_configured else -1
        sample_markers = SampleMarkersWithOffsets(smpl_model, list(get_all_offset_files().values()),
                                                  noise_level=noise_level)
    else:
        raise ValueError("We expect to use the real offsets.")

    noise_fn = get_noise_fn(config, randomize_if_configured)

    def _preprocess_fn(sample, mode='all', **noise_kwargs):
        if mode == 'all':
            sample = sample_markers(fk(normalize_root(sample)))
            return noise_fn(sample)
        elif mode == 'normalize_only':
            return normalize_root(sample)
        elif mode == 'after_normalize':
            return noise_fn(sample_markers(fk(sample)))
        else:
            raise ValueError("Mode '{}' unknown.".format(mode))

    return _preprocess_fn