Update app.py

app.py

@@ -5,7 +5,6 @@ import json
 from pathlib import Path
 import gradio as gr
 from openai import OpenAI
-from PIL import Image
 
 API_KEY = "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"
 MODEL = "gpt-5.1"
@@ -17,6 +16,7 @@ def upload_pdf(path):
     return client.files.create(file=open(path, "rb"), purpose="assistants").id
 
 
+# ---------------- Prompt (unchanged) ----------------
 def prompt():
     return (
         "Extract structured JSON from the attached logistics document. Return ONLY valid JSON.\n"
@@ -77,81 +77,62 @@ def prompt():
     )
 
 
-def extract_image(path):
-    """Process image via filepath"""
-    img_bytes = Path(path).read_bytes()
-    ext = Path(path).suffix.replace(".", "").lower()
 
-    b64 = base64.b64encode(img_bytes).decode()
+# ---------------- Extraction ----------------
+def extract(path):
+    suffix = Path(path).suffix.lower()
 
-    content = [
-        {"type": "text", "text": prompt()},
-        {
-            "type": "image_url",
-            "image_url": {"url": f"data:image/{ext};base64,{b64}"}
-        }
-    ]
+    if suffix == ".pdf":
+        fid = upload_pdf(path)
+        content = [
+            {"type": "text", "text": prompt()},
+            {"type": "file", "file": {"file_id": fid}}
+        ]
+    else:
+        b64 = base64.b64encode(Path(path).read_bytes()).decode()
+        ext = suffix[1:]
+        content = [
+            {"type": "text", "text": prompt()},
+            {"type": "image_url", "image_url": {"url": f"data:image/{ext};base64,{b64}"}}
+        ]
 
     r = client.chat.completions.create(
         model=MODEL,
         messages=[{"role": "user", "content": content}]
     )
 
-    t = r.choices[0].message.content
-    return t[t.find("{"): t.rfind("}") + 1]
+    text = r.choices[0].message.content
+    return text[text.find("{"): text.rfind("}") + 1]
 
 
-def extract_pdf(file):
-    """Process PDF"""
-    path = Path(file.name)
-    fid = upload_pdf(path)
-
-    content = [
-        {"type": "text", "text": prompt()},
-        {"type": "file", "file": {"file_id": fid}}
-    ]
-
-    r = client.chat.completions.create(
-        model=MODEL,
-        messages=[{"role": "user", "content": content}]
-    )
-
-    t = r.choices[0].message.content
-    return t[t.find("{"): t.rfind("}") + 1]
-
-
-def process(image_path, pdf_file):
-    if image_path:
-        return extract_image(image_path)
-
-    if pdf_file:
-        return extract_pdf(pdf_file)
-
+def ui(image_input, pdf_input):
+    if image_input:
+        return extract(image_input)
+    if pdf_input:
+        return extract(pdf_input.name)
     return "{}"
 
 
+# ---------------- UI ----------------
+
 with gr.Blocks() as demo:
     gr.Markdown("# **Logistics OCR Data Extractor (GPT-5.1)**")
 
     with gr.Row():
-        image_input = gr.Image(type="filepath", label="Upload Image")
-        pdf_input = gr.File(type="file", label="Upload PDF")
+        img = gr.Image(label="Upload Image", type="filepath")
+        pdf = gr.File(label="Upload PDF", file_types=["pdf"])
 
-    output = gr.JSON(label="Extracted JSON")
-    submit = gr.Button("Submit")
+    out = gr.JSON(label="Extracted JSON")
+    btn = gr.Button("Submit")
 
-    submit.click(
-        fn=process,
-        inputs=[image_input, pdf_input],
-        outputs=output
-    )
+    btn.click(fn=ui, inputs=[img, pdf], outputs=out)
 
     gr.Examples(
         examples=[
             ["IMG_0001.jpg", None],
             ["IMG_0002.jpg", None]
         ],
-        inputs=[image_input, pdf_input],
+        inputs=[img, pdf],
         label="Sample Images"
     )