cv-analysis-service/scripts/client_mock.py
2022-04-26 16:01:57 +02:00

62 lines
2.1 KiB
Python

# python client_mock.py --pdf_path=/home/iriley/Documents/pdfs/unscanned/06.pdf --operations=table-parsing
import argparse
import json
from multiprocessing.sharedctypes import Value
import requests
from cv_analysis.utils.preprocessing import open_pdf
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--pdf_path", required=True, help="path to PDF file")
parser.add_argument(
"--first_page", type=int, required=False, default=0, help="page number from which to start (starts at 0)"
)
parser.add_argument(
"--last_page",
type=int,
required=False,
default=None,
help="page number at which to stop (non-inclusive); specify None to go to the end",
)
parser.add_argument(
"--operations",
type=str,
required=False,
help="Comma-separated list of operations, any of the following: \ntable-parsing\nredaction-detection\
\nfigure-detection\nlayout-detection",
default="table-parsing",
)
args = parser.parse_args()
return args
def main(args):
operations = args.operations.split(",")
for operation in operations:
print("****************************")
print(f"{' '+operation+' ':^27}")
print("****************************")
if operation == "table-parsing":
response = requests.post("http://127.0.0.1:5000/tables", data=open(args.pdf_path, "rb"))
elif operation == "redaction-detection":
response = requests.post("http://127.0.0.1:5000/redactions", data=open(args.pdf_path, "rb"))
elif operation == "figure-detection":
response = requests.post("http://127.0.0.1:5000/figures", data=open(args.pdf_path, "rb"))
elif operation == "layout-parsing":
response = requests.post("http://127.0.0.1:5000/layout", data=open(args.pdf_path, "rb"))
else:
raise ValueError(f"{args.operation} is not a valid value.")
response.raise_for_status()
predictions = response.json()
print(json.dumps(predictions, indent=2))
if __name__ == "__main__":
args = parse_args()
main(args)