import io
import gradio as gr
from google.oauth2 import service_account
from googleapiclient.discovery import build
from apiclient.http import MediaFileUpload, MediaIoBaseDownload

service = build('drive', 'v3', credentials=service_account.Credentials.from_service_account_file('hf-ocr-33c0cb785de5.json'))

def infer(im):
  im.save('converted.png')
  id = service.files().create(body={'mimeType': 'application/vnd.google-apps.document'}, media_body=MediaFileUpload('converted.png')).execute()['id']
  with io.BytesIO() as f:
    downloader = MediaIoBaseDownload(f, service.files().export_media(fileId=id, mimeType='text/plain'))
    while downloader.next_chunk()[1] is False:
      pass
    service.files().delete(fileId=id).execute()
    return f.getvalue().decode('UTF-8').replace('\ufeff________________\r\n\r\n', '')

iface = gr.Interface(
    fn=infer,
    title="Google Drive OCR",
    description="When you convert an image to a Google doc, Drive uses Optical Character Recognition (OCR) to convert the image to text.",
    inputs=[gr.inputs.Image(label='image', type='pil')],
    outputs='text',
    examples=['testocr.png', 'receipt.webp', '20131216170659.jpg'],
    article="<a href=\"https://developers.google.com/drive/api/guides/manage-uploads#import_to_google_docs_types\">Docs</a>",
).launch()