From 9b7226d8fd98d2ee793a93c3e1b53c2ca23cf2a5 Mon Sep 17 00:00:00 2001 From: Ritwij Aryan Parmar Date: Tue, 26 May 2026 14:32:01 -0400 Subject: [PATCH 1/2] feat: expose layout confidence metadata --- .../partition/common/test_common.py | 24 ++++++++++++++++ test_unstructured/staging/test_base.py | 4 +++ unstructured/documents/elements.py | 8 ++++++ unstructured/partition/common/common.py | 28 ++++++++++++++----- unstructured/staging/base.py | 6 +++- 5 files changed, 62 insertions(+), 8 deletions(-) diff --git a/test_unstructured/partition/common/test_common.py b/test_unstructured/partition/common/test_common.py index 35825dedf4..a696183d64 100644 --- a/test_unstructured/partition/common/test_common.py +++ b/test_unstructured/partition/common/test_common.py @@ -87,6 +87,30 @@ def test_normalize_layout_element_dict(): ) +def test_normalize_layout_element_preserves_layout_confidence_metadata(): + layout_element = { + "type": "Title", + "coordinates": [[1, 2], [3, 4], [5, 6], [7, 8]], + "text": "Some uncertain text", + "prob": 0.0, + "source": "ocr", + } + coordinate_system = PixelSpace(width=10, height=20) + element = common.normalize_layout_element( + layout_element, + coordinate_system=coordinate_system, + ) + + assert isinstance(element, Title) + assert element.metadata.confidence_score == 0.0 + assert element.metadata.detection_class_prob == 0.0 + assert element.metadata.extraction_method == "ocr" + metadata_dict = element.metadata.to_dict() + assert metadata_dict["confidence_score"] == 0.0 + assert metadata_dict["detection_class_prob"] == 0.0 + assert metadata_dict["extraction_method"] == "ocr" + + def test_normalize_layout_element_dict_caption(): layout_element = { "type": "Figure", diff --git a/test_unstructured/staging/test_base.py b/test_unstructured/staging/test_base.py index 8f2b0c35f6..568651ab73 100644 --- a/test_unstructured/staging/test_base.py +++ b/test_unstructured/staging/test_base.py @@ -194,6 +194,8 @@ def test_default_pandas_dtypes(): attached_to_filename="attached_to_filename", parent_id="parent_id", category_depth=1, + confidence_score=0.0, + extraction_method="ocr", image_path="image_path", languages=["eng", "spa"], page_number=1, @@ -350,6 +352,8 @@ def test_convert_to_coco(): attached_to_filename="attached_to_filename", parent_id="parent_id", category_depth=1, + confidence_score=0.0, + extraction_method="ocr", image_path="image_path", languages=["eng", "spa"], page_number=1, diff --git a/unstructured/documents/elements.py b/unstructured/documents/elements.py index 05a64a592a..6f0dfc0fc8 100644 --- a/unstructured/documents/elements.py +++ b/unstructured/documents/elements.py @@ -162,6 +162,7 @@ class ElementMetadata: attached_to_filename: Optional[str] category_depth: Optional[int] + confidence_score: Optional[float] coordinates: Optional[CoordinatesMetadata] data_source: Optional[DataSourceMetadata] # -- Detection Model Class Probabilities from Unstructured-Inference Hi-Res -- @@ -170,6 +171,7 @@ class ElementMetadata: detection_origin: Optional[str] emphasized_text_contents: Optional[list[str]] emphasized_text_tags: Optional[list[str]] + extraction_method: Optional[str] file_directory: Optional[str] filename: Optional[str] filetype: Optional[str] @@ -236,9 +238,11 @@ def __init__( bcc_recipient: Optional[list[str]] = None, category_depth: Optional[int] = None, cc_recipient: Optional[list[str]] = None, + confidence_score: Optional[float] = None, coordinates: Optional[CoordinatesMetadata] = None, data_source: Optional[DataSourceMetadata] = None, detection_class_prob: Optional[float] = None, + extraction_method: Optional[str] = None, emphasized_text_contents: Optional[list[str]] = None, emphasized_text_tags: Optional[list[str]] = None, file_directory: Optional[str] = None, @@ -281,9 +285,11 @@ def __init__( self.bcc_recipient = bcc_recipient self.category_depth = category_depth self.cc_recipient = cc_recipient + self.confidence_score = confidence_score self.coordinates = coordinates self.data_source = data_source self.detection_class_prob = detection_class_prob + self.extraction_method = extraction_method self.emphasized_text_contents = emphasized_text_contents self.emphasized_text_tags = emphasized_text_tags @@ -515,10 +521,12 @@ def field_consolidation_strategies(cls) -> dict[str, ConsolidationStrategy]: "cc_recipient": cls.FIRST, "bcc_recipient": cls.FIRST, "category_depth": cls.DROP, + "confidence_score": cls.DROP, "coordinates": cls.DROP, "data_source": cls.FIRST, "detection_class_prob": cls.DROP, "detection_origin": cls.DROP, + "extraction_method": cls.DROP, "emphasized_text_contents": cls.LIST_CONCATENATE, "emphasized_text_tags": cls.LIST_CONCATENATE, "file_directory": cls.FIRST, diff --git a/unstructured/partition/common/common.py b/unstructured/partition/common/common.py index d18fc8c87b..71f8918be2 100644 --- a/unstructured/partition/common/common.py +++ b/unstructured/partition/common/common.py @@ -31,6 +31,17 @@ from unstructured_inference.inference.layoutelement import LayoutElement +def _layout_source_to_extraction_method(source: Any) -> str | None: + """Returns a serializable layout extraction method from inference source metadata.""" + if source is None: + return None + if isinstance(source, str): + return source + if hasattr(source, "value"): + return str(source.value) + return str(source) + + def normalize_layout_element( layout_element: LayoutElement | Element | dict[str, Any], coordinate_system: Optional[CoordinateSystem] = None, @@ -58,23 +69,26 @@ def normalize_layout_element( element_type = layout_dict.get("type") prob = layout_dict.get("prob") aux_origin = layout_dict.get("source", None) - origin = None + extraction_method = _layout_source_to_extraction_method(aux_origin) if isinstance(layout_dict.get("is_extracted"), Enum): is_extracted = layout_dict["is_extracted"].value else: is_extracted = None - if aux_origin: - origin = aux_origin.value - if prob and isinstance(prob, (int, str, float, numbers.Number)): - class_prob_metadata = ElementMetadata(detection_class_prob=float(prob)) # type: ignore + if prob is not None and isinstance(prob, (int, str, float, numbers.Number)): + confidence_score = float(prob) + class_prob_metadata = ElementMetadata( + confidence_score=confidence_score, + detection_class_prob=confidence_score, + extraction_method=extraction_method, + ) else: - class_prob_metadata = ElementMetadata() + class_prob_metadata = ElementMetadata(extraction_method=extraction_method) class_prob_metadata.is_extracted = is_extracted common_kwargs = { "coordinates": coordinates, "coordinate_system": coordinate_system, "metadata": class_prob_metadata, - "detection_origin": origin, + "detection_origin": extraction_method, } if element_type == ElementType.LIST: if infer_list_items: diff --git a/unstructured/staging/base.py b/unstructured/staging/base.py index 47d8c99f83..e21eda0058 100644 --- a/unstructured/staging/base.py +++ b/unstructured/staging/base.py @@ -493,7 +493,9 @@ def _fix_metadata_field_precision(elements: Iterable[Element]) -> list[Element]: rounded_points.append(rounded_point) el.metadata.coordinates.points = tuple(rounded_points) - if el.metadata.detection_class_prob: + if el.metadata.confidence_score is not None: + el.metadata.confidence_score = round(el.metadata.confidence_score, 5) + if el.metadata.detection_class_prob is not None: el.metadata.detection_class_prob = round(el.metadata.detection_class_prob, 5) out_elements.append(el) @@ -641,6 +643,8 @@ def get_default_pandas_dtypes() -> dict[str, Any]: "attached_to_filename": pd.StringDtype(), # Optional[str] # type: ignore "parent_id": pd.StringDtype(), # Optional[str], # type: ignore "category_depth": "Int64", # Optional[int] + "confidence_score": float, # Optional[float], + "extraction_method": pd.StringDtype(), # Optional[str] # type: ignore "image_path": pd.StringDtype(), # Optional[str] # type: ignore "languages": object, # Optional[list[str]] "page_number": "Int64", # Optional[int] From 1bb2ee39c96d39f18bbecf5ec2aa62c73885fd7c Mon Sep 17 00:00:00 2001 From: Ritwij Aryan Parmar Date: Tue, 26 May 2026 16:46:04 -0400 Subject: [PATCH 2/2] Ignore empty layout confidence values --- .../partition/common/test_common.py | 20 +++++++++++++++++++ unstructured/partition/common/common.py | 2 +- 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/test_unstructured/partition/common/test_common.py b/test_unstructured/partition/common/test_common.py index a696183d64..0899b35010 100644 --- a/test_unstructured/partition/common/test_common.py +++ b/test_unstructured/partition/common/test_common.py @@ -111,6 +111,26 @@ def test_normalize_layout_element_preserves_layout_confidence_metadata(): assert metadata_dict["extraction_method"] == "ocr" +def test_normalize_layout_element_ignores_empty_layout_confidence_metadata(): + layout_element = { + "type": "Title", + "coordinates": [[1, 2], [3, 4], [5, 6], [7, 8]], + "text": "Some text with no confidence score", + "prob": "", + "source": "ocr", + } + coordinate_system = PixelSpace(width=10, height=20) + element = common.normalize_layout_element( + layout_element, + coordinate_system=coordinate_system, + ) + + assert isinstance(element, Title) + assert element.metadata.confidence_score is None + assert element.metadata.detection_class_prob is None + assert element.metadata.extraction_method == "ocr" + + def test_normalize_layout_element_dict_caption(): layout_element = { "type": "Figure", diff --git a/unstructured/partition/common/common.py b/unstructured/partition/common/common.py index 71f8918be2..ba6ebd12e7 100644 --- a/unstructured/partition/common/common.py +++ b/unstructured/partition/common/common.py @@ -74,7 +74,7 @@ def normalize_layout_element( is_extracted = layout_dict["is_extracted"].value else: is_extracted = None - if prob is not None and isinstance(prob, (int, str, float, numbers.Number)): + if prob not in (None, "") and isinstance(prob, (int, str, float, numbers.Number)): confidence_score = float(prob) class_prob_metadata = ElementMetadata( confidence_score=confidence_score,