import streamlit as st
import numpy as np
import cv2
from tensorflow.keras.models import load_model


from tensorflow.keras.layers import InputLayer

# Custom object scope to handle InputLayer configuration
custom_objects = {'InputLayer': InputLayer}

@st.cache_resource
def load_ocular_model():
    return load_model('odir_cnn_model.h5', custom_objects=custom_objects)

model = load_ocular_model()


model = load_ocular_model()

# Define the labels
LABELS = ['Normal (N)', 'Diabetes (D)', 'Glaucoma (G)', 'Cataract (C)', 
          'Age related Macular Degeneration (A)', 'Hypertension (H)', 
          'Pathological Myopia (M)', 'Other diseases/abnormalities (O)']

# Preprocess the image
def preprocess_image(image, img_size=128):  # Change img_size to 128
    img = cv2.resize(image, (img_size, img_size))
    img = img / 255.0  # Normalize
    img = np.expand_dims(img, axis=0)  # Add batch dimension
    return img


# Predict diseases based on left and right images
def predict_diseases(left_image, right_image, model):
    left_img = preprocess_image(left_image)
    right_img = preprocess_image(right_image)
    
    left_predictions = model.predict(left_img)
    right_predictions = model.predict(right_img)
    
    combined_predictions = (left_predictions + right_predictions) / 2
    
    pred_labels = {label: float(pred) for label, pred in zip(LABELS, combined_predictions[0])}
    
    return pred_labels

# Streamlit UI
st.title("Ocular Disease Prediction")
st.write("Upload left and right eye images to predict ocular diseases.")

# Image upload widgets
left_image_file = st.file_uploader("Upload Left Eye Image", type=["jpg", "jpeg", "png"])
right_image_file = st.file_uploader("Upload Right Eye Image", type=["jpg", "jpeg", "png"])

if left_image_file is not None and right_image_file is not None:
    # Load images
    left_image = cv2.imdecode(np.fromstring(left_image_file.read(), np.uint8), 1)
    right_image = cv2.imdecode(np.fromstring(right_image_file.read(), np.uint8), 1)
    
    st.image(left_image, caption='Left Eye Image', use_column_width=True)
    st.image(right_image, caption='Right Eye Image', use_column_width=True)
    
    st.write("Processing images and predicting diseases...")
    
    # Predict diseases
    disease_predictions = predict_diseases(left_image, right_image, model)
    
    # Display predictions
    st.write("Predictions:")
    for disease, score in disease_predictions.items():
        st.write(f"{disease}: {score:.4f}")

# Information about the model
st.write("**Model Information:**")
st.write("This model predicts various ocular diseases based on left and right fundus images.")