import gradio as gr
import numpy as np
from tensorflow.keras.preprocessing import image
from tensorflow.keras.models import load_model
from huggingface_hub import hf_hub_download

# Define the path to the model on Hugging Face
repo_id = "johansetiawan17/intel-image-classification"  
filename = "model-intel-image-classification.h5"

# Download the model from Hugging Face
model_path = hf_hub_download(repo_id=repo_id, filename=filename) 

# Load the model
model = load_model(model_path)

# Define variables
class_labels = ['buildings', 'forest', 'glacier', 'mountain', 'sea', 'street']

# Define the prediction function
def predict_image(img_input):
  image.save_img('./image.jpeg', img_input)
  img = image.load_img('./image.jpeg', target_size=(150, 150))
  img_array = image.img_to_array(img)
  img_array = np.expand_dims(img_array, axis=0)
  img_array /= 255.

  prediction = model.predict(img_array)
  predicted_class = np.argmax(prediction)
  confidence = prediction[0][predicted_class] * 100

  result = f"Prediction: {class_labels[predicted_class]} ({confidence:.2f}%)"
  return result

# Create the Gradio interface
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=6),
    title="Natural Scenes Classification",
    description="Upload an natural scenes or landscape image to classify it into one of the six categories: buildings, forest, glacier, mountain, sea, or street.",
)

# Launch the Gradio interface
iface.launch(share=True)