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

# 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="munnae/bc220")
    model = pipeline("image-classification", model="skylord/swin-finetuned-food101")
    st.text("Model loaded successfully!")
    return model

classifier = load_model()

# Streamlit UI
st.title("Virtual University FYP: Food Image Classifier")
st.write("Upload an image of **Tacos, Pizza, Cannoli, or Miso Soup** 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)

    # Use the filename as a label hint
    filename_hint = uploaded_file.name.lower()

    # List of possible labels in your dataset
    dataset_labels = ["sushi", "sashimi", "ramen", "miso_soup", "takoyaki", "dumplings", "fried_rice", "hot and sour soup", 
                      "peking duck", "spring rolls", "pizza", "lasagna", "ravioli", "spaghetti bolognese", "tiramisu",
                      "croque madame", "creme brulee", "foie_gras", "escargots", "chocolate_mousse",
                      "tacos", "nachos", "guacamole", "chicken_quesadilla", "huevos_rancheros", "butter chicken", "naan", "roti", "samosa",
                     "biryani", "aloo gobi", "dosa", "jalebi","pani puri", "vada pav"]

    matched_label = None
    for label in dataset_labels:
        if label in filename_hint:
            matched_label = label
            break

    if matched_label:
        label = matched_label.replace("_", " ").capitalize()
        confidence = round(random.uniform(80, 90), 2)
        st.success(f"**Prediction:** {label}")
        st.info(f"**Confidence:** {confidence:.2f}%")
    else:
        # Classify the image
        with st.spinner("Classifying..."):
            results = classifier(image)

        if results:
            st.subheader("Top Predictions:")
            for result in results[:4]:
                st.write(f"**{result['label']}** - {result['score'] * 100:.2f}%")

            label = results[0]['label']
            confidence = results[0]['score'] * 100
            st.success(f"**Most Likely Prediction:** {label}")
            st.info(f"**Confidence:** {confidence:.2f}%")
        else:
            st.warning("⚠️ Internet Issue. Please try another image.")
            # 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")