Skip to content

Commit 1dc9092

Browse files
authored
Fix FFDetrDetector crash on single-item final batch (#31)
1 parent ab5fab9 commit 1dc9092

1 file changed

Lines changed: 12 additions & 6 deletions

File tree

commonforms/inference.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -74,9 +74,10 @@ def extract_widgets(
7474
results = []
7575
for b in batch([p.image for p in pages], n=batch_size):
7676
predictions = self.model.predict(b, threshold=confidence)
77-
if len(pages) == 1 or batch_size == 1:
78-
predictions = [predictions]
79-
results.extend(predictions)
77+
if isinstance(predictions, list):
78+
results.extend(predictions)
79+
else:
80+
results.append(predictions)
8081

8182
widgets = {}
8283

@@ -264,9 +265,14 @@ def prepare_form(
264265
except pypdfium2._helpers.misc.PdfiumError:
265266
raise EncryptedPdfError
266267

267-
results = detector.extract_widgets(
268-
pages, confidence=confidence, image_size=image_size
269-
)
268+
if isinstance(detector, FFDetrDetector):
269+
results = detector.extract_widgets(
270+
pages, confidence=confidence, image_size=image_size, batch_size=batch_size
271+
)
272+
else:
273+
results = detector.extract_widgets(
274+
pages, confidence=confidence, image_size=image_size
275+
)
270276

271277
writer = PyPdfFormCreator(input_path)
272278
if not keep_existing_fields:

0 commit comments

Comments
 (0)