kalpkanungo/SceneGraphNet

1import json2 3LABELS_PATH = "data/relationship_dataset/labels.json"4OUTPUT_PATH = "data/relationship_dataset/labels_encoded.json"5MAP_PATH = "data/relationship_dataset/label_map.json"6 7with open(LABELS_PATH) as f:8    labels = json.load(f)9 10predicates = sorted(list(set([item["predicate"] for item in labels])))11 12label_map = {p: i for i, p in enumerate(predicates)}13 14encoded = []15 16for item in labels:17    encoded.append({18        "image": item["image"],19        "label": label_map[item["predicate"]]20    })21 22with open(OUTPUT_PATH, "w") as f:23    json.dump(encoded, f)24 25with open(MAP_PATH, "w") as f:26    json.dump(label_map, f)27 28print("Classes:", len(label_map))