image-classification-service/scripts/run_pipeline.py

import argparse
import json
import os
from glob import glob
from operator import truth

from image_prediction.pipeline import load_pipeline
from image_prediction.utils import get_logger
from image_prediction.utils.pdf_annotation import annotate_pdf

logger = get_logger()


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument("input", help="pdf file or directory")
    parser.add_argument("--metadata", help="optional figure detection metadata")
    parser.add_argument("--print", "-p", help="print output to terminal", action="store_true", default=False)
    parser.add_argument("--page_interval", "-i", help="page interval [i, j), min index = 0", nargs=2, type=int)

    args = parser.parse_args()

    return args


def process_pdf(pipeline, pdf_path, metadata=None, page_range=None):
    if metadata:
        with open(metadata) as f:
            metadata = json.load(f)

    with open(pdf_path, "rb") as f:
        logger.info(f"Processing {pdf_path}")
        predictions = list(pipeline(f.read(), page_range=page_range, metadata_per_image=metadata))

    annotate_pdf(
        pdf_path, predictions, os.path.join("/tmp", os.path.basename(pdf_path.replace(".pdf", f"_{truth(metadata)}_annotated.pdf")))
    )

    return predictions


def main(args):
    pipeline = load_pipeline(verbose=True, tolerance=3)

    if os.path.isfile(args.input):
        pdf_paths = [args.input]
    else:
        pdf_paths = glob(os.path.join(args.input, "*.pdf"))
    page_range = range(*args.page_interval) if args.page_interval else None
    metadata = args.metadata if args.metadata else None

    for pdf_path in pdf_paths:
        predictions = process_pdf(pipeline, pdf_path, metadata, page_range=page_range)
        if args.print:
            print(pdf_path)
            print(json.dumps(predictions, indent=2))


if __name__ == "__main__":
    args = parse_args()
    main(args)