Update app.py

app.py

@@ -2,32 +2,43 @@ import gradio as gr
 import requests
 import json
 import base64
+import os
+
 
 API_URL = "https://image-modal.nebuia.com/extract_id_data"
 
-def send_image_to_api(image_path, json_prompt):
-    # Read the image file and convert to base64
-    with open(image_path, "rb") as image_file:
-        image_base64 = base64.b64encode(image_file.read()).decode('utf-8')
-    
+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": ("image.jpg", base64.b64decode(image_base64), "image/jpeg")
+        "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
@@ -44,22 +55,34 @@ def send_image_to_api(image_path, json_prompt):
         return f"Error sending request: {e}"
 
 # Define the Gradio interface
-def gradio_interface(image_path, json_prompt):
-    if image_path is None:
-        return "Please upload an image."
-    return send_image_to_api(image_path, json_prompt)
+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
-iface = gr.Interface(
-    fn=gradio_interface,
-    inputs=[
-        gr.Image(type="filepath", label="Sube un ID"),
-        gr.Code(label="JSON Prompt", language="json", lines=10, value='{\n    "nombre": ""\n}')
-    ],
-    outputs=gr.Textbox(label="API Response", lines=10),
-    title="NebuIA ID Estructure Extractor",
-    description="Sube un ID, modifica la estructura con los datos que deseas extraer y presiona Submit."
-)
+# 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
-iface.launch()
+demo.launch()