AuditEdge's picture
initial commit
81e13bb
raw
history blame contribute delete
938 Bytes
import argparse
from asyncio.log import logger
from Layoutlm_inference.ocr import prepare_batch_for_inference
from Layoutlm_inference.inference_handler import handle
import logging
import os
if __name__ == "__main__":
# try:
parser = argparse.ArgumentParser()
parser.add_argument("--model_path", type=str, required=True)
parser.add_argument("--image_path", type=str, required=True) # single image path
args = parser.parse_args()
# Expecting a single image file
image_path = args.image_path
# Ensure the file exists before processing
if not os.path.isfile(image_path):
raise FileNotFoundError(f"The provided image path does not exist: {image_path}")
# Prepare batch for a single image
inference_batch = prepare_batch_for_inference([image_path]) # pass as a list
context = {"model_dir": args.model_path}
# Handle the inference
handle(inference_batch, context)