Skip to content

Commit

Permalink
Merge pull request #136 from opendatalab/feat/labelme
Browse files Browse the repository at this point in the history 
fix: update attributes in xml and pascal_voc
  • Loading branch information
@gary-Shen
gary-Shen authored Nov 11, 2024
2 parents 7a4a4e1 + 0affd60 commit 1cf0def
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 42 deletions.
22 changes: 10 additions & 12 deletions labelu/internal/common/converter.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import csv
import json
import os
from xml.dom import minidom
from zipfile import ZipFile
from PIL import Image, ImageDraw
from enum import Enum
Expand Down Expand Up @@ -752,43 +753,40 @@ def convert_to_xml(self, config: dict, input_data: List[dict], out_data_file_nam

# result struct
xml_converter = XML_converter()
root = xml_converter.create_root("root")
sample_item = ET.Element("sample")
root = ET.Element("root")
sample_item = ET.SubElement(root, "sample")

for sample in input_data:
data = json.loads(sample.get("data"))
file = sample.get("file", {})

# skip invalid data
annotated_result = json.loads(data.get("result"))
result = ET.Element("result")
result = ET.SubElement(sample_item, "result")
if annotated_result and sample.get("state") == "SKIPPED":
ET.SubElement(result, "valid").text = "False"

ET.SubElement(sample_item, "id").text = str(sample.get("id"))
ET.SubElement(sample_item, "fileName").text = file.get("filename", "")[9:]
ET.SubElement(sample_item, "url").text = file.get("url")

ET.SubElement(sample_item, "folder").text = str(settings.MEDIA_ROOT)
ET.SubElement(sample_item, "path").text = file.get("path")
ET.SubElement(sample_item, "fileName").text = file.get("filename", "")

ET.SubElement(result, "width").text = str(annotated_result.get("width", 0))
ET.SubElement(result, "height").text = str(annotated_result.get("height", 0))
ET.SubElement(result, "rotate").text = str(annotated_result.get("rotate", 0))

# change result struct
if annotated_result:
annotations = ET.Element("annotations")
annotations = ET.SubElement(result, "annotations")
for tool in annotated_result.copy().keys():
tool_results = annotated_result.pop(tool)
if tool.endswith("Tool"):
tool_annotations = xml_converter.convert_tool_results(tool, tool_results)

for annotation in tool_annotations:
annotations.append(annotation)

result.append(annotations)

sample_item.append(result)

root.append(sample_item)

tree = ET.ElementTree(root)
tree.write(file_full_path, encoding="utf-8", xml_declaration=True)
logger.info("Export file path: {}", file_full_path)
Expand Down
49 changes: 19 additions & 30 deletions labelu/internal/common/xml_converter.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,6 @@
from labelu.internal.common.config import settings

class XML_converter:
def create_root(self, root_name: str):
root = ET.Element(root_name)

return root

def create_pascal_voc_xml(self, config: dict, file: dict, sample_result: dict):
annotation = ET.Element("annotation")

Expand Down Expand Up @@ -58,9 +53,11 @@ def get_label(_tool: str, _input_label: str):
name.text = get_label(tool, tool_result.get("label", ""))
# get value from attributes
truncated = ET.SubElement(obj_elem, "truncated")
truncated.text = str(tool_result.get("attributes", {}).get("truncated", 0))
truncated_value = tool_result.get("attributes", {}).get("truncated", [0])
truncated.text = str(truncated_value[0])
difficult = ET.SubElement(obj_elem, "difficult")
difficult.text = str(tool_result.get("attributes", {}).get("difficult", 0))
difficult_value = tool_result.get("attributes", {}).get("difficult", [0])
difficult.text = str(difficult_value[0])

bndbox = ET.SubElement(obj_elem, "bndbox")
xmin = ET.SubElement(bndbox, "xmin")
Expand Down Expand Up @@ -103,15 +100,14 @@ def create_attributes(self, attributes: dict):
result = ET.Element("attributes")

for key, value in attributes.items():
attribute = ET.Element("attribute")
attribute = ET.SubElement(result, "attribute")
ET.SubElement(attribute, "key").text = key
ET.SubElement(attribute, "value").text = value
result.append(attribute)
ET.SubElement(attribute, "value").text = value if isinstance(value, str) else ", ".join(value)

return result

def create_rect(self, anno: dict):
annotation = ET.Element("rect")
annotation = ET.Element("object")

ET.SubElement(annotation, "toolName").text = "rectTool"

Expand All @@ -131,7 +127,7 @@ def create_rect(self, anno: dict):
return annotation

def create_polygon(self, anno: dict):
annotation = ET.Element("polygon")
annotation = ET.Element("object")

ET.SubElement(annotation, "toolName").text = "polygonTool"

Expand Down Expand Up @@ -170,7 +166,7 @@ def create_polygon(self, anno: dict):
return annotation

def create_point(self, anno: dict):
annotation = ET.Element("point")
annotation = ET.Element("object")

ET.SubElement(annotation, "toolName").text = "pointTool"

Expand All @@ -187,7 +183,7 @@ def create_point(self, anno: dict):
return annotation

def create_line(self, anno: dict):
annotation = ET.Element("line")
annotation = ET.Element("object")

ET.SubElement(annotation, "toolName").text = "lineTool"

Expand Down Expand Up @@ -226,7 +222,7 @@ def create_line(self, anno: dict):
return annotation

def create_cuboid(self, anno: dict):
annotation = ET.Element("cuboid")
annotation = ET.Element("object")

ET.SubElement(annotation, "toolName").text = "cuboidTool"

Expand Down Expand Up @@ -260,42 +256,35 @@ def create_cuboid(self, anno: dict):
return annotation

def create_tag(self, anno: dict):
annotation = ET.Element("tag")

annotation = ET.Element("object")
result = ET.SubElement(annotation, "result")
ET.SubElement(annotation, "toolName").text = "tagTool"

result = ET.Element("result")

ET.SubElement(result, "order").text = str(anno.get("order", 0))

ET.SubElement(result, "id").text = anno.get("id", "")

values = ET.Element("values")
values = ET.SubElement(result, "values")

for key, value in anno.get("value", {}).items():
ET.SubElement(values, key).text = ", ".join(value) if value else ""


result.append(values)
annotation.append(result)

return annotation

def create_text(self, anno: dict):
annotation = ET.Element("text")
annotation = ET.Element("object")

result = ET.SubElement(annotation, "result")
ET.SubElement(annotation, "toolName").text = "textTool"

result = ET.Element("result")

ET.SubElement(result, "order").text = str(anno.get("order", 0))
ET.SubElement(result, "id").text = anno.get("id", "")

values = ET.Element("values")
values = ET.SubElement(result, "values")

for key, value in anno.get("value", {}).items():
ET.SubElement(values, key).text = value

result.append(values)
annotation.append(result)

return annotation

0 comments on commit 1cf0def

Please sign in to comment.