Add handling for PDFs in data_unit and label retrieval

encord-team · Oct 9, 2024 · 8c81651 · 8c81651
1 parent 9995247
commit 8c81651
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/encord/objects/ontology_labels_impl.py b/encord/objects/ontology_labels_impl.py
@@ -1523,6 +1523,7 @@ def _to_encord_data_unit(self, frame_level_data: FrameLevelImageGroupData) -> Di
             or data_type == DataType.DICOM
             or data_type == DataType.IMAGE
             or data_type == DataType.NIFTI
+            or data_type == DataType.PDF
         ):
             data_sequence = frame_level_data.frame_number
 
@@ -1558,7 +1559,7 @@ def _to_encord_labels(self, frame_level_data: FrameLevelImageGroupData) -> Dict[
         ret: Dict[str, Any] = {}
         data_type = self._label_row_read_only_data.data_type
 
-        if data_type == DataType.IMAGE or data_type == DataType.IMG_GROUP:
+        if data_type == DataType.IMAGE or data_type == DataType.IMG_GROUP or data_type == DataType.PDF:
             frame = frame_level_data.frame_number
             ret.update(self._to_encord_label(frame))