-
Notifications
You must be signed in to change notification settings - Fork 0
/
getlabel.py
66 lines (51 loc) · 1.51 KB
/
getlabel.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
import xlrd
import json
import random
from tqdm import tqdm
label_path = "data/LIVE-SJTU_AVQA/MOS.xlsx"
train_ratio = 0.8
wb = xlrd.open_workbook(label_path)
sheet1 = wb.sheet_by_name("allNames")
sheet2 = wb.sheet_by_name("refNames")
sheet3 = wb.sheet_by_name("disNames")
sheet4 = wb.sheet_by_name("MOS")
sheet5 = wb.sheet_by_name("frameNum")
sheet6 = wb.sheet_by_name("frameRate")
sample_num = sheet1.nrows
labels = []
for j in range(sample_num):
sample_info = {}
# sheet2
sample_info["ref_video"] = sheet2.cell_value(j, 0)
sample_info["ref_audio"] = sheet2.cell_value(j, 1)
# sheet3
sample_info["dis_video"] = sheet3.cell_value(j, 0)
sample_info["dis_audio"] = sheet3.cell_value(j, 1)
# sheet4
sample_info["MOS"] = sheet4.cell_value(j, 0)
# sheet1
sample_info["compress"] = sheet1.cell_value(j, 1)
sample_info["kbps"] = sheet1.cell_value(j, 2)
# sheet5
sample_info["frame_num"] = sheet5.cell_value(j, 0)
# sheet6
sample_info["frame_rate"] = sheet6.cell_value(j, 0)
labels.append(sample_info)
random.shuffle(labels)
train_num = int(len(labels) * train_ratio)
train_set = labels[: train_num]
test_set = labels[train_num: ]
train_dict = {}
test_dict = {}
for id, info in enumerate(train_set):
train_dict[id] = info
for id, info in enumerate(test_set):
test_dict[id] = info
train_dict = json.dumps(train_dict)
test_dict = json.dumps(test_dict)
f = open("train.json", "w")
f.write(train_dict)
f.close()
f = open("test.json", "w")
f.write(test_dict)
f.close()