# app.py
from fastai.vision.all import *
import gradio as gr
import pathlib

# Fix for "learn.predict cannot use Path objects directly" issue on some systems
# when loading models trained on Windows and deploying on Linux/Mac.
# This ensures pathlib.PosixPath is available for fastai.
# See: https://forums.fast.ai/t/model-deploy-error-on-huggingface-spaces/96386/2
if platform.system() == "Windows":
    temp = pathlib.PosixPath
    pathlib.PosixPath = pathlib.WindowsPath
else:
    temp = pathlib.WindowsPath
    pathlib.WindowsPath = pathlib.PosixPath


# Load the exported FastAI model
# Make sure 'model.pkl' is in the same directory as 'app.py'
try:
    learn = load_learner('model (1).pkl')
except Exception as e:
    print(f"Error loading model: {e}")
    print("Please ensure 'model.pkl' is in the same directory as 'app.py' and was exported correctly.")
    exit()

# Restore original pathlib for other operations if needed (optional, but good practice)
if platform.system() == "Windows":
    pathlib.PosixPath = temp
else:
    pathlib.WindowsPath = temp


# Define the prediction function
def predict_image(img):
    """
    Takes an image, runs prediction using the FastAI model,
    and returns a human-readable string.
    """
    if img is None:
        return "Please upload an image."

    # FastAI expects a PIL Image or similar object
    is_cat, _, probs = learn.predict(img)

    # is_cat will be True if it's a cat, False otherwise (based on your label_func)
    # probs are the probabilities for [not_cat, cat] if your labels were sorted that way,
    # or [False, True] based on your boolean label function.
    # We need to map the boolean output back to "Cat" or "Not a Cat" and get the probability for the predicted class.

    # Assuming `is_cat` maps to the "True" class (cat) and `not_cat` maps to "False" class (not a cat)
    # The probabilities will be indexed according to the sorted unique labels, which for a boolean label_func will be [False, True]
    # So probs[0] is probability of not being a cat, probs[1] is probability of being a cat.

    if is_cat:
        label = "Cat"
        confidence = probs[1].item() # probability of being a cat
    else:
        label = "Not a Cat"
        confidence = probs[0].item() # probability of not being a cat

    return f"Prediction: {label} (Confidence: {confidence:.4f})"


# Create Gradio Interface
# Inputs: gr.Image() for image upload
# Outputs: gr.Label() or gr.Textbox() for displaying the result
gr_interface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.Textbox(label="Prediction"),
    title="Cat vs. Not-a-Cat Classifier",
    description="Upload an image to determine if it's a cat or something else!",
    examples=[
        ["examples/cat_example.jpg"],  # Make sure these paths are correct
        ["examples/dog_example.jpg"]
    ]
)

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