eval.py

# -*- coding: utf-8 -*-
"""
Evaluate GRU model

@author: Denis Tome'

Copyright Epic Games, Inc. All Rights Reserved.

"""

from base import ArgParser
from torch.utils.data import DataLoader
from logger.console_logger import ConsoleLogger
from base import OutputData
from utils.io import process_args_modes
from Speech2Face.model import GRU
from Speech2Face.dataset_def import Dataset, ToTensor
from Speech2Face.model.modules import PoseError
from Speech2Face.test_model import EvalGRU


LOGGER = ConsoleLogger('Main')
PARSER = ArgParser('GRUEvaluation')


def main(args):
    """Main"""

    if args.verbose:
        LOGGER.info('Initializing model')

    # ------------------- model -------------------

    model = GRU(args.model_config_file,
                batch_size=args.batch_size)

    if args.verbose:
        model.summary()

    # ------------------- dataset -------------------

    dataset = Dataset(args.dataset_file,
                      transformation=ToTensor(),
                      data_selection=OutputData.FEAT | OutputData.POS,
                      window_size=args.window_size)

    data_loader = DataLoader(dataset,
                             batch_size=args.batch_size,
                             shuffle=False)

    # ------------------- metric -------------------

    metrics = [PoseError()]

    # ------------------------------------------------------
    # ------------------- evaluate model -------------------
    # ------------------------------------------------------

    tester = EvalGRU(
        model,
        metrics,
        test_loader=data_loader,
        batch_size=args.batch_size,
        output=args.output,
        cuda=args.cuda,
        resume=args.resume,
        name=args.name,
        verbose=not args.no_verbose)

    tester.test()
    LOGGER.info('Done.')


if __name__ == '__main__':
    main(PARSER.get_arguments())