import streamlit as st
import onnxruntime as ort
from PIL import Image
import numpy as np



MODEL_PATH = "models/mobilenetv3_catsdogs_quantized_static.onnx"

# Load ONNX model
@st.cache_resource
def load_model():
    session = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])
    return session

session = load_model()


def preprocess_image(image):

    image = image.resize((224, 224))
    

    img = np.array(image).astype(np.float32) / 255.0

    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    img = (img - mean) / std


    img = np.transpose(img, (2, 0, 1))


    img = np.expand_dims(img, axis=0).astype(np.float32)

    return img

CLASSES = ["Cat", "Dog"]

st.set_page_config(page_title="Cat vs Dog Classifier ", layout="centered")

st.title("Cat vs Dog Classifier")
st.markdown("Upload an image and let the model decide if it’s a **cat** or a **dog**.")

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

if uploaded_file:
    
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Uploaded Image", width="stretch")

 
    img_tensor = preprocess_image(image)


    with st.spinner("Running inference..."):
        outputs = session.run(None, {"input": img_tensor})
        probs = np.exp(outputs[0]) / np.sum(np.exp(outputs[0])) 
        pred_idx = int(np.argmax(probs))
        pred_class = CLASSES[pred_idx]
        confidence = probs[0][pred_idx] * 100

  
    st.markdown("---")
    st.subheader(f"Prediction: **{pred_class}**")
    st.write(f"Confidence: `{confidence:.2f}%`")