import gradio as gr
from SegCloth import segment_clothing

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

custom_css = """
/* Remove box shadow and border from example images */
.gr-example-image {
    box-shadow: none !important;
    border: none !important;
}

/* Center example images */
.gr-example-image-wrapper {
    display: flex;
    justify-content: center;
}
"""

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"),
    css=custom_css
)

iface.launch()