DeepWDM/post_process.py at master · adiyoss/DeepWDM · 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
137
138
import argparse
import os


def predict_duration_target(y):
    """
    This function get all the voice activities that were detected in the signal
    :param y: the network predictions
    :return: list of the voice activities
    """
    predictions = []
    onset_found = False
    onset = 0
    for i in range(1, len(y)):
        prev = y[i - 1]
        curr = y[i]
        if prev is 1 and curr is 2:
            onset_found = True
            onset = i
        if prev is 2 and curr is 1 and onset_found:
            onset_found = False
            predictions.append([onset, i - 1])
    return predictions


def max_duration(predictions):
    max_value = 1
    max_idx = -1
    for i, p in enumerate(predictions):
        tmp = p[1] - p[0]
        if tmp > max_value:
            max_value = tmp
            max_idx = i
    if max_idx != -1:
        return predictions[max_idx]
    else:
        print('No predictions')
        return [0, 0]


def smooth_duration(predictions):
    new_seq = list()
    onset = 0
    offset = 0
    for i, p in enumerate(predictions):
        if onset == 0 and offset == 0:
            onset = p[0]
            offset = p[1]
        else:
            if p[0] - offset < 10:
                offset = p[1]
            else:
                new_seq.append([onset, offset])
                onset = p[0]
                offset = p[1]
    new_seq.append([onset, offset])
    return max_duration(new_seq)

    # for icassp predictions
    # max_length = 0
    # onset = 0
    # offset = 0
    # for p in predictions:
    #     tmp_dur = p[1] - p[0]
    #     if tmp_dur > max_length:
    #         onset = p[0]
    #         offset = p[1]
    # return [onset, offset]

    # max_idx = -1
    # max_length = 0
    # onset = 0
    # offset = 0
    # final_b = 0
    # final_r = 0
    # for i, p in enumerate(predictions):
    #     if onset == 0 and offset == 0:
    #         onset = p[0]
    #         offset = p[1]
    #     else:
    #         if p[0] - offset < 10:
    #             offset = p[1]
    #         else:
    #             tmp_len = offset - onset
    #             if offset - onset > max_length:
    #                 max_idx = i
    #                 final_b = onset
    #                 final_r = offset
    #                 max_length = tmp_len
    #             onset = 0
    #             offset = 0
    # if max_idx == -1 and (offset != 0 or onset != 0):
    #     max_idx = 1
    #     final_b = onset
    #     final_r = offset

    # if max_idx != -1:
    #     return [final_b, final_r]
    # else:
    #     print('No predictions')
    #     return [0, 0]


def post_process(filename, output_path):
    """
    Computes the duration from the network predictions
    :param filename: the predictions file path
    :param output_path: the path to write the final duration
    """
    x_file = os.path.abspath(filename)
    output_path = os.path.abspath(output_path)

    classifications = list()
    # parsing the predictions file
    with open(x_file) as f:
        lines = f.readlines()
        for line in lines:
            classifications.append(int(line[:-1]))
    f.close()

    predictions = predict_duration_target(classifications)
    # prediction = max_duration(predictions)
    prediction = smooth_duration(predictions)
    with open(output_path, 'w') as fid:
        fid.write(str(prediction[0]) + ' ' + str(prediction[1]) + '\n')
    fid.close()


if __name__ == "__main__":
    # -------------MENU-------------- #
    # command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("filename", help="The path to the prediction file that were generated from the nn")
    parser.add_argument("output_path", help="The path to save the label")
    args = parser.parse_args()

    # main function
    post_process(args.filename, args.output_path)