import gradio as gr
from pred_color import *

example_sample = [
    "babyxiang_ai.png"
]

def pred_func(img):
    out = single_pred_features(img)
    if type(out) == type({}):
        return out["spiga_seg"]

gr=gr.Interface(fn=pred_func, inputs=['image',],
outputs=[gr.Image(label='output').style(height=512)],
examples=example_sample if example_sample else None,
)
gr.launch(share=False)