import streamlit as st
from transformers import pipeline
from PIL import Image
import os

# Load the image classification pipeline
@st.cache_resource
def load_image_classification_pipeline():
    """
    Load the image classification pipeline using a pretrained model.
    """
    return pipeline("image-classification", model="Shresthadev403/food-image-classification")

pipe_classification = load_image_classification_pipeline()

# Load the GPT-Neo model for ingredient generation
@st.cache_resource
def load_llama_pipeline():
    """
    Load the GPT-Neo model for ingredient generation.
    """
    return pipeline("text-generation", model="EleutherAI/gpt-neo-1.3B")

pipe_llama = load_llama_pipeline()

def get_ingredients_llama(food_name):
    """
    Generate a list of ingredients for the given food item using GPT-Neo.
    Returns a clean, comma-separated list of ingredients.
    """
    prompt = (
        f"List only the main ingredients typically used to prepare {food_name}. "
        "Respond with just the ingredients in a comma-separated list without any extra text or explanations."
    )
    try:
        response = pipe_llama(prompt, max_length=50, num_return_sequences=1)
        generated_text = response[0]["generated_text"].strip()
        
        # Process the response to ensure it's a clean, comma-separated list
        ingredients = generated_text.split(":")[-1].strip()  # Handle cases like "Ingredients: ..."
        ingredients = ingredients.replace(".", "").strip()  # Remove periods and extra spaces
        return ingredients
    except Exception as e:
        return f"Error generating ingredients: {e}"
# Streamlit app setup
st.title("Food Image Recognition with Ingredients")

# Add banner image
st.image("IR_IMAGE.png", caption="Food Recognition Model", use_column_width=True)

# Sidebar for model information
st.sidebar.title("Model Information")
st.sidebar.write("**Image Classification Model**: Shresthadev403/food-image-classification")
st.sidebar.write("**LLM for Ingredients**: EleutherAI/gpt-neo-1.3B")

# Upload image
uploaded_file = st.file_uploader("Choose a food image...", type=["jpg", "png", "jpeg"])

if uploaded_file is not None:
    # Display the uploaded image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)
    st.write("Classifying...")

    # Make predictions
    predictions = pipe_classification(image)

    # Display only the top prediction
    top_food = predictions[0]['label']
    st.header(f"Food: {top_food}")

    # Generate and display ingredients for the top prediction
    st.subheader("Ingredients")
    try:
        ingredients = get_ingredients_llama(top_food)
        st.write(ingredients)
    except Exception as e:
        st.error(f"Error generating ingredients: {e}")

# Footer
st.sidebar.markdown("Created with ❤️ using Streamlit and Hugging Face.")