fix: fix category name in pascal_voc and tfrecord format files

opendatalab · Nov 5, 2024 · b089074 · b089074
1 parent 4344f99
commit b089074
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/labelu/internal/common/tf_record_converter.py b/labelu/internal/common/tf_record_converter.py
@@ -16,6 +16,10 @@ def _get_label_id(self, label: str):
     def create_tf_examples(self, sample_results: List[dict], config: dict):
         examples = []
         label_text_dict = {}
+
+        for tool in config.get("tools", []):
+            label_text_dict[tool.get("tool")] = { attr.get("value"): attr.get("key") for attr in tool.get("config", {}).get("attributes", [])}
+
         common_attributes = { attr.get("value"): attr.get("key") for attr in config.get("attributes", [])}
 
         def get_label(_tool: str, _input_label: str):

diff --git a/labelu/internal/common/xml_converter.py b/labelu/internal/common/xml_converter.py
@@ -31,6 +31,10 @@ def create_pascal_voc_xml(self, config: dict, file: dict, sample_result: dict):
         depth.text = "3"
 
         label_text_dict = {}
+
+        for tool in config.get("tools", []):
+            label_text_dict[tool.get("tool")] = { attr.get("value"): attr.get("key") for attr in tool.get("config", {}).get("attributes", [])}
+
         common_attributes = { attr.get("value"): attr.get("key") for attr in config.get("attributes", [])}
 
         def get_label(_tool: str, _input_label: str):