import base64
import json
from io import BytesIO
import pandas as pd
from PIL import Image

import gradio as gr
import requests


def ocr(image):

    image = Image.open(image)
    img_buffer = BytesIO()
    image.save(img_buffer, format=image.format)
    byte_data = img_buffer.getvalue()
    base64_bytes = base64.b64encode(byte_data)  # bytes
    base64_str = base64_bytes.decode()
    url = "https://www.modelscope.cn/api/v1/studio/damo/ofa_ocr_pipeline/gradio/api/predict/"
    payload = json.dumps({
        "data": [f"data:image/jpeg;base64,{base64_str}"],
        "dataType": ["image"]
    })
    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload)
    jobj = json.loads(response.text)
    out_img_base64 = jobj['Data']['data'][0].replace('data:image/png;base64,','')
    out_img = Image.open(BytesIO(base64.urlsafe_b64decode(out_img_base64)))
    ocr_result = jobj['Data']['data'][1]['data']

    result = pd.DataFrame(ocr_result, columns=['Box ID', 'Text'])

    return out_img, result


title = "Chinese OCR"
description = "Gradio Demo for Chinese OCR based on OFA-Base. "\
              "Upload your own image or click any one of the examples, and click " \
              "\"Submit\" and then wait for the generated OCR result." \
              "\n中文OCR体验区。欢迎上传图片，静待检测文字返回~"
article = "<p style='text-align: center'><a href='https://github.com/OFA-Sys/OFA' target='_blank'>OFA Github " \
          "Repo</a></p> "
examples = [['shupai.png'], ['chinese.jpg'], ['gaidao.jpeg'],
             ['qiaodaima.png'], ['xsd.jpg']]
io = gr.Interface(fn=ocr, inputs=gr.inputs.Image(type='filepath', label='Image'),
                  outputs=[gr.outputs.Image(type='pil', label='Image'),
                           gr.outputs.Dataframe(headers=['Box ID', 'Text'], type='pandas', label='OCR Results')],
                  title=title, description=description, article=article)
io.launch()