From b14a341cfcfb673b0969ea0261de2c0c5c4de5e9 Mon Sep 17 00:00:00 2001
From: Julius Unverfehrt <julius.unverfehrt@iqser.com>
Date: Wed, 17 Aug 2022 14:53:07 +0200
Subject: [PATCH] readd annotate_pdf script

---
 scripts/annotate_pdf.py | 53 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 53 insertions(+)
 create mode 100644 scripts/annotate_pdf.py

diff --git a/scripts/annotate_pdf.py b/scripts/annotate_pdf.py
new file mode 100644
index 0000000..2f0ad83
--- /dev/null
+++ b/scripts/annotate_pdf.py
@@ -0,0 +1,53 @@
+import argparse
+from itertools import starmap
+from pathlib import Path
+
+from PIL import Image
+from funcy import lmap
+
+from cv_analysis.figure_detection.figure_detection import detect_figures
+from cv_analysis.layout_parsing import parse_layout
+from cv_analysis.table_parsing import parse_tables
+from cv_analysis.utils.draw import draw_rectangles
+from pdf2img.conversion import convert_pages_to_images
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("pdf_path")
+    parser.add_argument("--output_folder", default="/tmp")
+    parser.add_argument("--type", "-t", choices=["table", "layout", "figure"], required=True)
+    return parser.parse_args()
+
+
+def analyse_and_annotate(images, analysis_fn):
+    arrays = lmap(lambda image: image.asarray(), images)
+    result = map(analysis_fn, arrays)
+    annotated_images = starmap(draw_rectangles, zip(arrays, result))
+    return annotated_images
+
+
+def save_as_pdf(images, output_folder, file_name, operation):
+    Path(output_folder).mkdir(parents=True, exist_ok=True)
+    images = lmap(Image.fromarray, images)
+    images[0].save(f"{output_folder}/{file_name}_annotated_{operation}.pdf", save_all=True, append_images=images)
+
+
+def get_analysis_fn(analysis_type):
+    if analysis_type == "table":
+        return parse_tables
+    elif analysis_type == "layout":
+        return parse_layout
+    elif analysis_type == "figure":
+        return detect_figures
+    else:
+        raise
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    with open(args.pdf_path, "rb") as f:
+        pdf_bytes = f.read()
+        images = convert_pages_to_images(pdf_bytes)
+    annotated_pages = analyse_and_annotate(images=images, analysis_fn=get_analysis_fn(args.type))
+    save_as_pdf(annotated_pages, args.output_folder, Path(args.pdf_path).stem, args.type)