import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image

# -----------------------------
# 1. Instellingen
# -----------------------------
# Gebruik dezelfde image size als bij training!
IMG_SIZE = (224, 224)   # of (128, 128) als je model daarop is getraind

CLASS_NAMES = ['Karacadag', 'Basmati', 'Jasmine', 'Arborio', 'Ipsala']

st.set_page_config(page_title="Rice Classifier", page_icon="🌾")
st.title("🌾 Rice Classifier")
st.write("Upload een rijstkorrel-afbeelding om het type te laten voorspellen.")

# -----------------------------
# 2. Laad TFLite model
# -----------------------------
@st.cache_resource
def load_interpreter():
    interpreter = tf.lite.Interpreter(model_path="rice_model.tflite")
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    return interpreter, input_details, output_details

try:
    interpreter, input_details, output_details = load_interpreter()
except Exception as e:
    st.error(f"Kon TFLite-model niet laden: {e}")
    st.stop()

# -----------------------------
# 3. Upload afbeelding
# -----------------------------
uploaded = st.file_uploader("Upload een afbeelding", type=["jpg", "jpeg", "png"])

if uploaded is None:
    st.info("👆 Kies hierboven een afbeelding.")
else:
    # Toon de geüploade afbeelding
    img = Image.open(uploaded)
    st.image(img, caption="Geüploade afbeelding", use_container_width=True)

    # Zorg dat het echt RGB is
    if img.mode != "RGB":
        img = img.convert("RGB")

    # Resize + normaliseer
    img = img.resize(IMG_SIZE)
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0).astype(np.float32)

    # -----------------------------
    # 4. Voorspelling
    # -----------------------------
    interpreter.set_tensor(input_details[0]['index'], img_array)
    interpreter.invoke()
    prediction = interpreter.get_tensor(output_details[0]['index'])

    idx = int(np.argmax(prediction))
    confidence = float(prediction[0][idx])

    st.subheader("🔍 Resultaat")
    st.write(f"**Predicted class:** {CLASS_NAMES[idx]}")
    st.write(f"**Confidence:** {confidence:.4f}")