import tempfile
import os

import paddlehub as hub
from PIL import Image

#https://github.com/PaddlePaddle/PaddleHub/tree/b402610a6f0b382a978e82473b541ea1fc6cf09a/modules/image/text_recognition
#pp_ocrv3 = hub.Module(name="ch_pp-ocrv3")
module_name = "ch_pp-ocrv3" #"chinese_ocr_db_crnn_mobile"
pp_ocrv3 = hub.Module(name=module_name)

def inference_img(img):
    with tempfile.TemporaryDirectory() as tempdir_name:
        if module_name in ["ch_pp-ocrv3", "chinese_ocr_db_crnn_mobile"]:
            pp_ocrv3.recognize_text(images=[img], use_gpu=False, output_dir=tempdir_name, visualization=True)
        elif module_name == "multi_languages_ocr_db_crnn":
            pp_ocrv3.recognize_text(images=[img], output_dir=tempdir_name, visualization=True)
        result_names = os.listdir(tempdir_name)
        result_image = Image.open(os.path.join(tempdir_name, result_names[0]))
        return result_image
    
def inference_json(img):
        if module_name in ["ch_pp-ocrv3", "chinese_ocr_db_crnn_mobile"]:
            results = pp_ocrv3.recognize_text(images=[img], use_gpu=False, visualization=False)
        elif module_name == "multi_languages_ocr_db_crnn":
            results = pp_ocrv3.recognize_text(images=[img], visualization=False)
        return results