import gradio as gr
from utils import model_initialization, prediction
from PIL import Image
from typing import Dict, Any


def gradio_interface(image: Image.Image) -> Dict[str, Any]:
    """
    Perform image classification using a pre-trained model.

    Args:
        image (Image.Image): The input image uploaded by the user.

    Returns:
        Dict[str, Any]: A dictionary containing the classification result with the
        most promising label and confidence score.
    """
    # Initialize the pre-trained pipeline
    pipe = model_initialization()
    
    # Perform prediction on the uploaded image
    result = prediction(pipe, image)

    return result


# Define the Gradio interface
demo = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Image(type="pil", label="Upload Image"),  # Accepts PIL Image input
    outputs=gr.JSON(label="Prediction Details"),  # Outputs as JSON
    title="RESNET WILL NEVER DIE. Image Classification with ResNet-18",
    description=(
        "Welcome to the Image Classification Demo! Upload an image to classify it using"
        "ResNet-18 model. The model will predict the most likely label along with its confidence score."
    ),
    theme="soft",
    examples=[["artifacts/ball.png"], ["artifacts/panda.jpg"]],
)

# Launch the Gradio app
if __name__ == "__main__":
    demo.launch()