__all__ = [
    'learn', 'classify_image', 'categories', 'image', 'label', 'examples', 'intf'
]

from fastai.vision.all import *
import gradio as gr
import timm  # keep only if your model actually needs timm

def is_real(x):
    # Recreate the original function used during training
    # Example: return x > 0.5 or whatever logic you had
    pass  # Replace with actual logic


# ✅ Load the fastai model properly — no manual pickle.load()
learn = load_learner('model.pkl', cpu=True)

# Define your categories exactly as trained
categories = ('Virtual Staging', 'Real')

# Prediction function for Gradio
def classify_image(img):
    pred, idx, probs = learn.predict(img)
    # Cast to float so Gradio handles them cleanly
    return dict(zip(categories, map(float, probs)))

# Gradio UI components
image = gr.inputs.Image(shape=(192, 192))
label = gr.outputs.Label()
examples = ['virtual.jpg', 'real.jpg']  # sample files in your Space

# Create and launch interface
intf = gr.Interface(
    fn=classify_image,
    inputs=image,
    outputs=label,
    examples=examples,
    share=True
)

if __name__ == "__main__":
    intf.launch()