import streamlit as st
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForImageClassification
import torch

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

st.title("Food Classification")
st.markdown("AI-powered food recognition using a pretrained Food101 model")

# Load processor and model
processor = AutoImageProcessor.from_pretrained("nateraw/food")
model = AutoModelForImageClassification.from_pretrained("nateraw/food")
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()  # set model to evaluation mode

# Function to preprocess and classify image
def classify_image(image: Image.Image):
    inputs = processor(images=image.convert("RGB"), return_tensors="pt")
    # Move to device
    for k, v in inputs.items():
        inputs[k] = v.to(device)
    # Forward pass
    with torch.no_grad():
        outputs = model(**inputs)
    # Softmax probabilities
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    # Get top 5 predictions
    top_prob, top_class = torch.topk(probs, k=5)
    id2label = model.config.id2label
    results = []
    for prob, idx in zip(top_prob[0], top_class[0]):
        results.append({"label": id2label[idx.item()], "score": prob.item()})
    return results

# Streamlit layout
col1, col2 = st.columns(2)

with col1:
    st.subheader("Upload Image")
    uploaded_file = st.file_uploader("Choose a food image", type=["jpg", "jpeg", "png"])
    if uploaded_file:
        image = Image.open(uploaded_file)
        st.image(image, width=400)

with col2:
    st.subheader("Results")
    if uploaded_file:
        with st.spinner("Analyzing..."):
            results = classify_image(image)
            if results:
                for i, r in enumerate(results, 1):
                    st.write(f"{i}. {r['label']} — {r['score']*100:.2f}%")
                    st.progress(r['score'])
            else:
                st.info("Could not classify the image. Try again.")
    else:
        st.info("Upload an image to classify.")

st.markdown("---")
st.caption("Week 1 Project - Image Classification Pipeline")