Update app.py

app.py

@@ -3,19 +3,20 @@ import pandas as pd
 import matplotlib.pyplot as plt
 import io
 import google.generativeai as genai
+from PIL import Image
 
 def process_file(api_key, file, instructions):
-    # Set up Gemini API
+    # Configure Gemini API
     genai.configure(api_key=api_key)
     model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')
 
-    # Read the file
+    # Read uploaded file
     if file.name.endswith('.csv'):
         df = pd.read_csv(file.name)
     else:
         df = pd.read_excel(file.name)
 
-    # Analyze data and get visualization suggestions from Gemini
+    # Generate analysis prompt
     data_description = df.describe().to_string()
     columns_info = "\n".join([f"{col}: {df[col].dtype}" for col in df.columns])
     prompt = f"""
@@ -28,72 +29,62 @@ def process_file(api_key, file, instructions):
     
     User instructions: {instructions if instructions else 'No specific instructions provided.'}
     
-    Generate Python code for 3 different visualizations using matplotlib. Each visualization should be unique and provide insights into the data. Do not include any explanations or descriptions, only the Python code for each visualization. Ensure the code is not indented.
-
-    Format your response as:
+    Generate Python code for 3 different matplotlib visualizations. Only provide executable code without explanations.
+    Format as:
     # Visualization 1
-    # Your code here
-
+    [code]
     # Visualization 2
-    # Your code here
-
+    [code]
     # Visualization 3
-    # Your code here
+    [code]
     """
 
+    # Get Gemini response
     response = model.generate_content(prompt)
     code = response.text.strip()
 
-    print("Generated code:")
-    print(code)
-
+    # Process visualizations
     visualizations = []
     for i, viz_code in enumerate(code.split("# Visualization")[1:4], 1):
         plt.figure(figsize=(10, 6))
         try:
-            # Remove any leading spaces from each line
+            # Clean and execute generated code
             cleaned_code = '\n'.join(line.strip() for line in viz_code.split('\n') if line.strip())
-            print(f"\nVisualization {i} code:")
-            print(cleaned_code)
-            
             exec(cleaned_code, {'df': df, 'plt': plt})
             plt.title(f"Visualization {i}")
             
-            # Save the plot to a BytesIO object
+            # Convert plot to Gradio-compatible format
             buf = io.BytesIO()
-            plt.savefig(buf, format='png')
-            buf.seek(0)
+            plt.savefig(buf, format='png', bbox_inches='tight')
             plt.close()
-            
-            visualizations.append(buf)
+            buf.seek(0)
+            visualizations.append(Image.open(buf))
         except Exception as e:
-            print(f"Error in visualization {i}: {str(e)}")
+            print(f"Visualization {i} error: {str(e)}")
             visualizations.append(None)
 
-    # Ensure we always return 3 visualizations
-    while len(visualizations) < 3:
-        visualizations.append(None)
-
-    return visualizations
+    # Ensure 3 outputs for Gradio
+    return visualizations + [None]*(3-len(visualizations))
 
 # Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Data Visualization with Gemini")
-    api_key = gr.Textbox(label="Enter Gemini API Key", type="password")
-    file = gr.File(label="Upload Excel or CSV file")
-    instructions = gr.Textbox(label="Optional visualization instructions")
-    submit = gr.Button("Generate Visualizations")
+    gr.Markdown("""# AI-Powered Data Visualizer
+    Upload your data file and get automatic visualizations powered by Gemini""")
+    
+    with gr.Row():
+        api_key = gr.Textbox(label="Gemini API Key", type="password")
+        instructions = gr.Textbox(label="Custom Instructions (optional)")
+    
+    file = gr.File(label="Dataset (CSV/Excel)", file_types=[".csv", ".xlsx"])
+    submit = gr.Button("Generate Insights", variant="primary")
     
     with gr.Row():
-        output1 = gr.Image(label="Visualization 1")
-        output2 = gr.Image(label="Visualization 2")
-        output3 = gr.Image(label="Visualization 3")
+        outputs = [gr.Image(label=f"Visualization {i+1}") for i in range(3)]
 
     submit.click(
-        fn=process_file,
+        process_file,
         inputs=[api_key, file, instructions],
-        outputs=[output1, output2, output3],
-        show_progress=True,
+        outputs=outputs
     )
 
 demo.launch()