Skip to content

Commit ebd90b9

Browse files
authored
Update detectron2onnx.py
The proposed solution by @hegderavin at Unstructured-IO/unstructured#3718 (comment)
1 parent 5dc66b4 commit ebd90b9

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

unstructured_inference/models/detectron2onnx.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
from PIL import Image
1111

1212
from unstructured_inference.constants import Source
13-
from unstructured_inference.inference.layoutelement import LayoutElement
13+
from unstructured_inference.inference.layoutelement import LayoutElement, LayoutElements
1414
from unstructured_inference.logger import logger, logger_onnx
1515
from unstructured_inference.models.unstructuredmodel import (
1616
UnstructuredObjectDetectionModel,
@@ -88,7 +88,7 @@ def predict(self, image: Image.Image) -> List[LayoutElement]:
8888
logger_onnx.debug(
8989
"Ignoring runtime error from onnx (likely due to encountering blank page).",
9090
)
91-
return []
91+
return LayoutElements.from_list([])
9292
input_w, input_h = image.size
9393
regions = self.postprocess(bboxes, labels, confidence_scores, input_w, input_h)
9494

@@ -174,4 +174,4 @@ def postprocess(
174174
regions.append(region)
175175

176176
regions.sort(key=lambda element: element.bbox.y1)
177-
return cast(List[LayoutElement], regions)
177+
return LayoutElements.from_list(cast(List[LayoutElement], regions))

0 commit comments

Comments
 (0)