import os
import tensorflow as tf

# Disable all GPUS
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

# Load your pre-trained model
def load_model():
    model = tf.keras.models.load_model('depi-graduation-project.h5')  # Replace with your model's path
    return model

model = load_model()


# Define the labels (categories)
labels = ['Water', 'Cloudy', 'Desert', 'Green Area']

# Function to preprocess the image and predict the class
def classify_image(image):
    img = image.resize((128, 128))  # Resize the image
    img = np.array(img) / 255.0     # Normalize the image
    img = np.expand_dims(img, axis=0)  # Add batch dimension
    prediction = model.predict(img)
    predicted_class = labels[np.argmax(prediction)]
    
    # Prepare output with probabilities
    return {labels[i]: float(prediction[0][i]) for i in range(len(labels))}

# Define the Gradio interface
image_input = gr.inputs.Image(shape=(128, 128))
label_output = gr.outputs.Label(num_top_classes=4)

# Launch the interface
gr.Interface(fn=classify_image, 
             inputs=image_input, 
             outputs=label_output, 
             title="Satellite Image Classification", 
             description="Classify satellite images into four types: Water, Cloudy, Desert, Green Area").launch()