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

# Load the saved model
model = tf.saved_model.load('saved_model/embryo_classifier')

# Define image size (should match the input size of your model)
IMG_SIZE = (300, 300)

# Function to preprocess the image
def preprocess_image(image):
    image = image.resize(IMG_SIZE, Image.ANTIALIAS)
    inp_numpy = np.array(image)[None]
    inp = tf.constant(inp_numpy, dtype='float32')
    return inp

# Streamlit interface
st.title("Embryo Quality Assessment")

st.write("Upload an embryo image to classify its quality.")

# File uploader
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    image = Image.open(uploaded_file).convert('RGB')
    st.image(image, caption='Uploaded Image.', use_column_width=True)
    
    st.write("Classifying...")
    
    # Preprocess the image
    processed_image = preprocess_image(image)
    
    # Make predictions
    class_scores = model(processed_image)[0].numpy()
    predicted_class = class_scores.argmax()
    
    # Display the results
    classes = ['Low Quality', 'Medium Quality', 'High Quality']  # Adjust according to your classes
    st.write(f"Prediction: {classes[predicted_class]}")
    st.write(f"Confidence: {np.max(class_scores) * 100:.2f}%")