import gradio as gr
import tensorflow as tf
from PIL import Image
import numpy as np

labels = ['Birch Forest', 'Cave', 'Cherry Grove', 'Dark Forest', 'Deep Dark', 'Desert', 'End', 'Forest', 'Jungle', 'Mushroom Fields', 'Nether', 'Ocean', 'Plains', 'Savanna', 'Swamp', 'Taiga']

def predict_biome(uploaded_file):
    if uploaded_file is None:
        return "No file uploaded.", None, "No prediction"

    model = tf.keras.models.load_model('biomes-xception-model.keras')

    # Load the image from the file path
    with Image.open(uploaded_file).convert('RGB') as img:
        img = img.resize((150, 150))
        img_array = np.array(img)

        prediction = model.predict(np.expand_dims(img_array, axis=0))

        # Identify the most confident prediction
        confidences = {labels[i]: np.round(float(prediction[0][i]), 2) for i in range(len(labels))}
        return img, confidences

# Define the Gradio interface
iface = gr.Interface(
    fn=predict_biome,  # Function to process the input
    inputs=gr.File(label="Upload File"),  # File upload widget
    outputs=["image", "text"],  # Output types for image and text
    title="Minecraft Biomes Classifier",  # Title of the interface
    description="Upload a picture of a Minecraft Biome (preferably a Birch Forest, Cave, Cherry Grove, Dark Forest, Deep Dark, Desert, End, Forest, Jungle, Mushroom Fields, Nether, Ocean, Plains, Savanna, Swamp or Taiga) to see what Biome it is and the models confidence level."  # Description of the interface
)

# Launch the interface
iface.launch()