mirrors/docling
1
0
Fork 0
mirror of https://github.com/DS4SD/docling.git synced 2025-12-08 12:48:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: Replace the layout_predictor.predict_batch() with layout_predictor.predict() in a loop

Browse Source 
Signed-off-by: Nikos Livathinos <nli@zurich.ibm.com>
This commit is contained in:
Nikos Livathinos
2025-08-28 15:14:51 +02:00
parent 36d44f1225
commit 946ea1c2cb
1 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
docling/models/layout_model.py
View File

@@ -171,9 +171,12 @@ class LayoutModel(BasePageModel):
batch_predictions = [] batch_predictions = []
if valid_page_images: if valid_page_images:
with TimeRecorder(conv_res, "layout"): with TimeRecorder(conv_res, "layout"):
batch_predictions = self.layout_predictor.predict_batch( # type: ignore[attr-defined] # batch_predictions = self.layout_predictor.predict_batch( # type: ignore[attr-defined]
valid_page_images # valid_page_images
) # )
for vpi in valid_page_images:
all_predictions = list(self.layout_predictor.predict(vpi))
batch_predictions.append(all_predictions)
# Process each page with its predictions # Process each page with its predictions
valid_page_idx = 0 valid_page_idx = 0