import gradio as gr
from SegCloth import segment_clothing

def segment(img, clothes):
    return segment_clothing(img, clothes)

iface = gr.Interface(fn=segment, 
                     inputs=[gr.Image(type='pil', label='Image'),
                             gr.Dropdown(choices=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"],
                                         value=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"],
                                         multiselect=True,
                                         label='Clothing')], 
                     outputs=gr.Image(label='Clothes Extractor'),
                     title='Clothes Extractor',
                     description="Upload an image of a person wearing clothing and select below which piece of clothing to extract from the image.",
                     examples=[
                         ['./1.jpg', ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]],
                         ['./Model 1.jpg', ["Pants", "Upper-clothes"]],
                         ['./Model 2.jpg', ["Pants", "Upper-clothes"]],
                         ['./Model 3.jpg', ["Pants", "Upper-clothes"]],
                         ['./Model 4.jpg', ["Pants", "Upper-clothes"]],
                         ['./Model 5.jpg', ["Pants", "Upper-clothes"]],
                         ['./Model 6.jpg', ["Pants", "Upper-clothes"]],
                         ['./Model 7.jpg', ["Pants", "Upper-clothes"]],  
                     ],
                     theme=gr.themes.Base(primary_hue="teal", secondary_hue="teal", neutral_hue="slate"))

iface.launch()