diff --git a/scripts/parse_pdf.py b/scripts/parse_pdf.py new file mode 100644 index 0000000..e6052ca --- /dev/null +++ b/scripts/parse_pdf.py @@ -0,0 +1,15 @@ +import json + +from cv_analysis.server.pipeline import make_image_analysis_pipeline +from cv_analysis.table_inference import infer_lines + +pipe = make_image_analysis_pipeline(infer_lines) + +# FIXME: Implement argparsing + +pdf_bytes = open("/home/junverfehrt/Documents/rosario_test_file.pdf", "rb").read() +vlp_output = json.load(open("/home/junverfehrt/Documents/rosario_test_file_vlp.json", "r")) + +best_result = list(pipe(data={"pdf": pdf_bytes, "vlp_output": vlp_output})) + +print(best_result) diff --git a/src/cv_analysis/utils/annotate.py b/src/cv_analysis/utils/annotate.py new file mode 100644 index 0000000..5822186 --- /dev/null +++ b/src/cv_analysis/utils/annotate.py @@ -0,0 +1,58 @@ +from functools import singledispatch +from operator import itemgetter +from pathlib import Path +from typing import Union + +import fitz +from kn_utils.logging import logger + + +def annotate_pdf(pdf: Union[str, bytes, Path], predictions, output_path: Union[str, Path] = None): + pdf_bytes = provide_byte_stream(pdf) + with fitz.open(stream=pdf_bytes) as pdf_handle: + for prediction in predictions: + # FIXME: Adapt to line drawing + index = prediction["page_idx"] + annotate_page(pdf_handle[index], prediction) + output_path = output_path or "/tmp/annotated.pdf" + pdf_handle.save(output_path) + logger.info(f"Annotated PDF saved to {output_path}") + + +def annotate_page(page: fitz.Page, prediction): + for box in prediction["boxes"]: + bbox = itemgetter("x1", "y1", "x2", "y2")(box["box"]) + label, probability, uuid = itemgetter("label", "probability", "uuid")(box) + + bbox = mirror_on_x_axis(bbox, page.bound().height) + x0, y0, x1, y1 = bbox + page.draw_rect(fitz.Rect(x0, y0, x1, y1), color=(0, 0, 1), width=2) + label_x, label_y = x0, y0 - 5 + page.insert_text((label_x, label_y), f"{label} ({probability:.2f}), {uuid}", fontsize=12, color=(0.4, 0.4, 1)) + + return page + + +def mirror_on_x_axis(bbox, page_height): + x0, y0, x1, y1 = bbox + y0_new = page_height - y1 + y1_new = page_height - y0 + + return x0, y0_new, x1, y1_new + + +@singledispatch +def provide_byte_stream(pdf: Union[bytes, Path, str]) -> bytes: + pass + + +@provide_byte_stream.register(bytes) +def _(pdf): + return pdf + + +@provide_byte_stream.register(str) +@provide_byte_stream.register(Path) +def _(pdf): + with open(pdf, "rb") as pdf_file: + return pdf_file.read()