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)