-
Notifications
You must be signed in to change notification settings - Fork 30
/
mrr.py
89 lines (77 loc) · 2.67 KB
/
mrr.py
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
import argparse
import csv
import sys
import os
from typing import List
def read_meta_data(csv_path):
data_info = []
with open(csv_path, "r") as csv_file:
csv_reader = csv.reader(csv_file, delimiter=',')
line_count = 0
for row in csv_reader:
if line_count > 0:
music_id = row[0]
song_path = row[1]
hum_path = row[2]
data_info.append([music_id, song_path, hum_path])
line_count += 1
print(f'Processed {line_count} lines.')
return data_info
def create_meta_dict(train_meta):
sound2id = {}
id2hum = {}
id2song = {}
id2stt = {}
stt = 0
for row in train_meta:
music_id = row[0]
song_name = row[1].split("/")[-1].split(".")[0]
hum_name = row[2].split("/")[-1].split(".")[0]
sound2id[song_name] = music_id
sound2id[hum_name] = music_id
if music_id not in id2hum.keys():
id2hum[music_id] = [hum_name]
id2song[music_id] = [song_name]
else:
id2hum[music_id].append(hum_name)
id2song[music_id].append(song_name)
if music_id not in id2stt.keys():
id2stt[music_id] = stt
stt += 1
return sound2id, id2hum, id2song, id2stt
def mean_reciprocal_rank(preds: List[str], gt: str, k: int=10):
preds = preds[: min(k, len(preds))]
score = 0
for rank, pred in enumerate(preds):
if pred == gt:
score = 1 / (rank + 1)
break
return score
def main(csv_pred, csv_gt):
train_meta = read_meta_data(csv_gt)
sound2id, _, _, _ = create_meta_dict(train_meta)
acc = 0
num = 0
with open(csv_pred, "r") as csv_file:
csv_reader = csv.reader(csv_file, delimiter=',')
line_count = 0
for row in csv_reader:
if line_count > 0:
hum_id = row[0].split(".")[0]
preds = []
for col in row[1:]:
preds.append(sound2id[str(col)])
print(hum_id, mean_reciprocal_rank(preds, sound2id[str(hum_id)]))
acc += mean_reciprocal_rank(preds, sound2id[str(hum_id)])
num += 1
line_count += 1
print(f'Processed {line_count} lines.')
return acc / num
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--csv_pred", type=str, required=True, help="path to predict csv")
parser.add_argument("--csv_gt", type=str, required=True, help="path to ground-truth csv")
args = parser.parse_args()
mrr = main(args.csv_pred, args.csv_gt)
print("-----------------------------")
print(f"MRR: {mrr}")