Update app.py

app.py

@@ -10,9 +10,9 @@ from io import BytesIO
 def get_segmentation_mask(image_url):
     client = Client("facebook/sapiens-seg")
     result = client.predict(image=handle_file(image_url), model_name="1b", api_name="/process_image")
-    return np.load(result[1])  # Result[2] contains the .npy mask
+    return np.load(result[2])  # Result[2] contains the .npy mask
 
-def process_image(image, category_to_hide):
+def process_image(image, categories_to_hide):
     # Convert uploaded image to a PIL Image
     image = Image.open(image.name).convert("RGB")
     
@@ -33,9 +33,10 @@ def process_image(image, category_to_hide):
     image_array = np.array(image)
     masked_image = image_array.copy()
     
-    # Black out selected category
-    for idx in grouped_mapping[category_to_hide]:
-        masked_image[mask_data == idx] = [0, 0, 0]
+    # Black out selected categories
+    for category in categories_to_hide:
+        for idx in grouped_mapping.get(category, []):
+            masked_image[mask_data == idx] = [0, 0, 0]
     
     # Convert back to PIL Image
     result_image = Image.fromarray(masked_image)
@@ -47,13 +48,13 @@ demo = gr.Interface(
     fn=process_image,
     inputs=[
         gr.File(label="Upload an Image"),
-        gr.Radio([
+        gr.CheckboxGroup([
             "Background", "Clothes", "Face", "Hair", "Skin (Hands, Feet, Body)"
-        ], label="Select Category to Hide")
+        ], label="Select Categories to Hide")
     ],
     outputs=gr.Image(label="Masked Image"),
     title="Segmentation Mask Editor",
-    description="Upload an image, generate a segmentation mask, and select a category to black out."
+    description="Upload an image, generate a segmentation mask, and select categories to black out."
 )
 
 if __name__ == "__main__":