-
Notifications
You must be signed in to change notification settings - Fork 12
/
get_counts.py
71 lines (60 loc) · 2.02 KB
/
get_counts.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
from utils.vg_utils import entity_to_aliases, predicate_to_aliases
import json
# file to read
SCENE_GRAPH_FILE = "data/raw_data/scene_graphs.json"
ATTRIBUTES_FILE = "data/raw_data/attributes.json"
entity_counts = {}
predicate_counts = {}
attribute_counts = {}
def get_entity_predicate_counts(scene_graph_data):
for image in scene_graph_data:
for entity in image["objects"]:
entity_name = entity["name"] if "name" in entity else entity["names"]
entity_name = entity_to_aliases(entity_name)
if str(entity_name) in entity_counts:
entity_counts[str(entity_name)] += 1
else:
entity_counts[str(entity_name)] = 1
for predicate in image["relationships"]:
predicate_name = predicate["predicate"]
predicate_name = predicate_to_aliases(predicate_name)
if str(predicate_name) in predicate_counts:
predicate_counts[str(predicate_name)] += 1
else:
predicate_counts[str(predicate_name)] = 1
def get_attribute_counts(attributes_data):
for image in attributes_data:
for entity in image["attributes"]:
if "attributes" in entity:
for attribute in entity["attributes"]:
if attribute in attribute_counts:
attribute_counts[attribute] += 1
else:
attribute_counts[attribute] = 1
def main():
print("Loading scene graph data...")
f = open(SCENE_GRAPH_FILE, "rb")
scene_graph_data = json.load(f)
f.close()
print("Done!")
get_entity_predicate_counts(scene_graph_data)
print("Saving entity counts...")
with open("entity_counts.json", "w") as outfile:
json.dump(entity_counts, outfile)
print("Done!")
print("Saving predicate counts...")
with open("predicate_counts.json", "w") as outfile:
json.dump(predicate_counts, outfile)
print("Done!")
print("Loading attributes file...")
f = open(ATTRIBUTES_FILE, "rb")
attributes_data = json.load(f)
f.close()
print("Done!")
get_attribute_counts(attributes_data)
print("Saving attribute counts...")
with open("attribute_counts.json", "w") as outfile:
json.dump(attribute_counts, outfile)
print("Done!")
if __name__ == "__main__":
main()