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

def label_func(fn): return path/'masks1b-binary'/f'{fn.stem}.png'

repo_id = "hugginglearners/kvasir-seg"
learn = from_pretrained_fastai(repo_id)
#labels = learn.dls.vocab

def predict(img):
    img = PILImage.create(img)
    pred, _, _ = learn.predict(img)
    return PILMask.create(pred*255)

interface_options = {
    "title": "kvasir-seg fastai segmentation",
    "description": "tbd",
    "interpretation": "default",
    "layout": "horizontal",
#    "examples": [
#        "100098.jpg",
#        "100002.jpg",
#        "100048.jpg"
#    ],
    "allow_flagging": "never",
}

demo = gr.Interface(
    fn=predict,
    inputs=gr.inputs.Image(shape=(224, 224)),
    outputs=gr.inputs.Image(shape=(224, 224)),
    **interface_options,
)

launch_options = {
    "enable_queue": True,
    "share": False,
}

demo.launch(**launch_options)