-
Notifications
You must be signed in to change notification settings - Fork 0
/
label_replacement.py
117 lines (86 loc) · 3.13 KB
/
label_replacement.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
#!/usr/bin/env python3
"""
Author : Travis Simmons
Date : 2022-03-14
Purpose: Replace labels in labelbox export
"""
import argparse
import os
import sys
import json
# --------------------------------------------------
def get_args():
"""Get command-line arguments"""
parser = argparse.ArgumentParser(
description='Rock the Casbah',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
# parser.add_argument('positional',
# metavar='str',
# help='A positional argument')
parser.add_argument('-i',
'--input_json',
help='JSON exported from labelbox labels',
metavar='str',
type=str,
required=True)
parser.add_argument('-ol',
'--original_labels',
help='A list of labels you want to replace',
nargs='+',
required=True)
parser.add_argument('-rl',
'--replacement_labels',
help='replacement labels',
nargs='+',
required=True)
parser.add_argument('-of',
'--output_filename',
help='Output file name.',
metavar='str',
type=str,
default='label_replaced.json')
return parser.parse_args()
# --------------------------------------------------
def main():
"""Make a jazz noise here"""
args = get_args()
input_filename = args.input_json
ol_list = args.original_labels
rl_list = args.replacement_labels
output_filename = args.output_filename
# create replacement dict
r_dict = {}
# if only one replacement label was provided,
# replace all labels with it
if len(rl_list) == 0:
replacement_label = rl_list[0]
for i in ol_list:
r_dict[i] = replacement_label
# if more than one label is given, then do a 1:1 replacement
else:
for index, row in enumerate(ol_list):
r_dict[row] = rl_list[index]
with open(input_filename) as f:
data = json.load(f)
for index, i in enumerate(data):
try:
# print(data[index].keys())
labels = data[index]['Label']['objects']
for lab_num, i in enumerate(labels):
# lookup old label, if it is not one you want to replace, then skip it
old_label = data[index]['Label']['objects'][lab_num]['title']
try:
new_label = r_dict[old_label]
# repalce the label
data[index]['Label']['objects'][lab_num]['title'] = new_label
data[index]['Label']['objects'][lab_num]['value'] = new_label
except:
pass
except:
print('Finished.')
break
with open(output_filename, 'w') as f:
json.dump(data, f)
# --------------------------------------------------
if __name__ == '__main__':
main()