import streamlit as st
import numpy as np
import cv2
import joblib
import tensorflow as tf
from tensorflow.keras.applications import (
    ResNet50, VGG16, EfficientNetV2B0, InceptionV3, ResNet101, DenseNet201
)
from tensorflow.keras.preprocessing import image
import os

# Define available models
MODELS = {
    'ResNet50': ResNet50(weights='imagenet', include_top=False, input_shape=(224, 224, 3), pooling="avg"),
    'VGG16': VGG16(weights='imagenet', include_top=False, input_shape=(224, 224, 3), pooling="avg"),
    'EfficientNetV2B0': EfficientNetV2B0(weights='imagenet', include_top=False, input_shape=(224, 224, 3), pooling="avg"),
    'InceptionV3': InceptionV3(weights='imagenet', include_top=False, input_shape=(224, 224, 3), pooling="avg"),
    'ResNet101': ResNet101(weights='imagenet', include_top=False, input_shape=(224, 224, 3), pooling="avg"),
    'DenseNet201': DenseNet201(weights='imagenet', include_top=False, input_shape=(224, 224, 3), pooling="avg")
}

# Load trained models
MODEL_PATHS = {model_name: f"{model_name}_catboost.pkl" for model_name in MODELS}
trained_models = {}

# Load the trained models into memory
for model_name, path in MODEL_PATHS.items():
    if os.path.exists(path):
        trained_models[model_name] = joblib.load(path)

# Define class names (modify based on dataset)
CLASS_NAMES = ['ADI','DEB','LYM','MUC','MUS','NOR','STR','TUM']  # Update with actual labels

# Streamlit UI
st.title("Multi-Model Image Classifier")
st.markdown("Upload an image and select which models and classes to use for prediction.")

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

# Select Models
selected_models = st.multiselect("Select models for prediction:", list(trained_models.keys()))

# Select Classes
selected_classes = st.multiselect("Select classes to predict:", CLASS_NAMES, default=CLASS_NAMES)

# Function to preprocess image
def preprocess_image(img_path):
    img = image.load_img(img_path, target_size=(224, 224))
    img_array = image.img_to_array(img) / 255.0  # Normalize
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

# Perform prediction
if uploaded_file and selected_models:
    # Read and preprocess the image
    image_np = np.array(bytearray(uploaded_file.read()), dtype=np.uint8)
    image_np = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
    image_np = cv2.resize(image_np, (224, 224)) / 255.0
    image_np = np.expand_dims(image_np, axis=0)

    # Extract Features
    extracted_features = {}
    for model_name in selected_models:
        extracted_features[model_name] = MODELS[model_name].predict(image_np)

    # Predict using each selected CatBoost model
    predictions = {}
    for model_name in selected_models:
        X_input = extracted_features[model_name]
        catboost_model = trained_models[model_name]
        y_pred = catboost_model.predict_proba(X_input)[0]  # Get probabilities
        predictions[model_name] = y_pred

    # Display results
    st.subheader("Prediction Results")

    for model_name, y_pred in predictions.items():
        st.write(f"**Model: {model_name}**")
        for i, class_name in enumerate(CLASS_NAMES):
            if class_name in selected_classes:
                st.write(f"  - {class_name}: **{y_pred[i]:.4f}**")