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)

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

interface_options = {
    "title": "kvasir-seg fastai segmentation",
    "description": "Demonstration of segmentation of gastrointestinal polyp images. This app is for reference only. It should not be used for medical diagnosis. Model was trained on Kvasir SEG dataset (https://datasets.simula.no/kvasir-seg/)",
    "layout": "horizontal",
    "examples": [
        "cju5eftctcdbj08712gdp989f.jpg",
        "cju42qet0lsq90871e50xbnuv.jpg",
        "cju8b0jr0r2oi0801jiquetd5.jpg"
    ],
    "allow_flagging": "never"
}

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

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

demo.launch(**launch_options)