add tests for ffdetr vs. ffdnet

Author: jbarrow

commonforms/inference.py

@@ -25,21 +25,18 @@
     ("FFDNET-S", False): ("jbarrow/FFDNet-S", "FFDNet-S.pt"),
     ("FFDNET-L", True): ("jbarrow/FFDNet-L-cpu", "FFDNet-L.onnx"),
     ("FFDNET-L", False): ("jbarrow/FFDNet-L", "FFDNet-L.pt"),
-    ("FFDETR", False): ("jbarrow/FFDetr", "FFDetr.pth")
+    ("FFDETR", False): ("jbarrow/FFDetr", "FFDetr.pth"),
 }
 
 
-
 def batch(lst: list, n: int = 8):
     l = len(lst)
     for ndx in range(0, l, n):
-        yield lst[ndx:min(ndx + n, l)]
+        yield lst[ndx : min(ndx + n, l)]
 
 
 class FFDetrDetector:
-    def __init__(
-        self, model_or_path: str, device: int | str = "cpu"
-    ) -> None:
+    def __init__(self, model_or_path: str, device: int | str = "cpu") -> None:
         self.device = device
         self.model = RFDETRMedium(pretrain_weights=self.get_model_path(model_or_path))
 
@@ -49,8 +46,8 @@ def get_model_path(self, model_or_path: str) -> str:
         model_upper = model_or_path.upper()
         if model_upper in ["FFDETR"]:
             # download the model, will just use the cached version if it already exists
-            repo_id, filename = models[(model_upper, False)] 
-            model_path = hf_hub_download(repo_id=repo_id, filename=filename) 
+            repo_id, filename = models[(model_upper, False)]
+            model_path = hf_hub_download(repo_id=repo_id, filename=filename)
         else:
             model_path = model_or_path
 
@@ -75,7 +72,7 @@ def extract_widgets(
     ) -> dict[int, list[Widget]]:
         image_size = 1024
         results = []
-        for b in batch([p.image for p in pages], n=batch_size):
+        for b in batch([self.resize(p.image, image_size) for p in pages], n=batch_size):
             predictions = self.model.predict(b, threshold=confidence)
             if len(pages) == 1 or batch_size == 1:
                 predictions = [predictions]
@@ -131,8 +128,8 @@ def get_model_path(
         model_upper = model_or_path.upper()
         if model_upper in ["FFDNET-S", "FFDNET-L"]:
             # download the model, will just use the cached version if it already exists
-            repo_id, filename = models[(model_upper, fast)] 
-            model_path = hf_hub_download(repo_id=repo_id, filename=filename) 
+            repo_id, filename = models[(model_upper, fast)]
+            model_path = hf_hub_download(repo_id=repo_id, filename=filename)
         else:
             model_path = model_or_path
 
@@ -247,17 +244,20 @@ def prepare_form(
     input_path: str | Path,
     output_path: str | Path,
     *,
-    model_or_path: str = "FFDNet-L",
+    model_or_path: str = "FFDetr",
     keep_existing_fields: bool = False,
     use_signature_fields: bool = False,
     device: int | str = "cpu",
-    image_size: int = 1600,
-    confidence: float = 0.3,
+    image_size: int = 1024,
+    confidence: float = 0.4,
     fast: bool = False,
     multiline: bool = False,
+    batch_size: int = 4,
 ):
-    # detector = FFDNetDetector(model_or_path, device=device, fast=fast)
-    detector = FFDetrDetector("FFDetr")
+    if "FFDNET" in model_or_path.upper():
+        detector = FFDNetDetector(model_or_path, device=device, fast=fast)
+    else:
+        detector = FFDetrDetector(model_or_path)
 
     try:
         pages = render_pdf(input_path)
@@ -277,7 +277,9 @@ def prepare_form(
             name = f"{widget.widget_type.lower()}_{widget.page}_{i}"
 
             if widget.widget_type == "TextBox":
-                writer.add_text_box(name, page_ix, widget.bounding_box, multiline=multiline)
+                writer.add_text_box(
+                    name, page_ix, widget.bounding_box, multiline=multiline
+                )
             elif widget.widget_type == "ChoiceButton":
                 writer.add_checkbox(name, page_ix, widget.bounding_box)
             elif widget.widget_type == "Signature":

tests/inference_test.py

@@ -8,7 +8,7 @@
 def test_inference(tmp_path):
     # tmp_path is a built-in pythest fixture where we'll write the outputs
     output_path = tmp_path / "output.pdf"
-    commonforms.prepare_form("./tests/resources/input.pdf", output_path)
+    commonforms.prepare_form("./tests/resources/input.pdf", output_path, model_or_path="FFDetr")
 
     assert output_path.exists()
 
@@ -20,7 +20,7 @@ def test_inference(tmp_path):
 
 def test_inference_fast(tmp_path):
     output_path = tmp_path / "output.pdf"
-    commonforms.prepare_form("./tests/resources/input.pdf", output_path, fast=True)
+    commonforms.prepare_form("./tests/resources/input.pdf", output_path, fast=True, model_or_path="FFDNet-L")
 
     assert output_path.exists()
 
@@ -32,7 +32,9 @@ def test_inference_fast(tmp_path):
 
 def test_mutlinline(tmp_path):
     output_path = tmp_path / "output.pdf"
-    commonforms.prepare_form("./tests/resources/input.pdf", output_path, fast=True, multiline=True)
+    commonforms.prepare_form(
+        "./tests/resources/input.pdf", output_path, fast=True, multiline=True
+    )
 
     assert output_path.exists()
 
@@ -42,7 +44,6 @@ def test_mutlinline(tmp_path):
     doc.document.close()
 
 
-
 def test_encrypted_failure(tmp_path):
     # Reminder to future Joe: password for encrypted PDF is "kanbanery"
     output_path = tmp_path / "output.pdf"
@@ -51,7 +52,22 @@ def test_encrypted_failure(tmp_path):
         commonforms.prepare_form("./tests/resources/encrypted.pdf", output_path)
 
 
+def test_inference_ffdetr(tmp_path):
+    # tmp_path is a built-in pythest fixture where we'll write the outputs
+    output_path = tmp_path / "output.pdf"
+    commonforms.prepare_form(
+        "./tests/resources/input.pdf", output_path, model_or_path="FFDetr"
+    )
+
+    assert output_path.exists()
+
+    doc = formalpdf.open(output_path)
+    assert len(doc[0].widgets()) > 0
+
+    doc.document.close()
+
+
 # TODO(joe): future tests around handling encrypted PDFs
 #   1. add a --password flag and test that inference doesn't fail
-#   2. if a password is provided, ensure that the _output_ PDF remains encrpyted 
-#      with the same password 
+#   2. if a password is provided, ensure that the _output_ PDF remains encrpyted
+#      with the same password