Update unstructured inference detection models to output class probs (#191)

Updated the detection models we currently support to output
classification probabilities and persist those through
unstructured-inference LayoutElement. Next step once this version is
released is to update Unstructured package to allow the probabilities
from unstructured-inference to persist all the way to the end of
partition through metadata.

---------

Co-authored-by: Alan Bertl <[email protected]>

Author: pravin-unstructured

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.5.22
+
+* Add object-detection classification probabilities to LayoutElement for all currently implemented object detection models
+
 ## 0.5.21
 
 * adds `safe_division` to replae 0 with machine epsilon for `float` to avoid division by 0

test_unstructured_inference/models/test_detectron2onnx.py

@@ -70,6 +70,14 @@ def test_inference():
         assert len(elements) == 1
         element = elements[0]
         (x1, y1), _, (x2, y2), _ = element.coordinates
+        assert hasattr(
+            element,
+            "prob",
+        )  # NOTE(pravin) New Assertion to Make Sure element has probabilities
+        assert isinstance(
+            element.prob,
+            float,
+        )  # NOTE(pravin) New Assertion to Make Sure Populated Probability is Float
         # NOTE(alan): The bbox coordinates get resized, so check their relative proportions
         assert x2 / x1 == pytest.approx(3.0)  # x1 == 1, x2 == 3 before scaling
         assert y2 / y1 == pytest.approx(2.0)  # y1 == 2, y2 == 4 before scaling

test_unstructured_inference/models/test_model.py

@@ -63,6 +63,12 @@ def test_model_initializes_once():
         "models",
         {},
     ):
-        doc = layout.DocumentLayout.from_file("sample-docs/layout-parser-paper.pdf")
-
+        doc = layout.DocumentLayout.from_file("sample-docs/loremipsum.pdf")
         doc.pages[0].detection_model.initializer.assert_called_once()
+        assert hasattr(
+            doc.pages[0].elements[0],
+            "prob",
+        )  # NOTE(pravin) New Assertion to Make Sure Elements have probability attribute
+        assert (
+            doc.pages[0].elements[0].prob is None
+        )  # NOTE(pravin) New Assertion to Make Sure Uncategorized Text has None Probability

test_unstructured_inference/models/test_yolox.py

@@ -15,6 +15,14 @@ def test_layout_yolox_local_parsing_image():
     assert len(document_layout.pages) == 1
     # NOTE(benjamin) The example sent to the test contains 13 detections
     assert len(document_layout.pages[0].elements) == 13
+    assert hasattr(
+        document_layout.pages[0].elements[0],
+        "prob",
+    )  # NOTE(pravin) New Assertion to Make Sure LayoutElement has probabilities
+    assert isinstance(
+        document_layout.pages[0].elements[0].prob,
+        float,
+    )  # NOTE(pravin) New Assertion to Make Sure Populated Probability is Float
 
 
 @pytest.mark.slow()
@@ -26,6 +34,14 @@ def test_layout_yolox_local_parsing_pdf():
     assert len(document_layout.pages) == 1
     # NOTE(benjamin) The example sent to the test contains 5 detections
     assert len(document_layout.pages[0].elements) == 5
+    assert hasattr(
+        document_layout.pages[0].elements[0],
+        "prob",
+    )  # NOTE(pravin) New Assertion to Make Sure LayoutElement has probabilities
+    assert isinstance(
+        document_layout.pages[0].elements[0].prob,
+        float,
+    )  # NOTE(pravin) New Assertion to Make Sure Populated Probability is Float
 
 
 @pytest.mark.slow()
@@ -51,6 +67,14 @@ def test_layout_yolox_local_parsing_image_soft():
     assert len(document_layout.pages) == 1
     # NOTE(benjamin) Soft version of the test, run make test-long in order to run with full model
     assert len(document_layout.pages[0].elements) > 0
+    assert hasattr(
+        document_layout.pages[0].elements[0],
+        "prob",
+    )  # NOTE(pravin) New Assertion to Make Sure LayoutElement has probabilities
+    assert isinstance(
+        document_layout.pages[0].elements[0].prob,
+        float,
+    )  # NOTE(pravin) New Assertion to Make Sure Populated Probability is Float
 
 
 def test_layout_yolox_local_parsing_pdf_soft():
@@ -61,6 +85,13 @@ def test_layout_yolox_local_parsing_pdf_soft():
     assert len(document_layout.pages) == 1
     # NOTE(benjamin) Soft version of the test, run make test-long in order to run with full model
     assert len(document_layout.pages[0].elements) > 0
+    assert hasattr(
+        document_layout.pages[0].elements[0],
+        "prob",
+    )  # NOTE(pravin) New Assertion to Make Sure LayoutElement has probabilities
+    assert (
+        document_layout.pages[0].elements[0].prob is None
+    )  # NOTE(pravin) New Assertion to Make Sure Uncategorized Text has None Probability
 
 
 def test_layout_yolox_local_parsing_empty_pdf_soft():

unstructured_inference/__version__.py

@@ -1 +1 @@
-__version__ = "0.5.21"  # pragma: no cover
+__version__ = "0.5.22"  # pragma: no cover

unstructured_inference/inference/layoutelement.py

@@ -21,6 +21,7 @@
 @dataclass
 class LayoutElement(TextRegion):
     type: Optional[str] = None
+    prob: Optional[float] = None
 
     def extract_text(
         self,
@@ -48,6 +49,7 @@ def to_dict(self) -> dict:
             "coordinates": self.coordinates,
             "text": self.text,
             "type": self.type,
+            "prob": self.prob,
         }
         return out_dict
 
@@ -65,7 +67,8 @@ def from_lp_textblock(cls, textblock: TextBlock):
         x1, y1, x2, y2 = textblock.coordinates
         text = textblock.text
         type = textblock.type
-        return cls(x1, y1, x2, y2, text, type)
+        score = textblock.score
+        return cls(x1, y1, x2, y2, text, type, prob=score)
 
 
 def interpret_table_block(text_block: TextRegion, image: Image.Image) -> str:

unstructured_inference/models/detectron2onnx.py

@@ -157,6 +157,7 @@ def postprocess(
                     y2 * height_conversion,
                     text=None,
                     type=detected_class,
+                    prob=conf,
                 )
 
                 regions.append(region)

unstructured_inference/models/yolox.py

@@ -141,9 +141,9 @@ def image_processing(
         for det in dets:
             # Each detection should have (x1,y1,x2,y2,probability,class) format
             # being (x1,y1) the top left and (x2,y2) the bottom right
-            x1, y1, x2, y2, _, class_id = det.tolist()
+            x1, y1, x2, y2, prob, class_id = det.tolist()
             detected_class = self.layout_classes[int(class_id)]
-            region = LayoutElement(x1, y1, x2, y2, text=None, type=detected_class)
+            region = LayoutElement(x1, y1, x2, y2, text=None, type=detected_class, prob=prob)
 
             regions.append(region)