import gradio as gr
from SegCloth import segment_clothing

def segment(img, clothes, face, hair, skin):
    selected_items = clothes + face + hair + skin  # Combine all selections into one list
    return segment_clothing(img, selected_items)

iface = gr.Interface(
    fn=segment, 
    inputs=[
        gr.Image(type='pil', label='Image'),
        
        # Clothes Category
        gr.CheckboxGroup(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"],
                         label='Clothes'),
        
        # Face Category
        gr.CheckboxGroup(choices=["Face", "Sunglasses"],
                         value=["Face", "Sunglasses"],
                         label='Face'),
        
        # Hair Category
        gr.CheckboxGroup(choices=["Hair"],
                         value=["Hair"],
                         label='Hair'),
        
        # Skin Category
        gr.CheckboxGroup(choices=["Left-arm", "Right-arm", "Left-leg", "Right-leg"],
                         value=["Left-arm", "Right-arm", "Left-leg", "Right-leg"],
                         label='Skin')
    ], 
    outputs=gr.Image(label='Clothing Crop'),
    examples=[['./1.jpg', ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]]],
    theme="Nymbo/Alyx_Theme"
)

iface.launch()