37 lines
816 B
Python
37 lines
816 B
Python
import argparse
|
|
import json
|
|
import os
|
|
|
|
from image_prediction.pipeline import load_pipeline
|
|
from image_prediction.transformer.transformers.coordinate.pdfnet import PDFNetCoordinateTransformer
|
|
from image_prediction.utils.pdf_annotation import annotate_pdf
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("pdf")
|
|
|
|
args = parser.parse_args()
|
|
|
|
return args
|
|
|
|
|
|
def main(args):
|
|
pipeline = load_pipeline(verbose=True)
|
|
|
|
pdf_path = args.pdf
|
|
|
|
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__":
|
|
args = parse_args()
|
|
main(args)
|