app.py

# AUTOGENERATED! DO NOT EDIT! File to edit: app.ipynb.

# %% auto 0
__all__ = ['learn', 'categories', 'image', 'label', 'examples', 'intf', 'is_cat', 'classify_img']

# %% app.ipynb 1
from fastai.vision.all import *
import gradio as gr

# def is_cat(x): return x[0].isupper()

# %% app.ipynb 3
learn_swin = load_learner("model_swin.pkl")
learn_conv = load_learner("model_conv.pkl")

# %% app.ipynb 5
categories = [
    'Bathroom', 'Bedroom', 'Floor plan', 'Front', 'Home Office', 'Kitchen', 
    'Laundry', 'Living room', 'Parking', 'Porch', 'Swimming pool', 'Views', 
    'Walk In Closet', 'Yard'
]

def classify_img(img, use_conv):
    pred,idx,probs = learn_conv.predict(img) if use_conv else learn_swin.predict(img)
    return dict(zip(categories, map(float, probs)))


# %% app.ipynb 7

examples = [
    ["kitchen.jpg", False], 
    ["living_room.jpg", False], 
    ["living_room2.jpg", False],
    ["kitchen.jpg", True], 
    ["living_room.jpg", True], 
    ["living_room2.jpg", True]
]

intf = gr.Interface(
    fn=classify_img, 
    inputs=[
        gr.components.Image(shape=(640, 480)), 
        gr.components.Checkbox(label="Use conv model", value=False), 
    ], 
    outputs=[
        gr.components.Label()
    ], 
    examples=examples
)
intf.launch(inline=False)