import gradio as gr
import numpy as np
from PIL import Image
import tensorflow as tf, keras
from huggingface_hub import hf_hub_download
from keras.applications.resnet50 import preprocess_input

model_path = hf_hub_download(
    repo_id="Branden28/ResNet50_CUB-200",
    filename="resnet50_model_final.keras"
)
model = keras.models.load_model(model_path)

# Dummy prediction function
def predict(image):
  if image is None:
    return "Error: No image"

  image_resized = image.resize((224,224))

  image_array=np.asarray(image_resized)
  image_array=np.expand_dims(image_array, axis=0)

  preprocessed_input = preprocess_input(image_array)

  raw_predictions = model.predict(preprocessed_input)[0]
  predicted_index = int(np.argmax(raw_predictions))
  top_confidence_score = float(raw_predictions[predicted_index])

  return predicted_index, top_confidence_score
  # return predicted_label_string, top_confidence_score

with gr.Blocks(title="🧠 Image Classification") as demo:
    gr.Markdown("## 🧠 Image Classification")
    # gr.Markdown("## Upload an image and provide the true label to compare predictions.")

    with gr.Row():
        with gr.Column():
            gr.Markdown("### 🧾 Upload Image")
            gr.Markdown("Upload the image to make inference")
            image = gr.Image(type="pil")
            submit = gr.Button("Predict")

        with gr.Column():
            gr.Markdown("### 🧾 Result")
            gr.Markdown("Prediction will give the predicted label and ground truth")
            pred = gr.Number(label="Predicted Label", interactive=False)
            confidence_out = gr.Number(label="Ground Truth", interactive=False)

    submit.click(
        fn=predict,
        inputs=[image],
        outputs=[pred, confidence_out]
    )


# demo.launch(server_name="127.0.0.1", server_port=7860, show_error=True)
if __name__ == "__main__":
    demo.launch()