Update app.py

app.py

@@ -1,47 +1,227 @@
-import os
 import gradio as gr
+import torch
 from transformers import Pix2StructProcessor, Pix2StructForConditionalGeneration
 from PIL import Image
+import requests
+import io
+import re
+import pandas as pd
 import json
 
-# Fix threading error
-os.environ["OMP_NUM_THREADS"] = "1"
-
-# Load DePlot
-model_id = "google/deplot"
-processor = Pix2StructProcessor.from_pretrained(model_id)
-model = Pix2StructForConditionalGeneration.from_pretrained(model_id)
-
-def extract_chart(image):
-    # Step 1: Run DePlot
-    inputs = processor(images=image, text="Generate table from chart.", return_tensors="pt")
-    predictions = model.generate(**inputs, max_new_tokens=512)
-    table = processor.decode(predictions[0], skip_special_tokens=True)
-
-    # Step 2: Dummy structured JSON
-    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": []}
-    }
-
-    # Step 3: Merge outputs
-    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)
-
-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."
-)
+# Load the DePlot model and processor
+MODEL_NAME = "google/deplot"
 
+def load_model():
+    """Load the DePlot model and processor"""
+    try:
+        processor = Pix2StructProcessor.from_pretrained(MODEL_NAME)
+        model = Pix2StructForConditionalGeneration.from_pretrained(MODEL_NAME)
+        return processor, model
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        return None, None
+
+processor, model = load_model()
+
+def extract_chart_data(image, question="Generate underlying data table of the figure below:"):
+    """
+    Extract data from chart image using DePlot model
+    
+    Args:
+        image: PIL Image or file path
+        question: Question to ask about the chart
+    
+    Returns:
+        Extracted data as text and structured format
+    """
+    if processor is None or model is None:
+        return "Error: Model not loaded properly", None
+    
+    try:
+        # Ensure image is PIL Image
+        if isinstance(image, str):
+            image = Image.open(image)
+        elif hasattr(image, 'name'):  # Gradio file object
+            image = Image.open(image.name)
+        
+        # Convert to RGB if necessary
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        # Process the image and question
+        inputs = processor(images=image, text=question, return_tensors="pt")
+        
+        # Generate predictions
+        predictions = model.generate(**inputs, max_new_tokens=512)
+        
+        # Decode the output
+        extracted_text = processor.decode(predictions[0], skip_special_tokens=True)
+        
+        # Try to parse the extracted text into structured data
+        structured_data = parse_extracted_data(extracted_text)
+        
+        return extracted_text, structured_data
+    
+    except Exception as e:
+        return f"Error processing image: {str(e)}", None
+
+def parse_extracted_data(text):
+    """
+    Parse the extracted text to create structured data
+    This is a basic parser - you might need to customize based on your needs
+    """
+    try:
+        # Look for table-like patterns
+        lines = text.strip().split('\n')
+        data = []
+        
+        # Try to find header and data rows
+        for line in lines:
+            if '|' in line:  # Table format with pipes
+                row = [cell.strip() for cell in line.split('|') if cell.strip()]
+                if row:
+                    data.append(row)
+            elif '\t' in line:  # Tab-separated
+                row = [cell.strip() for cell in line.split('\t') if cell.strip()]
+                if row:
+                    data.append(row)
+            elif ',' in line and not line.startswith('The'):  # CSV-like
+                row = [cell.strip() for cell in line.split(',') if cell.strip()]
+                if row:
+                    data.append(row)
+        
+        if data:
+            # Create DataFrame
+            if len(data) > 1:
+                df = pd.DataFrame(data[1:], columns=data[0])
+            else:
+                df = pd.DataFrame(data)
+            return df
+        
+        return None
+    
+    except Exception as e:
+        print(f"Error parsing data: {e}")
+        return None
+
+def process_chart(image, custom_question):
+    """
+    Main function to process chart and return results
+    """
+    if image is None:
+        return "Please upload an image", None, None
+    
+    # Use custom question if provided, otherwise use default
+    question = custom_question if custom_question.strip() else "Generate underlying data table of the figure below:"
+    
+    # Extract data
+    raw_output, structured_data = extract_chart_data(image, question)
+    
+    # Prepare outputs
+    if structured_data is not None and not structured_data.empty:
+        # Convert DataFrame to HTML for display
+        table_html = structured_data.to_html(index=False, classes='table table-striped')
+        # Convert DataFrame to CSV string for download
+        csv_output = structured_data.to_csv(index=False)
+    else:
+        table_html = "Could not parse data into structured format"
+        csv_output = None
+    
+    return raw_output, table_html, csv_output
+
+# Create Gradio interface
+def create_interface():
+    with gr.Blocks(title="DePlot Chart Data Extractor", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("""
+        # 📊 DePlot Chart Data Extractor
+        
+        Upload any chart or plot image to extract the underlying data, even without visible data labels!
+        This tool uses Google's DePlot model to understand and extract data from various types of charts.
+        
+        **Supported chart types:** Bar charts, line graphs, scatter plots, pie charts, and more!
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                # Input section
+                image_input = gr.Image(
+                    type="pil",
+                    label="Upload Chart Image",
+                    height=400
+                )
+                
+                custom_question = gr.Textbox(
+                    label="Custom Question (optional)",
+                    placeholder="e.g., 'What are the values for each category?' or leave empty for default",
+                    lines=2
+                )
+                
+                extract_btn = gr.Button("Extract Data", variant="primary", size="lg")
+            
+            with gr.Column(scale=1):
+                # Output section
+                with gr.Tab("Raw Output"):
+                    raw_output = gr.Textbox(
+                        label="Extracted Text",
+                        lines=10,
+                        show_copy_button=True
+                    )
+                
+                with gr.Tab("Structured Data"):
+                    structured_output = gr.HTML(
+                        label="Parsed Data Table"
+                    )
+                
+                # Download section
+                csv_download = gr.File(
+                    label="Download CSV",
+                    visible=False
+                )
+        
+        # Example images
+        gr.Markdown("### 📋 Try these examples:")
+        
+        example_images = [
+            ["examples/bar_chart.png", "Extract the data from this bar chart"],
+            ["examples/line_graph.png", "What are the trend values over time?"],
+            ["examples/pie_chart.png", "Give me the percentage breakdown"]
+        ]
+        
+        # Note: You'll need to add example images to your space
+        
+        # Event handlers
+        def process_and_download(image, question):
+            raw, table, csv_data = process_chart(image, question)
+            
+            if csv_data:
+                # Create temporary CSV file for download
+                csv_file = io.StringIO()
+                csv_file.write(csv_data)
+                csv_file.seek(0)
+                return raw, table, gr.File(value=csv_data, visible=True)
+            else:
+                return raw, table, gr.File(visible=False)
+        
+        extract_btn.click(
+            fn=process_and_download,
+            inputs=[image_input, custom_question],
+            outputs=[raw_output, structured_output, csv_download]
+        )
+        
+        gr.Markdown("""
+        ### 💡 Tips for better results:
+        - Use clear, high-resolution images
+        - Ensure chart elements are visible and not too small
+        - Try different custom questions for specific data you need
+        - Works best with standard chart types (bar, line, scatter, pie)
+        
+        ### 🔧 Model Information:
+        This space uses Google's DePlot model, which is specifically trained to extract data from plots and figures.
+        """)
+    
+    return demo
+
+# Create and launch the interface
 if __name__ == "__main__":
-    demo.launch()
+    demo = create_interface()
+    demo.launch()