Add PDF handling to initialise_labels

encord-team · Oct 8, 2024 · 9995247 · 9995247
1 parent 9633b9e
commit 9995247
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/encord/constants/enums.py b/encord/constants/enums.py
@@ -24,6 +24,7 @@ class DataType(StringEnum):
     IMAGE = "image"
     DICOM_STUDY = "dicom_study"
     NIFTI = "nifti"
+    PDF = "pdf"
 
     # will be displayed if the Encord platform has a new data type that is not present in this SDK version. Please upgrade your SDK version
     MISSING_DATA_TYPE = "_MISSING_DATA_TYPE_"

diff --git a/encord/objects/ontology_labels_impl.py b/encord/objects/ontology_labels_impl.py
@@ -1752,7 +1752,7 @@ def _parse_label_row_dict(self, label_row_dict: dict) -> LabelRowReadOnlyData:
         frame_to_image_hash = {item.frame_number: item.image_hash for item in frame_level_data.values()}
         data_type = DataType(label_row_dict["data_type"])
 
-        if data_type == DataType.VIDEO or data_type == DataType.IMAGE:
+        if data_type == DataType.VIDEO or data_type == DataType.IMAGE or data_type == DataType.PDF:
             data_dict = list(label_row_dict["data_units"].values())[0]
             data_link = data_dict["data_link"]
             # Dimensions should be always there
@@ -1819,7 +1819,7 @@ def _parse_labels_from_dict(self, label_row_dict: dict):
         for data_unit in label_row_dict["data_units"].values():
             data_type = DataType(label_row_dict["data_type"])
 
-            if data_type == DataType.IMG_GROUP or data_type == DataType.IMAGE:
+            if data_type == DataType.IMG_GROUP or data_type == DataType.IMAGE or data_type == DataType.PDF:
                 frame = int(data_unit["data_sequence"])
                 self._add_object_instances_from_objects(data_unit["labels"].get("objects", []), frame)
                 self._add_classification_instances_from_classifications(