From 39c111fd4220359de27f36dd98a13a66196b3506 Mon Sep 17 00:00:00 2001 From: Matthias Bisping Date: Sun, 3 Apr 2022 04:08:00 +0200 Subject: [PATCH] integrated PDFNet coordinate transformer into pipeline --- image_prediction/default_objects.py | 5 +- .../coordinate/coordinate_transformer.py | 7 +- image_prediction/utils/pdf_annotation.py | 99 +++++++++++++++++++ scripts/run_pipeline.py | 19 ++-- ...ca794fccb8cd38b95f2bf6ba9_predictions.json | 4 +- 5 files changed, 123 insertions(+), 11 deletions(-) create mode 100644 image_prediction/utils/pdf_annotation.py diff --git a/image_prediction/default_objects.py b/image_prediction/default_objects.py index e7b3856..2bb8cb6 100644 --- a/image_prediction/default_objects.py +++ b/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 diff --git a/image_prediction/transformer/transformers/coordinate/coordinate_transformer.py b/image_prediction/transformer/transformers/coordinate/coordinate_transformer.py index c033e09..b725a51 100644 --- a/image_prediction/transformer/transformers/coordinate/coordinate_transformer.py +++ b/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) diff --git a/image_prediction/utils/pdf_annotation.py b/image_prediction/utils/pdf_annotation.py new file mode 100644 index 0000000..a463a53 --- /dev/null +++ b/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}") diff --git a/scripts/run_pipeline.py b/scripts/run_pipeline.py index 7972ba4..f0d3be8 100644 --- a/scripts/run_pipeline.py +++ b/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) diff --git a/test/data/f2dc689ca794fccb8cd38b95f2bf6ba9_predictions.json b/test/data/f2dc689ca794fccb8cd38b95f2bf6ba9_predictions.json index 76fbcb8..1e95721 100644 --- a/test/data/f2dc689ca794fccb8cd38b95f2bf6ba9_predictions.json +++ b/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": {