TonSpeech/eval.py at main · tungedng2710/TonSpeech · 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
134
135
136
# Copyright (c) 2022 TonAI Research, MetaTon
# Author: Tung Ng
# Our GitHub: https://github.com/MetaTon-AI-Research

import argparse
import os
from src.eval_utils import load_sample, eval_pesq, eval_stoi
import pandas as pd
from tqdm import tqdm

def get_args():
    parser = argparse.ArgumentParser(description='Process some integers.')
    parser.add_argument('--eval_on_dataset', type=int, help='0: single sample; 1: dataset', default=0)
    parser.add_argument('--eval_voicebank', type=int, help='0: single sample; 1: voicebank', default=0)
    parser.add_argument('--down_sample', type=int, help='down sample rate into 16k Hz', default=1)
    parser.add_argument('--clean', type=str, help='path/to/clean/voice')
    parser.add_argument('--denoised', type=str, help='path/to/denoised/voice (single) or folder (dataset)')
    parser.add_argument('--metric', type=str, help='pesq or stoi', default='pesq')
    parser.add_argument('--trimmed_duration', type=int, default=-1)
    parser.add_argument('--to_csv', type=int, help='save the results to csv file', default=0)
    parser.add_argument('--verbose', type=int, help='show the progress', default=1)

    return parser.parse_args()

def eval_single_sample(args):
    trimmed_duration = args.trimmed_duration
    if args.down_sample == 1:
        down_sample = True
    else:
        down_sample = False

    clean, fs = load_sample(path=args.clean, down_sample=down_sample)
    denoised, fs = load_sample(path=args.denoised, down_sample=down_sample)

    min_length = min(len(clean), len(denoised))
    clean = clean[:min_length]
    denoised = denoised[:min_length]

    # Print the result to console
    print("Clean voice: ", args.clean)
    print("Denoised voice: ", args.denoised)
    if args.metric == 'pesq':
        print("Calculating PESQ score...")
        print("PESQ score: ", eval_pesq(fs, clean, denoised, trimmed_duration))
    elif args.metric == 'stoi':
        print("STOI score: ", eval_stoi(fs, clean, denoised))
    else:
        print("Metric should be pesq or stoi")

def eval_dataset(args):
    trimmed_duration = args.trimmed_duration
    if args.down_sample == 1:
        down_sample = True
    else:
        down_sample = False

    root_dir = args.denoised
    sample_files = sorted(os.listdir(root_dir))
    clean, fs = load_sample(path=args.clean, down_sample=down_sample)
    scores = []
    fnames = []
    for file in tqdm(sample_files):
        fnames.append(file)
        fpath = root_dir+'/'+file
        denoised, fs = load_sample(path=fpath, down_sample=down_sample)
        min_length = min(len(clean), len(denoised))
        clean = clean[:min_length]
        denoised = denoised[:min_length]
        if args.metric == 'pesq':
            score = eval_pesq(fs, clean, denoised, trimmed_duration)
            scores.append(score)
        elif args.metric == 'stoi':
            score = eval_stoi(fs, clean, denoised)
            scores.append(score)
        else:
            print("Metric should be pesq or stoi")
            break

    if args.verbose > 0:
        print("---------------------------------------------")
        for i in range(len(fnames)):
            print("{file}: PESQ: {score}".format(file=fnames[i], score=round(scores[i], 4)))
        print("Mean PESQ: {score}".format(score=round(sum(scores)/len(scores), 4)))

    if args.to_csv == 1:
        print("Saving result to csv file...")
        df_dict = {"file_name": fnames,
                   args.metric+"_score": scores}
        df = pd.DataFrame(df_dict)
        saved_name = args.metric+"_results.csv"
        df.to_csv(saved_name, index=False)

def eval_voicebank(args):
    trimmed_duration = args.trimmed_duration
    if args.down_sample == 1:
        down_sample = True
    else:
        down_sample = False
    scores = []
    fnames = []
    for file in tqdm(os.listdir(args.clean)):
        fnames.append(file)
        denoised_path = args.denoised+'/'+file
        clean_path = args.clean+'/'+file

        clean, fs = load_sample(path=clean_path, down_sample=down_sample)
        denoised, fs = load_sample(path=denoised_path, down_sample=down_sample)
        min_length = min(len(clean), len(denoised))
        clean = clean[:min_length]
        denoised = denoised[:min_length]
        if args.metric == 'pesq':
            score = eval_pesq(fs, clean, denoised, trimmed_duration)
            scores.append(score)
        elif args.metric == 'stoi':
            score = eval_stoi(fs, clean, denoised)
            scores.append(score)
        else:
            print("Metric should be pesq or stoi")
            break
    if args.verbose > 0:
        print("---------------------------------------------")
        for i in range(len(fnames)):
            print("{file}: PESQ: {score}".format(file=fnames[i], score=round(scores[i], 4)))
        print("Mean PESQ: {score}".format(score=round(sum(scores)/len(scores), 4)))

if __name__ == "__main__":
    print("Evaluating...")
    args = get_args()
    if args.eval_on_dataset == 0:
        eval_single_sample(args)
    elif args.eval_on_dataset == 1:
        eval_dataset(args)
    elif args.eval_on_dataset == 2:
        eval_voicebank(args)
    else:
        print("args.eval_on_dataset should be 0, 1 or 2")