import gradio as gr
from keras.models import load_model
from keras.preprocessing import image
from keras.applications.vgg16 import preprocess_input
import numpy as np
import cv2
from PIL import Image
import uuid

# Load your pre-trained model
model = KerasModel

# Function to preprocess the image for prediction
def preprocess_image(img_array, s=224):
    img = (img_array * 255).astype(np.uint8)  # Convert back to uint8
    img_array = cv2.resize(img, (s, s))
    img_array = img_array.astype(np.uint8)
    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
    return img_array

# Function to make predictions
def predict_mineral(img):
    # Convert the PIL Image to a NumPy array
    img_array = np.array(img)
    # Generate a unique filename for the temporary image
    img_path = f"temp_image_{uuid.uuid4()}.jpg"
    Image.fromarray((img_array * 255).astype(np.uint8)).save(img_path)
    # Preprocess the image and make predictions
    processed_img = preprocess_image(img_array)
    prediction = model.predict(processed_img)
    class_idx = np.argmax(prediction)
    # Update this with your class labels
    class_labels = ['biotite', 'granite', 'olivine', 'plagioclase', 'staurolite']
    print(f"Prediction scores: {prediction[0]}")
    print(f"Predicted class: {class_labels[class_idx]}")
    # Convert prediction scores to percentages
    prediction_scores = [f"{score*100:.2f}%" for score in prediction[0]]
    # Create a formatted string to display the class labels and prediction scores
    class_info = "\n".join([f"{label}: {score}" for label, score in zip(class_labels, prediction_scores)])
    return class_labels[class_idx], float(prediction[0][class_idx]), class_info

# Load CSS and JavaScript files


# Create a Gradio interface
with gr.Blocks() as iface:
    with gr.Row():
        image_input = gr.Image(label="Input Image")
        output_text = gr.Textbox(label="Predicted Class")
        output_number = gr.Number(label="Prediction Score")
    with gr.Row():
        class_output = gr.Textbox(label="Class Predictions")
        run_button = gr.Button("Predict")

    run_button.click(predict_mineral, inputs=image_input, outputs=[output_text, output_number, class_output])

iface.launch(share=True)