Add cats to hierarchy

bddl/data_generation/get_hierarchy_full.py

@@ -10,16 +10,17 @@
 from nltk.corpus import wordnet as wn
 
 HIERARCHY_OUTPUT_FN = pathlib.Path(__file__).parents[1] / "generated_data" / "output_hierarchy.json"
+HIERARCHY_PROPERTIES_OUTPUT_FN = pathlib.Path(__file__).parents[1] / "generated_data" / "output_hierarchy_properties.json"
 CATEGORY_MAPPING_FN = pathlib.Path(__file__).parents[1] / "generated_data" / "category_mapping.csv"
+SYN_PROP_PARAM_FN = pathlib.Path(__file__).parents[1] / "generated_data" / "propagated_annots_params.json"
 
 
 def add_igibson_objects(node, synset_to_cat):
   '''
   Go through the hierarchy and add the words associated with the synsets as attributes.
   '''
-  # categories = 
   if node["name"] in synset_to_cat:
-    node["igibson_categories"] = sorted(synset_to_cat[node["name"]])
+    node["categories"] = sorted(synset_to_cat[node["name"]])
 
   if "children" in node:
     for child_node in node["children"]:
@@ -55,6 +56,13 @@ def generate_paths(paths, path, synset, syn_prop_dict):
       generate_paths(paths, path + [hypernym], hypernym, syn_prop_dict)
 
 
+def add_properties(node, syn_prop_param_dict):
+  node["abilities"] = syn_prop_param_dict[node["name"]]
+  if "children" in node: 
+    for child_node in node["children"]:
+      add_properties(child_node, syn_prop_param_dict)
+
+
 # API
 
 def get_hierarchy(syn_prop_dict): 
@@ -85,6 +93,15 @@ def get_hierarchy(syn_prop_dict):
   return hierarchy
 
 
+def create_get_save_hierarchy_with_properties(hierarchy):
+  with open(SYN_PROP_PARAM_FN) as f:
+    syn_prop_param_dict = json.load(f)
+  add_properties(hierarchy, syn_prop_param_dict)
+  with open(HIERARCHY_PROPERTIES_OUTPUT_FN, "w") as f:
+    json.dump(hierarchy, f, indent=2)
+  return hierarchy
+
+
 if __name__ == "__main__":
   pass
 

bddl/data_generation/run_everything.py

@@ -1,5 +1,5 @@
 import pathlib
-from bddl.data_generation.get_hierarchy_full import get_hierarchy
+from bddl.data_generation.get_hierarchy_full import get_hierarchy, create_get_save_hierarchy_with_properties
 from bddl.data_generation.get_syn_prop_annots_canonical import create_get_save_annots_canonical, create_get_save_properties_to_synsets, create_get_save_synsets_to_descriptors
 from bddl.data_generation.propagate_by_intersection import create_get_save_propagated_canonical
 from bddl.data_generation.process_prop_param_annots import create_get_save_propagated_annots_params
@@ -20,8 +20,6 @@
 # owned_models.to_csv("owned_models.csv", index=False)
 
 def main():
-    nltk.download('wordnet')
-
     # Get full hierarchy (it's created and saved on import)
     with open(SYN_PROP_DATA_FN, "r") as f:
         syn_prop_dict = {}
@@ -44,6 +42,9 @@ def main():
     # Add parameter info to syns-to-props
     create_get_save_propagated_annots_params(propagated_canonical, props_to_syns)
 
+    # Add prop-param info to hierarchy 
+    create_get_save_hierarchy_with_properties(hierarchy)
+
     # # Create and save activity-specific hierarchies (no getting because that will get complicated)
     # create_save_activity_specific_hierarchies()