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

# Set Streamlit page config
st.set_page_config(page_title="Food Image Classifier", layout="centered")

# Load the model
@st.cache_resource
def load_model():
    st.text("Loading model...")
    model = pipeline("image-classification", model="Xenova/mobilenet_v2_1.0_224")
    st.text("Model loaded successfully!")
    return model

classifier = load_model()

# Streamlit UI
st.title("🍕🥖 Food Image Classifier")
st.write("Upload an image of **roti, pizza, naan, or tofu** to classify.")

uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Convert file to PIL image
    image = Image.open(uploaded_file)

    # Display the uploaded image
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Classify the image
    with st.spinner("Classifying..."):
        results = classifier(image)

    # Display results
    if results:
        label = results[0]['label']
        confidence = results[0]['score'] * 100  # Convert to percentage

        st.success(f"**Prediction:** {label}")
        st.info(f"**Confidence:** {confidence:.2f}%")

        # Option to classify another image
        st.button("Classify Another Image", on_click=lambda: st.experimental_rerun())

# Footer
st.markdown("---")
st.markdown("Made by **Muneeb Sahaf** | Final Year Project 2025")