Update app.py

app.py

@@ -1,12 +1,20 @@
 import gradio as gr
 from SegCloth import segment_clothing
 
-def segment(img, clothes, face, hair, skin):
-    clothes = clothes if clothes is not None else []
-    face = face if face is not None else []
-    hair = hair if hair is not None else []
-    skin = skin if skin is not None else []
-    selected_items = clothes + face + hair + skin  # Combine all selections into one list
+# Define items for each category
+CATEGORY_ITEMS = {
+    "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(
@@ -14,28 +22,13 @@ iface = gr.Interface(
     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')
+        # Category Selection CheckboxGroup
+        gr.CheckboxGroup(choices=["Clothes", "Face", "Hair", "Skin"],
+                         label='Select Categories',
+                         value=["Clothes", "Face", "Hair", "Skin"])  # Default selection
     ], 
     outputs=gr.Image(label='Clothing Crop'),
-    examples=[['./1.jpg', ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]]],
+    examples=[['./1.jpg', ["Clothes", "Face"]]],
     theme="Nymbo/Alyx_Theme"
 )