19 lines
991 B
Python
19 lines
991 B
Python
|
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="")
|
||
|
|
|