import gradio as gr
from fastai.vision.all import PILImage



def predict(image):
    # Transform the image
    pil_image = PILImage.create(image)
    
    # Predict
    preds, _, probs = learn.predict(pil_image)
    
    # Apply the threshold
    threshold = 0.425
    classes = [learn.dls.vocab[i] for i in range(len(probs)) if probs[i] > threshold]

    pet_type = "dog" if "dog" in classes else "cat" if "cat" in classes else None
    breeds = [breed for breed in classes if breed not in ["cat", "dog"]]

    if pet_type:
        breed_info = f"The breed is {'/'.join(breeds)}." if breeds else "The breed is not identified."
        return f"This is a {pet_type}. {breed_info}"
    else:
        return "This is not a cat, nor a dog."


# Define the Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.inputs.Image(shape=(224, 224)),
    outputs="text",
    live=True,
    title="Cat and Dog Image Classifier",
    description="Upload an image of a cat or a dog, and the model will identify the type and breed.",
    article="This model has been trained on the Oxford Pets dataset and might not recognize all types dog and cat breeds. For best results, use clear images."
)


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