import gradio as gr
from SegCloth import segment_clothing

# Define items for each category, including Background
CATEGORY_ITEMS = {
    "Background": ["Background"],
    "Clothes": ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"],
    "Face": ["Face", "Sunglasses"],
    "Hair": ["Hair"],
    "Skin": ["Left-arm", "Right-arm", "Left-leg", "Right-leg"]
}

def segment(img, selected_categories):
    # Get the items based on the selected categories
    selected_items = []
    for category in selected_categories:
        selected_items.extend(CATEGORY_ITEMS.get(category, []))
    
    return segment_clothing(img, selected_items)

iface = gr.Interface(
    fn=segment, 
    inputs=[
        gr.Image(type='pil', label='Image'),
        
        # Category Selection CheckboxGroup including Background
        gr.CheckboxGroup(choices=["Background", "Clothes", "Face", "Hair", "Skin"],
                         label='Select Categories',
                         value=["Background", "Clothes", "Face", "Hair", "Skin"])  # Default selection
    ], 
    outputs=gr.Image(label='Clothing Crop'),
    examples=[['./1.jpg', ["Background", "Clothes", "Face"]]]
)

iface.launch()