app.py

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


st.set_page_config(page_title="Facial Emotion Classification", layout="centered")
st.title("Facial Emotion Classification")

CLASS_NAMES = ['angry', 'disgust', 'fear', 'happy', 'neutral', 'sad', 'surprise']

MODEL_PATHS = {
    "enhanced_cnn": "ML_project_/enhanced_cnn_20251127_115643/best_model.keras",
    "stanford_cnn": "ML_project_/stanford_cnn_20251127_125730/stanford_five_layer_final_0.7158.keras",
    "resnet50":     "ML_project_/resnet50_fixed_20251204_100110/resnet50_fer_final.keras",
}

#load the models 
@st.cache_resource
def load_models(paths: dict):
    out = {}
    for name, p in paths.items():
        out[name] = tf.keras.models.load_model(p)
    return out

try:
    models = load_models(MODEL_PATHS)
except Exception as e:
    st.error(f"Failed to load one of the .keras models.\n\n{e}")
    st.stop()
# pre-processing strategies used to train and test the models.
def pil_to_rgb_np(pil_img):
    return np.array(pil_img.convert("RGB"), dtype=np.uint8)  

def resize_np_uint8(img_uint8, size_hw):
    # use PIL for simple resizing
    pil = Image.fromarray(img_uint8)
    pil = pil.resize((size_hw[1], size_hw[0])) 
    return np.array(pil, dtype=np.uint8)

def preprocess_for_model(model_name: str, img_uint8_rgb: np.ndarray):
    """
    Returns float32 batch for each model.
    """
    if "resnet50" in model_name:
        x = resize_np_uint8(img_uint8_rgb, (224, 224))
        x = x.astype(np.float32)
        x = tf.keras.applications.resnet50.preprocess_input(x)
        x = np.expand_dims(x, axis=0)  
        return x

    if "stanford_cnn" in model_name:
        x = resize_np_uint8(img_uint8_rgb, (48, 48)).astype(np.float32)
        gray = np.dot(x[..., :3], [0.2989, 0.5870, 0.1140])
        gray = np.expand_dims(gray, axis=-1)     
        gray = gray / 255.0
        gray = np.expand_dims(gray, axis=0)      
        return gray.astype(np.float32)

    if "enhanced_cnn" in model_name:
        x = resize_np_uint8(img_uint8_rgb, (48, 48)).astype(np.float32) / 255.0
        x = np.expand_dims(x, axis=0)  
        return x.astype(np.float32)

    x = resize_np_uint8(img_uint8_rgb, (48, 48)).astype(np.float32) / 255.0
    x = np.expand_dims(x, axis=0)
    return x.astype(np.float32)


def ensure_probs(y):
    """
    Converts model output to a 1D probability vector [C].
    """
    y = np.asarray(y)
    if y.ndim == 2:
        y = y[0]
    if not (np.all(y >= 0) and np.all(y <= 1.0) and np.isclose(np.sum(y), 1.0, atol=1e-3)):
        y = tf.nn.softmax(y).numpy()
    return y.astype(np.float32)


#predicition function 
def predict_ensemble(pil_img):
    img_uint8 = pil_to_rgb_np(pil_img)

    per_model_results = {}
    for name, model in models.items():
        x = preprocess_for_model(name, img_uint8)
        y = model.predict(x, verbose=0)
        probs = ensure_probs(y)

        if probs.shape[0] != len(CLASS_NAMES):
            raise ValueError(
                f"{name} outputs {probs.shape[0]} classes, but CLASS_NAMES has {len(CLASS_NAMES)}."
            )
        
        # Store both probabilities and max confidence info
        max_idx = int(np.argmax(probs))
        max_conf = float(probs[max_idx])
        
        per_model_results[name] = {
            'probs': probs,
            'predicted_class': CLASS_NAMES[max_idx],
            'max_confidence': max_conf,
            'predicted_idx': max_idx
        }
#model with highest probability
    best_model = max(per_model_results.keys(), 
                    key=lambda k: per_model_results[k]['max_confidence'])
    
    ensemble_result = per_model_results[best_model]
    
    return per_model_results, ensemble_result, best_model


# Basic UI to upload image and get predicitions
uploaded = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg", "webp", "bmp"])

if uploaded is not None:
    img = Image.open(io.BytesIO(uploaded.read()))
    st.image(img, caption="Uploaded image", use_container_width=True)

    if st.button("Predict", type="primary"):
        try:
            per_model_results, ensemble_result, best_model = predict_ensemble(img)

            # Show ensemble result (from best model)
            st.subheader("Ensemble Prediction")
            col1, col2 = st.columns([2, 1])
            
            with col1:
                st.success(f"**{ensemble_result['predicted_class']}**")
            with col2:
                st.metric("Confidence", f"{ensemble_result['max_confidence']:.3f}")

            st.subheader("All Model Predictions")
            
            for name in ["enhanced_cnn", "stanford_cnn", "resnet50"]:
                if name in per_model_results:
                    result = per_model_results[name]
                    st.write(f"**{name}** → {result['predicted_class']} (conf: {result['max_confidence']:.3f})")
                
            st.subheader(f"Detailed Probabilities ({best_model})")
            winning_probs = ensemble_result['probs']
            
            for i, cname in enumerate(CLASS_NAMES):
                prob_val = float(winning_probs[i])
                
                if i == ensemble_result['predicted_idx']:
                    st.write(f"**{cname}**: {prob_val:.4f}")
                else:
                    st.write(f"{cname}: {prob_val:.4f}")
                
                st.progress(prob_val)

        except Exception as e:
            st.error(f"Prediction failed: {e}")