demo

papermage/recipes/core_recipe.py

@@ -6,7 +6,8 @@
 
 import logging
 from pathlib import Path
-from typing import Union
+from typing import Dict, List, Union
+import warnings
 from papermage.predictors.sklearn_predictors.word_predictor import make_text
 
 from papermage.utils.annotate import group_by
@@ -21,6 +22,7 @@
     Document,
     EntitiesFieldName,
     Entity,
+    Box,
     EquationsFieldName,
     FiguresFieldName,
     FootersFieldName,
@@ -90,7 +92,10 @@ def __init__(
         self.parser = PDFPlumberParser()
         self.rasterizer = PDF2ImageRasterizer()
 
-        self.word_predictor = SVMWordPredictor.from_path(svm_word_predictor_path)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            self.word_predictor = SVMWordPredictor.from_path(svm_word_predictor_path)
+
         self.effdet_publaynet_predictor = LPBlockPredictor.from_pretrained(effdet_publaynet_predictor_path)
         # self.effdet_mfd_predictor = LPBlockPredictor.from_pretrained(effdet_mfd_predictor_path)
         self.ivila_predictor = IVILATokenClassificationPredictor.from_pretrained(ivila_predictor_path)
@@ -126,6 +131,7 @@ def from_path(self, pdfpath: str) -> Document:
 
     def from_doc(self, doc: Document) -> Document:
         self.logger.info("Predicting words...")
+
         words = self.word_predictor.predict(doc=doc)
         doc.annotate_entity(field_name=WordsFieldName, entities=words)
 
@@ -134,15 +140,37 @@ def from_doc(self, doc: Document) -> Document:
         doc.annotate_entity(field_name=SentencesFieldName, entities=sentences)
 
         self.logger.info("Predicting blocks...")
-        blocks = self.effdet_publaynet_predictor.predict(doc=doc)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            blocks = self.effdet_publaynet_predictor.predict(doc=doc)
         doc.annotate_entity(field_name=BlocksFieldName, entities=blocks)
 
         self.logger.info("Predicting vila...")
         vila_entities = self.ivila_predictor.predict(doc=doc)
+        doc.annotate_entity(field_name="vila_entities", entities=vila_entities)
+
         for entity in vila_entities:
+            entity.boxes = [
+                Box.create_enclosing_box(
+                    [b for t in doc.find_by_span(entity, field_name=TokensFieldName) for b in t.boxes]
+                )
+            ]
             entity.text = make_text(entity=entity, document=doc)
-        doc.annotate_entity(field_name="vila_entities", entities=vila_entities)
         preds = group_by(entities=vila_entities, metadata_field="label", metadata_values_map=VILA_LABELS_MAP)
         doc.annotate(*preds)
-
         return doc
+
+
+if __name__ == "__main__":
+    import argparse
+    import json
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--pdf", required=True, type=str, help="Path to PDF file.")
+    parser.add_argument("--output", type=str, help="Path to output JSON file.")
+    args = parser.parse_args()
+
+    recipe = CoreRecipe()
+    doc = recipe.from_path(pdfpath=args.pdf)
+    with open(args.output, "w") as f:
+        json.dump(doc.to_json(), f, indent=2)

papermage/visualizers/__init__.py

@@ -0,0 +1,3 @@
+from .visualizer import plot_entities_on_page
+
+__all__ = ["plot_entities_on_page"]