import os
os.environ["TF_USE_LEGACY_KERAS"] = "1"

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

from tensorflow.keras.models import load_model
from tensorflow.keras.applications.efficientnet import preprocess_input

# -----------------------------------
# CONFIG
# -----------------------------------
IMG_SIZE = 300
CLASS_NAMES = ['cup', 'fork', 'glass', 'knife', 'plate', 'spoon']

# -----------------------------------
# MODEL LOADING
# -----------------------------------
@st.cache_resource
def load_model_file():
    return keras.models.load_model("src/kitchen1.keras", compile=False)

model = load_model_file()

# -----------------------------------
# IMAGE PREPROCESS
# -----------------------------------
def process_image(img):
    img = img.resize((IMG_SIZE, IMG_SIZE))
    img = np.array(img)

    if img.ndim == 2:
        img = np.stack([img]*3, axis=-1)
    elif img.shape[-1] == 4:
        img = img[..., :3]

    img = np.expand_dims(img, 0)
    img = preprocess_input(img)
    return img

# -----------------------------------
# UI LAYOUT
# -----------------------------------
st.title("🍽️ Kitchenware Classifier")
st.write("Upload an image to classify kitchen items using EfficientNetB3.")

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

if uploaded:
    img = Image.open(uploaded).convert("RGB")

    col1, col2 = st.columns(2)

    with col1:
        st.image(img, caption="Uploaded Image", use_container_width=True)

    with col2:
        st.write("Classifying...")

        image_tensor = process_image(img)
        preds = model.predict(image_tensor)

        predicted_index = np.argmax(preds)
        predicted_class = CLASS_NAMES[predicted_index]
        confidence = float(np.max(preds))

        st.success(f"Prediction: **{predicted_class.upper()}**")
        st.metric("Confidence", f"{confidence * 100:.2f}%")
        st.progress(int(confidence * 100))