from fastai.vision.all import load_learner, PILImage
import gradio as gr
from PIL import Image

# Load the trained model
model = load_learner('model.pkl')

# Define a function to make predictions
def predict(image):
    try:
        print("📸 Received image for prediction")

        # Convert to Fastai's expected PILImage format
        image = PILImage.create(image)

        # Run prediction
        pred, _, probs = model.predict(image)
        print(f"✅ Prediction successful: {pred}, Confidence: {probs.max():.2f}")
        
        return f"Prediction: {pred} (Confidence: {probs.max():.2f})"
    
    except Exception as e:
        print(f"❌ Error during prediction: {e}")
        return f"Error: {e}"

# Create the Gradio web interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Textbox(),
    title="Cat vs Dog Classifier",
    description="Upload an image of a cat or dog and let the model classify it!"
)

# Launch the Gradio app
interface.launch()