cvat to coco fix

airo-ugent · Oct 23, 2024 · 9104221 · 9104221
1 parent f476625
commit 9104221
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/airo-dataset-tools/airo_dataset_tools/cvat_labeling/convert_cvat_to_coco.py b/airo-dataset-tools/airo_dataset_tools/cvat_labeling/convert_cvat_to_coco.py
@@ -120,7 +120,13 @@ def _validate_coco_categories_are_in_cvat(
     # gather the annotation from CVAT
     cvat_categories_dict = defaultdict(list)
 
-    for annotation_category in cvat_parsed.annotations.meta.get_job_or_task().labels.label:
+    labels = cvat_parsed.annotations.meta.get_job_or_task().labels.label
+
+    # Handle both single LabelItem and list of LabelItems
+    if isinstance(labels, LabelItem):
+        labels = [labels]
+
+    for annotation_category in labels:
         assert isinstance(annotation_category, LabelItem)
         category_str, annotation_name = annotation_category.name.split(".")
         cvat_categories_dict[category_str].append(annotation_name)

diff --git a/airo-dataset-tools/airo_dataset_tools/data_parsers/coco.py b/airo-dataset-tools/airo_dataset_tools/data_parsers/coco.py
@@ -35,7 +35,9 @@
 from pydantic import BaseModel, field_validator, model_validator
 
 # Used by CocoInfo and CocoImage
-Datetime = str  # COCO uses both YYYY-MM-DD and YYYY-MM-DD HH:MM:SS for datetime
+Datetime = Union[
+    str, int
+]  # COCO uses both YYYY-MM-DD and YYYY-MM-DD HH:MM:SS for datetime, CVAT places 0 for missing values
 Url = str
 
 # Used by CocoImage and CocoLicenses