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

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

# MODEL LOADING
@st.cache_resource
def load_model_file():
    model_path = os.path.join(os.path.dirname(__file__), "kitchen1.keras")
    
    return tf.keras.models.load_model(model_path, compile=False)

try:
    model = load_model_file()
except Exception as e:
    st.error(f"Error: {e}")
    st.stop()

# 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 = tf.keras.applications.efficientnet.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:
    try:
        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))
            
    except Exception as e:
        st.error(f"Error: {e}")