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

@st.cache_resource
def load_model():
    processor = AutoImageProcessor.from_pretrained("microsoft/beit-base-patch16-224")
    model = AutoModelForImageClassification.from_pretrained("microsoft/beit-base-patch16-224")
    return processor, model

processor, model = load_model()

st.title("Kalorieestimering med AI")

uploaded_file = st.file_uploader("Upload et billede af din mad", type=["jpg", "jpeg", "png"])

if uploaded_file:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Dit billede", use_column_width=True)

    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    logits = outputs.logits
    probs = torch.nn.functional.softmax(logits, dim=1)
    top_probs, top_labels = torch.topk(probs, k=1)

    confidence = top_probs[0].item() * 100
    label = model.config.id2label[top_labels[0].item()]

    st.markdown(f"### Identificeret: **{label}** ({confidence:.1f}% sikkerhed)")

    if confidence < 70:
        st.warning("Usikker klassificering – vælg manuelt")
        manual = st.selectbox("Vælg fødevare manuelt", sorted(model.config.id2label.values()))
        label = manual

    st.markdown("### Antaget portionsstørrelse og estimeret energiindhold")
    st.markdown(f"- 1 portion **{label}** (ca. 200g)")
    st.markdown("- Estimeret energi: **~300 kcal** *(eksempelværdi)*")

    feedback = st.radio("Er forslaget korrekt?", ["Ja", "Nej", "Ved ikke"])
    if feedback == "Nej":
        korrekt_label = st.text_input("Hvad forestiller billedet egentlig?")
        if korrekt_label:
            st.success(f"Tak for dit input: *{korrekt_label}* gemt.")