@@ -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