import os
from sdsvkvu import load_engine, process_img, process_pdf, process_dir
from sdsvkvu.modules.run_ocr import load_ocr_engine
os.environ["CUDA_VISIBLE_DEVICES"]="1"
# os.environ["NLTK_DATA"]="/mnt/ssd1T/tuanlv/02-KVU/sdsvkvu/nltk_data"

if __name__ == "__main__":
    # ocr_engine = load_ocr_engine({"device": "cuda:0"})
    kwargs = {"device": "cuda:0", "ocr_engine": None}
    img_dir = "/mnt/hdd4T/OCR/tuanlv/02-KVU/sdsvkvu/visualize/test_manulife"
    save_dir = "/mnt/hdd4T/OCR/tuanlv/02-KVU/sdsvkvu/visualize/test_manulife"
    engine = load_engine(kwargs)
    # option: "vat" for vat invoice outputs, "sbt": sbt invoice outputs, else for raw outputs
    # outputs = process_img(img_dir, save_dir, engine, export_all=False, option="vat") 
    # outputs = process_pdf(img_dir, save_dir, engine, export_all=True, option="vat")
    process_dir(img_dir, save_dir, engine, export_all=True, option="manulife")
    # process_dir(img_dir, save_dir, engine, export_all=True, option="")