import streamlit as st
import numpy as np
from PIL import Image
import tensorflow as tf
import os
from tensorflow.keras import layers, models
from tensorflow.keras.applications import EfficientNetB3

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

# -----------------------------------
# MODEL DEFINITION
# -----------------------------------
def create_model():
    base_model = EfficientNetB3(
        weights=None, 
        include_top=False, 
        input_shape=(IMG_SIZE, IMG_SIZE, 3)
    )
    
    base_model.trainable = True 
    
    model = models.Sequential([
        base_model,
        layers.GlobalAveragePooling2D(),
        layers.BatchNormalization(),
        layers.Dropout(0.3),
        layers.Dense(256, activation='relu'),
        layers.Dropout(0.2),
        layers.Dense(NUM_CLASSES, activation='softmax') 
    ])
    
    return model

# -----------------------------------
# MODEL LOADING (WEIGHTS ONLY)
# -----------------------------------
@st.cache_resource
def load_model_from_weights():
    model = create_model()
    
    weights_path = os.path.join(os.path.dirname(__file__), "kitchen_weights.weights.h5")
    
    try:
        model.build((None, IMG_SIZE, IMG_SIZE, 3))
        
        model.load_weights(weights_path)
        return model
    except Exception as e:
        st.error(f"Ağırlıklar yüklenemedi: {e}")
        return None

model = load_model_from_weights()

# -----------------------------------
# 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 and model is not None:
    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"Hata oluştu: {e}")