integrated PDFNet coordinate transformer into pipeline

image_prediction/default_objects.py

@@ -13,6 +13,7 @@ from image_prediction.label_mapper.mappers.probability import ProbabilityMapper
 from image_prediction.model_loader.loader import ModelLoader
 from image_prediction.model_loader.loaders.mlflow import MlflowConnector
 from image_prediction.redai_adapter.mlflow import MlflowModelReader
+from image_prediction.transformer.transformers.coordinate.pdfnet import PDFNetCoordinateTransformer
 
 
 def get_mlflow_model_loader(mlruns_dir):
@@ -40,5 +41,7 @@ def get_extractor_classifier(model_loader, model_identifier, **kwargs):
 
 
 def get_formatter():
-    formatter = TransformerCompositor(EnumFormatter(), ResponseTransformer(), Snake2CamelCaseKeyFormatter())
+    formatter = TransformerCompositor(
+        PDFNetCoordinateTransformer(), EnumFormatter(), ResponseTransformer(), Snake2CamelCaseKeyFormatter()
+    )
     return formatter

image_prediction/transformer/transformers/coordinate/coordinate_transformer.py

@@ -1,7 +1,9 @@
 import abc
 
+from image_prediction.transformer.transformer import Transformer
 
-class CoordinateTransformer:
+
+class CoordinateTransformer(Transformer):
 
     @abc.abstractmethod
     def _forward(self, metadata):
@@ -22,3 +24,6 @@ class CoordinateTransformer:
             return self._backward(metadata)
         except TypeError:
             return map(self._backward, metadata)
+
+    def transform(self, metadata):
+        return self.forward(metadata)

image_prediction/utils/pdf_annotation.py

@@ -0,0 +1,99 @@
+"""Defines utilities for PDF processing."""
+
+import json
+from operator import itemgetter
+
+from PDFNetPython3.PDFNetPython import (
+    PDFDoc,
+    PDFNet,
+    Square,
+    Rect,
+    ColorPt,
+    BorderStyle,
+    SDFDoc,
+    Point,
+    Text,
+)
+
+from image_prediction.utils import get_logger
+
+logger = get_logger()
+
+
+def annotate_image(doc, image_info):
+    def draw_box():
+        sq = Square.Create(doc.GetSDFDoc(), Rect(*coords))
+        sq.SetColor(ColorPt(*color), 3)
+        sq.SetBorderStyle(BorderStyle(BorderStyle.e_dashed, 2, 0, 0, [4, 2]))
+        sq.SetPadding(4)
+        sq.RefreshAppearance()
+        page.AnnotPushBack(sq)
+
+    def add_note():
+        txt = Text.Create(doc.GetSDFDoc(), Point(*coords[:2]))
+        txt.SetContents(json.dumps(image_info, indent=2, ensure_ascii=False))
+        txt.SetColor(ColorPt(*color))
+        page.AnnotPushBack(txt)
+        txt.RefreshAppearance()
+
+    red = (1, 0, 0)
+    green = (0, 1, 0)
+    blue = (0, 0, 1)
+
+    if image_info["filters"]["allPassed"]:
+        color = green
+    elif image_info["filters"]["probability"]["unconfident"]:
+        color = red
+    else:
+        color = blue
+
+    page = doc.GetPage(image_info["position"]["pageNumber"])
+    coords = itemgetter("x1", "y1", "x2", "y2")(image_info["position"])
+
+    draw_box()
+    add_note()
+
+
+
+def init():
+    PDFNet.Initialize(
+        "Knecon AG(en.knecon.swiss):OEM:DDA-R::WL+:AMS(20211029):BECC974307DAB4F34B513BC9B2531B24496F6FCB83CD8AC574358A959730B622FABEF5C7")
+
+
+def draw_metadata_box(pdf_path, metadata, store_path):
+
+    init()
+
+    doc = PDFDoc(pdf_path)
+
+    color = (1, 0, 0)
+
+    print(metadata)
+
+    coords = itemgetter("x1", "y1", "x2", "y2")(metadata)
+    page = doc.GetPage(1)
+
+    sq = Square.Create(doc.GetSDFDoc(), Rect(*coords))
+    sq.SetColor(ColorPt(*color), 3)
+    sq.SetBorderStyle(BorderStyle(BorderStyle.e_dashed, 2, 0, 0, [4, 2]))
+    sq.SetPadding(4)
+    sq.RefreshAppearance()
+    page.AnnotPushBack(sq)
+
+    doc.Save(store_path, SDFDoc.e_linearized)
+
+    logger.info(f"Saved annotated PDF to {store_path}")
+
+
+def annotate_pdf(pdf_path, responses, store_path):
+
+    init()
+
+    doc = PDFDoc(pdf_path)
+
+    for image_info in responses:
+        annotate_image(doc, image_info)
+
+    doc.Save(store_path, SDFDoc.e_linearized)
+
+    logger.info(f"Saved annotated PDF to {store_path}")

scripts/run_pipeline.py

@@ -1,8 +1,10 @@
 import argparse
 import json
+import os
 
 from image_prediction.pipeline import load_pipeline
-from image_prediction.utils import get_logger
+from image_prediction.transformer.transformers.coordinate.pdfnet import PDFNetCoordinateTransformer
+from image_prediction.utils.pdf_annotation import annotate_pdf
 
 
 def parse_args():
@@ -17,15 +19,18 @@ def parse_args():
 def main(args):
     pipeline = load_pipeline(verbose=True)
 
-    with open(args.pdf, "rb") as f:
-        predictions = pipeline(f.read())
+    pdf_path = args.pdf
 
-    for prd in predictions:
-        print(json.dumps(prd, indent=2))
+    with open(pdf_path, "rb") as f:
+        predictions = list(pipeline(f.read()))
+
+    print(json.dumps(predictions, indent=2))
+
+    annotate_pdf(
+        pdf_path, predictions, os.path.join("/tmp", os.path.basename(pdf_path.replace(".pdf", "_annotated.pdf")))
+    )
 
 
 if __name__ == "__main__":
-    logger = get_logger()
-
     args = parse_args()
     main(args)

test/data/f2dc689ca794fccb8cd38b95f2bf6ba9_predictions.json

@@ -12,8 +12,8 @@
     "position": {
       "x1": 321,
       "x2": 515,
-      "y1": 300,
-      "y2": 494,
+      "y1": 348,
+      "y2": 542,
       "pageNumber": 2
     },
     "geometry": {