Spaces:
Sleeping
Sleeping
File size: 3,728 Bytes
6607e95 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import gradio as gr
import requests
import os
import urllib
import time
document_types = {"Annual Financial Statement": "afs", "Invoice": "i", "Purchase Order": "po", "Bank Statement": "bs"}
API_KEY = os.getenv("API_KEY", None)
SERVER_URL = os.getenv("SERVER_URL", None)
with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.red, secondary_hue=gr.themes.colors.pink)) as demo:
task_id = gr.State()
if SERVER_URL is not None:
file_input = gr.File(label="Upload a file", file_types=['.pdf'])
document_type = gr.Dropdown(label="Select Document Type", choices=list(document_types.keys()))
api_key = gr.Text(label="Enter API key", value=API_KEY if API_KEY else None, placeholder="Enter your API key here", visible=API_KEY is None)
with gr.Row():
clear_button = gr.Button("Clear")
process_button = gr.Button("Submit")
status_output = gr.Markdown(label="Current Status")
file_upload_status = gr.Text(label="File Upload Status", interactive=False)
else:
gr.Markdown("Server URL is not configured. Please set the SERVER_URL environment variable.")
def get_result(task_id, api_key, progress=gr.Progress()):
if task_id is not None:
current_status = ""
total_steps = 1000
progress(0, desc="Starting...")
for i in progress.tqdm(range(total_steps)):
response = requests.get(f"{SERVER_URL}/verify/get_results/{task_id}", headers={"X-Api-Key": api_key})
current_status = response.json()['current_status']
if current_status != "In Progress":
break
time.sleep(1)
return response.json()
@gr.render(inputs=task_id)
def show_results(progress_id):
if progress_id is not None:
view_results = gr.Button("View Results")
view_results.click(fn=get_result, inputs=[task_id, api_key], outputs=gr.Json())
def process_file(file, api_key, document_type, progress=gr.Progress()):
if file is not None and document_type is not None:
params = {"document-type": document_types[document_type]}
url_encoded_file_name = urllib.parse.quote(os.path.basename(file.name))
print(f"filename encoded {url_encoded_file_name}")
response = requests.get(f"{SERVER_URL}/verify/get_upload_url/{url_encoded_file_name}", headers={"X-Api-Key": api_key, "Accept": "application/json"}, params=params)
if response.status_code == 200:
upload_url = response.json()['url']
task_id = response.json()['task_id']
print(f"presigned url: {upload_url}, Task Id: {task_id}")
with open(file.name, 'rb') as f:
file_content = f.read()
file_size = len(file_content)
# Upload the entire file
upload_response = requests.put(upload_url, data=file_content, headers={'Content-Type': 'application/pdf'})
# Update the progress bar
progress(1, desc="Uploading...")
print(f"Uploaded {file_size} bytes")
return task_id, "Upload successful"
return None, "Upload failed"
def clear_fields():
return None, None, None
if SERVER_URL is not None:
process_button.click(fn=process_file, inputs=[file_input, api_key, document_type], outputs=[task_id, file_upload_status])
clear_button.click(fn=clear_fields, inputs=None, outputs=[file_input, document_type, task_id])
demo.launch() |