import mmocr
import gradio as gr
import os
from huggingface_hub import hf_hub_url
from mmocr.utils.ocr import MMOCR


# The default config path is based the current woking directory.
config_dir = os.path.join(os.path.dirname(
    os.path.dirname(mmocr.__file__)), 'configs/')

# TODO: Put the model on HF hub.
ocr = MMOCR(config_dir=config_dir)
results = ocr.readtext('demo_text_ocr.jpg',
                       output='.', print_result=True, imshow=True)

def infer(image):
  # TODO: Also display bounding boxes
  return ocr.readtext(image, output='.', print_result=True, imshow=False)

# TODO: Add a dropbox for model selection
iface = gr.Interface(fn=infer, inputs="image", outputs="json")
iface.launch()