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

# Assuming model.pkl is in the same directory as this script, or specify the full path.
# For Hugging Face Spaces, ensure model.pkl is in the root directory of your Space.
MODEL_PATH = 'model.pkl' # Adjusted for Hugging Face Spaces where model.pkl will be in the root
learn = load_learner(MODEL_PATH)

# IMPORTANT: For inference, aggressively remove all augmentation transforms from the validation DataLoaders.
# We want to preserve only the Normalize transform if it exists, which is crucial for pre-trained models.
# This comprehensive cleanup targets item_tfms, before_batch.tfms, and filters after_batch.tfms.
if hasattr(learn.dls, 'valid_dl'):
    # 1. Clear item_tfms: These are transforms applied to individual items (like Resize from your training config)
    learn.dls.valid_dl.item_tfms = L()
    print(f"DEBUG: After cleanup, learn.dls.valid_dl.item_tfms: {learn.dls.valid_dl.item_tfms}")

    # 2. Clear all transforms from before_batch: This is where aug_transforms typically reside
    learn.dls.valid_dl.before_batch.tfms = L()
    print(f"DEBUG: After cleanup, learn.dls.valid_dl.before_batch.tfms: {learn.dls.valid_dl.before_batch.tfms}")

    # 3. Filter after_batch.tfms to only keep Normalize
    new_after_batch_tfms = L()
    if hasattr(learn.dls.valid_dl, 'after_batch'):
        for tfm in learn.dls.valid_dl.after_batch.tfms:
            if isinstance(tfm, Normalize):
                new_after_batch_tfms.append(tfm)
    learn.dls.valid_dl.after_batch.tfms = new_after_batch_tfms
    print(f"DEBUG: After cleanup, learn.dls.valid_dl.after_batch.tfms: {learn.dls.valid_dl.after_batch.tfms}")

# Define the prediction function
def predict_image(img):
    # Handle the case where no image is provided (img is None)
    if img is None:
        return {f"Error: No image provided. ": 1.0} # Return a default error message in the expected format

    # Ensure the input is a fastai PILImage object
    img_fastai = PILImage.create(img)

    pred, pred_idx, probs = learn.predict(img_fastai)

    # Convert predictions to a dictionary with class names and their probabilities
    return {learn.dls.vocab[i]: float(probs[i]) for i in range(len(learn.dls.vocab))}

# Create a Gradio interface
if __name__ == '__main__':
    gr.Interface(fn=predict_image,
                 inputs=gr.Image(type="pil"),
                 outputs=gr.Label(),
                 title="My Fastai Image Classifier",
                 description="Upload an image to get a classification prediction."
                ).launch()