import requests
import gradio as gr
import pandas as pd
from PIL import Image, ImageDraw

def infer(im):
  im.save('converted.png')
  url = 'https://ajax.thehive.ai/api/demo/classify?endpoint=text_recognition'
  files = {
      'image': ('converted.png', open('converted.png', 'rb'), 'image/png'),
      'model_type': (None, 'detection'),
      'media_type': (None, 'photo'),
  }
  res = requests.post(url, files=files).json()
  img = im.convert('RGB')

  words = []
  draw = ImageDraw.Draw(img,'RGBA')
  for output in res['response']['output']:
    for poly in output['bounding_poly']:
      words += [c['class'] for c in poly['classes']]
      draw.rectangle((poly['dimensions']['left']-2,poly['dimensions']['top']-2,poly['dimensions']['right']+2,poly['dimensions']['bottom']+2), outline=(0,255,0,255), fill=(0,255,0,50),width=2)

  img.save('result.png')
  return 'result.png', '\n'.join([o['block_text'] for o in res['response']['output']]), pd.DataFrame(words)

iface = gr.Interface(
    fn=infer,
    title="Hive OCR",
    description="Demo for Hive OCR.Transcribe and analyze media depicting typed, written, or graphic text",
    inputs=[gr.inputs.Image(label='image', type='pil')],
    outputs=['image', 'text', gr.outputs.Dataframe(headers=['word'])],
    examples=['testocr.png', 'receipt.webp', '20131216170659.jpg'],
    article="<a href=\"https://thehive.ai/hive-ocr-solutions\">Hive OCR</a>",
).launch()