Normalize list-returning detection models

CHANGELOG.md

@@ -1,3 +1,8 @@
+## 1.6.14
+
+### Fixes
+- **Normalize object-detection model results**: the object-detection boundary now converts legacy list results to `LayoutElements`, so page-layout routing works consistently across built-in detection models.
+
 ## 1.6.13
 
 ### Fixes

test_unstructured_inference/inference/test_layout.py

@@ -122,6 +122,42 @@ def test_get_page_elements(monkeypatch, mock_final_layout):
     assert elements == page.elements_array
 
 
+class ListPredictingDetectionModel(UnstructuredObjectDetectionModel):
+    def __init__(self, predictions):
+        super().__init__()
+        self.predictions = predictions
+
+    def initialize(self, *args, **kwargs):
+        pass
+
+    def predict(self, x):
+        return self.predictions
+
+
+@pytest.mark.parametrize(
+    ("predictions", "expected_texts"),
+    [
+        ([layoutelement.LayoutElement.from_coords(0, 0, 1, 1, text="detected")], ["detected"]),
+        ([], []),
+    ],
+)
+def test_get_page_elements_with_detection_model_returning_a_list(
+    mock_image,
+    predictions,
+    expected_texts,
+):
+    page = layout.PageLayout(
+        number=1,
+        image=mock_image,
+        detection_model=ListPredictingDetectionModel(predictions),
+    )
+
+    result = page.get_elements_with_detection_model(inplace=False)
+
+    assert isinstance(result, layoutelement.LayoutElements)
+    assert result.texts.tolist() == expected_texts
+
+
 class MockPool:
     def map(self, f, xs):
         return [f(x) for x in xs]

test_unstructured_inference/models/test_detectron2onnx.py

@@ -6,6 +6,7 @@
 
 import unstructured_inference.models.base as models
 import unstructured_inference.models.detectron2onnx as detectron2
+from unstructured_inference.inference.layoutelement import LayoutElements
 
 
 class MockDetectron2ONNXLayoutModel:
@@ -50,7 +51,7 @@ def test_unstructured_detectron_model():
     model.model = 1
     with patch.object(detectron2.UnstructuredDetectronONNXModel, "predict", return_value=[]):
         result = model(None)
-    assert isinstance(result, list)
+    assert isinstance(result, LayoutElements)
     assert len(result) == 0
 
 
@@ -67,8 +68,9 @@ def test_inference():
             image = Image.open(fp)
             image.load()
         elements = model(image)
+        assert isinstance(elements, LayoutElements)
         assert len(elements) == 1
-        element = elements[0]
+        element = elements.as_list()[0]
         (x1, y1), _, (x2, y2), _ = element.bbox.coordinates
         assert hasattr(
             element,

unstructured_inference/__version__.py

@@ -1 +1 @@
-__version__ = "1.6.13"  # pragma: no cover
+__version__ = "1.6.14"  # pragma: no cover

unstructured_inference/models/unstructuredmodel.py

@@ -61,7 +61,10 @@ def predict(self, x: Image) -> LayoutElements | list[LayoutElement]:
 
     def __call__(self, x: Image) -> LayoutElements:
         """Inference using function call interface."""
-        return super().__call__(x)
+        predictions = super().__call__(x)
+        if isinstance(predictions, LayoutElements):
+            return predictions
+        return LayoutElements.from_list(predictions)
 
     @staticmethod
     def enhance_regions(