Spaces:
Runtime error
Runtime error
| import streamlit as st | |
| import google.generativeai as genai | |
| import os | |
| import pandas as pd | |
| import json | |
| import re | |
| from PIL import Image | |
| # Configure Gemini using environment variable | |
| def configure_gemini(): | |
| api_key = os.getenv("GEMINI_API_KEY") | |
| if not api_key: | |
| raise ValueError("GEMINI_API_KEY environment variable not set") | |
| genai.configure(api_key=api_key) | |
| return genai.GenerativeModel('gemini-1.5-flash-latest') | |
| # Process image with Gemini | |
| def process_image(model, image): | |
| prompt = """Analyze this restaurant menu image and extract structured data in this exact JSON format. Image could have dishes in multiple columns and rows. Verify and make sure none is missed. For example this image having around 4 to 5 categories and in total around 40 to 50 dishes in total. The dishes could have name, description, price, likes percentage, rating percentage. The extracted data should be in JSON format as shown below: | |
| [ | |
| { | |
| "category": "Category Name", | |
| "items": [ | |
| { | |
| "name": "Dish Name", | |
| "description": "Dish description (optional)", | |
| "price": "Price" | |
| } | |
| ] | |
| } | |
| ] | |
| Return ONLY the JSON, no other text. If no menu found, return empty JSON.""" | |
| response = model.generate_content([prompt, image]) | |
| return response.text | |
| def extract_json(response_text): | |
| # Debugging: Log the response text to see its format before processing | |
| #st.write("Received response_text:", response_text) | |
| response_text = response_text.strip('`') | |
| response_text = response_text.strip('`') | |
| # Remove any non-JSON prefixes | |
| prefix = 'json' | |
| if response_text.startswith(prefix): | |
| response_text = response_text[len(prefix):].strip() | |
| # Attempt to parse JSON | |
| try: | |
| json_data = json.loads(response_text) | |
| #st.write("Parsed JSON data:", json_data) | |
| return json_data | |
| except json.JSONDecodeError as e: | |
| st.error(f"Failed to decode JSON from the response. Error: {str(e)}") | |
| return [] | |
| def main(): | |
| st.title("🍽️ Menu Digitizer") | |
| st.write("Upload menu images to generate consolidated CSV") | |
| try: | |
| model = configure_gemini() | |
| except ValueError as e: | |
| st.error(f"Configuration error: {str(e)}") | |
| st.info("Please set the GEMINI_API_KEY environment variable") | |
| return | |
| uploaded_files = st.file_uploader("Upload menu images", type=["jpg", "jpeg", "png"], accept_multiple_files=True) | |
| all_items = [] # Initialize outside the loop to accumulate items from all files | |
| if uploaded_files: | |
| with st.spinner("Processing menus..."): | |
| for uploaded_file in uploaded_files: | |
| image = Image.open(uploaded_file) | |
| try: | |
| response = process_image(model, image) | |
| #st.write(f"Raw response for {uploaded_file.name}:") | |
| #st.code(response, language='json') # Display raw response | |
| menu_data = extract_json(response) | |
| #st.write(f"menu_data response for {menu_data}:") | |
| #st.code(menu_data) # Display raw response | |
| for category in menu_data: | |
| category_name = category.get('category', 'NA') # Use 'NA' if category is missing | |
| for item in category['items']: | |
| item['category'] = category_name | |
| all_items.append(item) | |
| except Exception as e: | |
| st.error(f"Error processing {uploaded_file.name}: {str(e)}") | |
| continue | |
| if all_items: | |
| df = pd.DataFrame(all_items) | |
| st.subheader("Extracted Menu Items") | |
| st.dataframe(df) | |
| csv = df.to_csv(index=False).encode('utf-8') | |
| st.download_button("Download CSV", csv, "consolidated_menu.csv", "text/csv") | |
| else: | |
| st.warning("No menu items found in the uploaded images") | |
| if __name__ == "__main__": | |
| main() |