IDExtractor / app.py
xellDart13's picture
Update app.py
1b3d8d5 verified
import gradio as gr
import requests
import json
import base64
import os
API_URL = "https://image-modal.nebuia.com/extract_id_data"
def send_file_to_api(file_path, json_prompt):
# Read the file and convert to base64
with open(file_path, "rb") as file:
file_base64 = base64.b64encode(file.read()).decode('utf-8')
# Parse the JSON prompt
try:
json_prompt_dict = json.loads(json_prompt)
except json.JSONDecodeError:
return "Error: Invalid JSON prompt"
# Determine file type
file_extension = os.path.splitext(file_path)[1].lower()
if file_extension in ['.jpg', '.jpeg', '.png']:
mime_type = "image/jpeg"
elif file_extension == '.pdf':
mime_type = "application/pdf"
else:
return "Error: Unsupported file type"
# Create a dictionary with the file data and JSON prompt
files = {
"file": (os.path.basename(file_path), base64.b64decode(file_base64), mime_type)
}
data = {
"json_prompt": json.dumps(json_prompt_dict)
}
try:
# Send POST request to the API
response = requests.post(API_URL, files=files, data=data)
# Check if the request was successful
if response.status_code == 200:
# Parse the JSON response
result = response.json()
if result.get("success"):
return json.dumps(result["data"], indent=2)
else:
error_message = f"Error in processing:\n{result.get('error', 'Unknown error')}\n"
error_message += f"Raw output: {result.get('raw_output', 'No raw output available')}"
return error_message
else:
return f"Error: Received status code {response.status_code}\n{response.text}"
except requests.RequestException as e:
return f"Error sending request: {e}"
# Define the Gradio interface
def gradio_interface(file, json_prompt):
if file is None:
return "Please upload an image or PDF file."
return send_file_to_api(file.name, json_prompt)
# Custom color for the theme
custom_purple = "#7f56d9"
# Create the Gradio interface using Blocks
with gr.Blocks(theme=gr.themes.Default(primary_hue="purple", secondary_hue="purple")) as demo:
gr.Markdown(
"""
<div style="text-align: center;">
<img src="https://copilot.nebuia.com/images/logo_white.png" style="max-height: 400px; max-width: 120px; object-fit: contain;">
</div>
## NebuIA ID Structure Extractor
Sube un ID (imagen o PDF), modifica la estructura con los datos que deseas extraer y presiona Submit.
"""
)
with gr.Row():
file_input = gr.File(label="Sube un ID (imagen o PDF)", type="filepath")
json_input = gr.Code(label="JSON Prompt", language="json", lines=10, value='{\n "nombre": ""\n}')
output = gr.Textbox(label="API Response", lines=10)
submit_btn = gr.Button("Submit")
submit_btn.click(fn=gradio_interface, inputs=[file_input, json_input], outputs=output)
# Launch the interface
demo.launch()