import gradio as gr
from craft_hw_ocr import OCR

trocr = OCR.load_TrOCRmodel()

def do_ocr(inp):
  img, results = OCR.craft_detection(inp)
  bboxes, text = OCR.text_recoginition(img, results, trocr[0], trocr[1])
  return OCR.visualize(img, results), text
  
inputs = gr.inputs.Image()
o1 = gr.outputs.Image()
o2 = gr.outputs.Textbox()

title = "CRAFT-OCR"
description = "OCR of both handwriting and printed text using CRAFT Text detector and TrOCR recoginition"
examples=[['example_1.png'],['example_2.jpg']]

article = "<p style='text-align: center'><a href='https://github.com/Vishnunkumar/craft_hw_ocr' target='_blank'>craft_hw_ocr</a></p><p style='text-align: center'> <p style='text-align: center'><a href='https://github.com/fcakyon/craft-text-detector' target='_blank'>craft-text-detector</a></p><p style='text-align: center'>"
gr.Interface(fn=do_ocr, inputs=inputs, outputs=[o1, o2], title=title, description=description, article=article, examples=examples, enable_queue=True).launch()