From 63f1176d995ad803e0a01a24bd9869700c268bef Mon Sep 17 00:00:00 2001 From: Keetarp Date: Mon, 29 Dec 2025 00:45:31 -0500 Subject: [PATCH] Fix FFDetrDetector crash on single-item final batch --- commonforms/inference.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/commonforms/inference.py b/commonforms/inference.py index ce8cc60..9eed964 100644 --- a/commonforms/inference.py +++ b/commonforms/inference.py @@ -74,9 +74,10 @@ def extract_widgets( results = [] for b in batch([p.image for p in pages], n=batch_size): predictions = self.model.predict(b, threshold=confidence) - if len(pages) == 1 or batch_size == 1: - predictions = [predictions] - results.extend(predictions) + if isinstance(predictions, list): + results.extend(predictions) + else: + results.append(predictions) widgets = {} @@ -264,9 +265,14 @@ def prepare_form( except pypdfium2._helpers.misc.PdfiumError: raise EncryptedPdfError - results = detector.extract_widgets( - pages, confidence=confidence, image_size=image_size - ) + if isinstance(detector, FFDetrDetector): + results = detector.extract_widgets( + pages, confidence=confidence, image_size=image_size, batch_size=batch_size + ) + else: + results = detector.extract_widgets( + pages, confidence=confidence, image_size=image_size + ) writer = PyPdfFormCreator(input_path) if not keep_existing_fields: