AItool's picture
Update app.py
94d7ce8 verified
__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()