-
Notifications
You must be signed in to change notification settings - Fork 26
/
generate_mapping.py
121 lines (80 loc) · 3.6 KB
/
generate_mapping.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
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
import os, glob
import json
import argparse
def public_paths_labeled(root):
"""Map paths for public datasets as dictionary list"""
images_raw = sorted(glob.glob(os.path.join(root, "Public/images/*")))
labels_raw = sorted(glob.glob(os.path.join(root, "Public/labels/*")))
data_dicts = []
for image_path, label_path in zip(images_raw, labels_raw):
name1 = image_path.split("/")[-1].split(".")[0]
name2 = label_path.split("/")[-1].split("_label")[0]
assert name1 == name2
data_item = {
"img": image_path.split("MEDIAR/")[-1],
"label": label_path.split("MEDIAR/")[-1],
}
data_dicts.append(data_item)
map_dict = {"public": data_dicts}
return map_dict
def official_paths_labeled(root):
"""Map paths for official labeled datasets as dictionary list"""
image_path = os.path.join(root, "Official/Training/images/*")
label_path = os.path.join(root, "Official/Training/labels/*")
images_raw = sorted(glob.glob(image_path))
labels_raw = sorted(glob.glob(label_path))
data_dicts = []
for image_path, label_path in zip(images_raw, labels_raw):
name1 = image_path.split("/")[-1].split(".")[0]
name2 = label_path.split("/")[-1].split("_label")[0]
assert name1 == name2
data_item = {
"img": image_path.split("MEDIAR/")[-1],
"label": label_path.split("MEDIAR/")[-1],
}
data_dicts.append(data_item)
map_dict = {"official": data_dicts}
return map_dict
def official_paths_tuning(root):
"""Map paths for official tuning datasets as dictionary list"""
image_path = os.path.join(root, "Official/Tuning/images/*")
images_raw = sorted(glob.glob(image_path))
data_dicts = []
for image_path in images_raw:
data_item = {"img": image_path.split("MEDIAR/")[-1]}
data_dicts.append(data_item)
map_dict = {"official": data_dicts}
return map_dict
def add_mapping_to_json(json_file, map_dict):
"""Save mapped dictionary as a json file"""
if not os.path.exists(json_file):
with open(json_file, "w") as file:
json.dump({}, file)
with open(json_file, "r") as file:
data = json.load(file)
for map_key, map_item in map_dict.items():
if map_key not in data.keys():
data[map_key] = map_item
else:
print('>>> "{}" already exists in path map keys...'.format(map_key))
with open(json_file, "w") as file:
json.dump(data, file)
if __name__ == "__main__":
# [!Caution] The paths should be overrided for the local environment!
parser = argparse.ArgumentParser(description="Mapping files and paths")
parser.add_argument("--root", default=".", type=str)
args = parser.parse_args()
MAP_DIR = "./train_tools/data_utils/"
print("\n----------- Path Mapping for Labeled Data is Started... -----------\n")
map_labeled = os.path.join(MAP_DIR, "mapping_labeled.json")
map_dict = official_paths_labeled(args.root)
add_mapping_to_json(map_labeled, map_dict)
print("\n----------- Path Mapping for Tuning Data is Started... -----------\n")
map_labeled = os.path.join(MAP_DIR, "mapping_tuning.json")
map_dict = official_paths_tuning(args.root)
add_mapping_to_json(map_labeled, map_dict)
print("\n----------- Path Mapping for Public Data is Started... -----------\n")
map_public = os.path.join(MAP_DIR, "mapping_public.json")
map_dict = public_paths_labeled(args.root)
add_mapping_to_json(map_public, map_dict)
print("\n-------------- Path Mapping is Ended !!! ---------------------------\n")