From ebd90b91e19ba2c3e562d40f68cf42bc9603bbd4 Mon Sep 17 00:00:00 2001 From: Yasiru Date: Mon, 28 Apr 2025 15:55:55 +1000 Subject: [PATCH] Update detectron2onnx.py The proposed solution by @hegderavin at https://github.com/Unstructured-IO/unstructured/issues/3718#issuecomment-2644039613 --- unstructured_inference/models/detectron2onnx.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/unstructured_inference/models/detectron2onnx.py b/unstructured_inference/models/detectron2onnx.py index 79cd0a1a..d4e9d5f4 100644 --- a/unstructured_inference/models/detectron2onnx.py +++ b/unstructured_inference/models/detectron2onnx.py @@ -10,7 +10,7 @@ from PIL import Image from unstructured_inference.constants import Source -from unstructured_inference.inference.layoutelement import LayoutElement +from unstructured_inference.inference.layoutelement import LayoutElement, LayoutElements from unstructured_inference.logger import logger, logger_onnx from unstructured_inference.models.unstructuredmodel import ( UnstructuredObjectDetectionModel, @@ -88,7 +88,7 @@ def predict(self, image: Image.Image) -> List[LayoutElement]: logger_onnx.debug( "Ignoring runtime error from onnx (likely due to encountering blank page).", ) - return [] + return LayoutElements.from_list([]) input_w, input_h = image.size regions = self.postprocess(bboxes, labels, confidence_scores, input_w, input_h) @@ -174,4 +174,4 @@ def postprocess( regions.append(region) regions.sort(key=lambda element: element.bbox.y1) - return cast(List[LayoutElement], regions) + return LayoutElements.from_list(cast(List[LayoutElement], regions))