Custom

app.py

@@ -74,10 +74,15 @@ def prepare_prompt(image_path):
     return messages, text
 
 @spaces.GPU
-def process_image(image):
+def process_image(image, schema_str):
     if image is None:
         return "No image provided."
     
+    try:
+        custom_schema = json.loads(schema_str)
+    except json.JSONDecodeError:
+        return "Invalid JSON schema provided."
+    
     base64_str = encode_image_from_pil(image)
     messages = [
         {
@@ -89,7 +94,7 @@ def process_image(image):
     ]
     text = processor.tokenizer.apply_chat_template(
         messages,
-        template=json.dumps(invoice_schema, indent=4),
+        template=json.dumps(custom_schema, indent=4),
         tokenize=False,
         add_generation_prompt=True
     )
@@ -129,10 +134,18 @@ def process_image(image):
 # Gradio interface
 iface = gr.Interface(
     fn=process_image,
-    inputs=gr.Image(type="pil", label="Upload Invoice Image"),
-    outputs=gr.Textbox(label="Extracted Invoice Data (JSON)"),
+    inputs=[
+        gr.Image(type="pil", label="Upload Invoice Image"),
+        gr.Textbox(
+            label="Custom Schema (JSON)",
+            value=json.dumps(invoice_schema, indent=4),
+            lines=10,
+            placeholder="Enter your custom JSON schema here..."
+        )
+    ],
+    outputs=gr.Textbox(label="Extracted Data (JSON)"),
     title="Invoice Parser with NuExtract",
-    description="Upload an invoice image to extract structured data using AI."
+    description="Upload an invoice image and provide a custom JSON schema to extract structured data using AI."
 )
 
 iface.launch()