cv-analysis-service/scripts/annotate.py

import argparse

from vidocp.table_parsing import annotate_tables_in_pdf
from vidocp.redaction_detection import annotate_redactions_in_pdf
from vidocp.layout_parsing import annotate_layout_in_pdf
from vidocp.figure_detection import detect_figures_in_pdf


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("pdf_path")
    parser.add_argument("page_index", type=int)
    parser.add_argument("--type", choices=["table", "redaction", "layout", "figure"])

    args = parser.parse_args()

    return args


if __name__ == "__main__":
    args = parse_args()
    if args.type == "table":
        annotate_tables_in_pdf(args.pdf_path, page_index=args.page_index)
    elif args.type == "redaction":
        annotate_redactions_in_pdf(args.pdf_path, page_index=args.page_index)
    elif args.type == "layout":
        annotate_layout_in_pdf(args.pdf_path, page_index=args.page_index)
    elif args.type == "figure":
        detect_figures_in_pdf(args.pdf_path, page_index=args.page_index)