Update app.py

app.py

@@ -1,49 +1,42 @@
 import gradio as gr
-from transformers import Pix2StructProcessor, Pix2StructForConditionalGeneration
+from transformers import AutoProcessor, AutoModelForSeq2SeqLM
 from PIL import Image
-import torch
+import json
 
-# Load model & processor
-processor = Pix2StructProcessor.from_pretrained("google/deplot")
-model = Pix2StructForConditionalGeneration.from_pretrained("google/deplot")
+# Load DePlot (vision-to-table)
+deplot_model = AutoModelForSeq2SeqLM.from_pretrained("google/deplot")
+deplot_processor = AutoProcessor.from_pretrained("google/deplot")
 
-# Move model to GPU if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+def extract_chart(image):
+    # Step 1: Run DePlot
+    inputs = deplot_processor(images=image, text="Generate table from chart.", return_tensors="pt")
+    outputs = deplot_model.generate(**inputs, max_new_tokens=512)
+    table = deplot_processor.decode(outputs[0], skip_special_tokens=True)
 
-# def extract_chart_data(img: Image.Image):
-#     inputs = processor(images=img, return_tensors="pt").to(device)
-#     outputs = model.generate(**inputs)
-#     result = processor.decode(outputs[0], skip_special_tokens=True)
-#     return result
+    # Step 2: Create dummy structured JSON (later we plug in detection+OCR)
+    structured_json = {
+        "metadata": {"title": "Demo Chart", "chart_type": "bar", "confidence": 0.5},
+        "axes": {"x_axis": {"label": "X", "ticks": []}, "y_axis": {"label": "Y", "ticks": []}},
+        "series": [],
+        "legend": {"entries": []}
+    }
 
-def extract_chart_data(img: Image.Image, header_text:str):
-    # Provide a header prompt for VQA
-    # header_text = "Convert this bar chart into a table of quarter and cost values"
-    print("Header text:", header_text)
-    
-    inputs = processor(images=img, text=header_text, return_tensors="pt").to(device)
-    outputs = model.generate(**inputs)
-    all_results = [processor.decode(seq, skip_special_tokens=True) for seq in outputs]
-    print(all_results)
-    result = processor.decode(outputs[0], skip_special_tokens=True)
-    return result
+    # Step 3: Merge (for now, just attach DePlot’s output)
+    merged_output = {
+        "structured_json": structured_json,
+        "deplot_table": table,
+        "fusion_notes": "Fusion layer not implemented yet, just showing both outputs."
+    }
 
+    return json.dumps(merged_output, indent=2)
 
-# Gradio interface
-iface = gr.Interface(
-    fn=extract_chart_data,
-    # inputs=gr.Image(type="pil"),
-    inputs=[
-        gr.Image(type="pil", label="Upload Chart Image"),  # Image input
-        gr.Textbox(
-            label="Instruction / Header Text",
-            placeholder="Enter instructions like 'Convert this bar chart into a table of quarter and cost values'"
-        )  # Text input
-    ],
-    outputs="text",
-    title="DePlot Chart Data Extractor 1",
-    description="Upload a chart image and get its extracted data/description."
+demo = gr.Interface(
+    fn=extract_chart,
+    inputs=gr.Image(type="pil"),
+    outputs="json",
+    title="Chart-to-JSON Extractor (Prototype)",
+    description="Uploads a chart, extracts structured JSON (dummy) and DePlot table side-by-side."
 )
 
-iface.launch()
+if __name__ == "__main__":
+    demo.launch()