StrengthNet/test.py at main · codening99/StrengthNet · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
#!/usr/bin/env python3
# -*- coding: utf-8 -*-


import os
import time
import numpy as np
from tqdm import tqdm
import argparse
import fnmatch
import pandas as pd

from statistics import mean

import tensorflow as tf
from tensorflow import keras
from model import CNN_BLSTM

import utils
import random


def find_files(root_dir, query="*.wav", include_root_dir=True):
    """Find files recursively.

    Args:
        root_dir (str): Root root_dir to find.
        query (str): Query to find.
        include_root_dir (bool): If False, root_dir name is not included.

    Returns:
        list: List of found filenames.

    """
    files = []
    for root, dirnames, filenames in os.walk(root_dir, followlinks=True):
        for filename in fnmatch.filter(filenames, query):
            files.append(os.path.join(root, filename))
    if not include_root_dir:
        files = [file_.replace(root_dir + "/", "") for file_ in files]

    return files


def main():
    parser = argparse.ArgumentParser(
        description="Evaluate custom waveform files using pretrained MOSnet.")
    parser.add_argument("--rootdir", default=None, type=str,
                        help="rootdir of the waveforms to be evaluated")
    parser.add_argument("--pretrained_model", default="./output/strengthnet.h5", type=str,
                        help="pretrained model file")
    args = parser.parse_args()

    #### tensorflow & gpu settings ####

    # 0 = all messages are logged (default behavior)
    # 1 = INFO messages are not printed
    # 2 = INFO and WARNING messages are not printed
    # 3 = INFO, WARNING, and ERROR messages are not printed
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'  # or any {'0', '1', '2'}

    tf.debugging.set_log_device_placement(False)
    # set memory growth
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        try:
            # Currently, memory growth needs to be the same across GPUs
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)

            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
            print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
        except RuntimeError as e:
            # Memory growth must be set before GPUs have been initialized
            print(e)

    ###################################

    # find waveform files
    wavfiles = sorted(find_files(args.rootdir, "*.wav"))

    # init model
    print("Loading model weights")
    StrengthNet = CNN_BLSTM()
    model = StrengthNet.build()
    model.load_weights(args.pretrained_model)

    # evaluation
    print("Start evaluating {} waveforms...".format(len(wavfiles)))
    # results = []

    results_frame = []

    for wavfile in tqdm(wavfiles):

        # spectrogram
        mel_sgram = utils.get_melspectrograms(wavfile)
        timestep = mel_sgram.shape[0]
        mel_sgram = np.reshape(mel_sgram, (1, timestep, utils.n_mels))
        # make prediction
        Strength_score, Frame_score, emo_class = model.predict(mel_sgram, verbose=0, batch_size=1)

        # print(Frame_score.shape)

        # write to list
        # result = wavfile + " {:.3f}".format(Strength_score[0][0])
        # results.append(result)

        # save frame_score
        # wavfile = wavfile.replace('../../dataset/Emotional Speech Dataset (ESD)/', '')
        result_frame = ""
        for item in range(Frame_score.shape[1]):
            result_frame += " {:.3f}".format(Frame_score[0][item][0])
        results_frame.append(result_frame)
        # print(results_frame)

    # Dump data
    data_frame = pd.DataFrame(
        data={'file_path': [wf.replace('../../dataset/ESDcorpus/', '') for wf in wavfiles],
              'score': [fc for fc in results_frame]})

    data_frame.to_csv(os.path.join(args.rootdir, 'frame_score.csv'))

    # write final raw result
    # resultrawpath = os.path.join(args.rootdir, "StrengthNet_result_raw.csv")
    # with open(resultrawpath, "w") as outfile:
    #     outfile.write("\n".join(sorted(results_frame)))

    print('All done, Save at', args.rootdir, 'exit.')


if __name__ == '__main__':
    main()