import gradio as gr
# from fastai.vision.all import *
import requests
from PIL import Image
import os

# learn = load_learner('bedroom_or_kitchen.pkl', 'rb')

# categories = ("Bedroom", "Kitchen")

# def classify_image(img):
#     pred, idx, probs = learn.predict(img)
#     return dict(zip(categories, map(float, probs)))

def get_ocr_prediction(img):
    url = 'https://app.nanonets.com/api/v2/OCR/Model/50440c27-8fb2-4644-a156-75ce1a3c7586/LabelFile/?async=false'
    
    img = Image.fromarray(img, 'RGB')
    img.save('predict.jpg')

    data = {'file': open('predict.jpg', 'rb')}

    response = requests.post(url, auth=requests.auth.HTTPBasicAuth('Dpe3_Fz3XZoYtzMN8nUtsmcDh5t4QYnG', ''), files=data)

    buyer_name = [a for a in response.json()["result"][0]["prediction"] if a["label"]=="buyer_name"][0]["ocr_text"]
    buyer_address = [a for a in response.json()["result"][0]["prediction"] if a["label"]=="buyer_address"][0]["ocr_text"]
    text = buyer_name+";"+buyer_address
    
    os.remove('predict.jpg')
    
    return text


image = gr.inputs.Image()
text = gr.outputs.Textbox()

examples = ['teste.jpg']

intf = gr.Interface(fn=get_ocr_prediction, inputs=image, outputs=text, examples=examples)

intf.launch()