Update app.py

app.py

@@ -14,7 +14,7 @@ def get_segmentation_mask(image_url):
 
 def process_image(image, categories_to_hide):
     # Convert uploaded image to a PIL Image
-    image = Image.open(image.name).convert("RGB")
+    image = Image.open(image.name).convert("RGBA")
     
     # Save temporarily and get the mask
     image.save("temp_image.png")
@@ -29,17 +29,23 @@ def process_image(image, categories_to_hide):
         "Skin (Hands, Feet, Body)": [4, 5, 6, 7, 10, 11, 13, 14, 15, 16, 19, 20, 21]  # Hands, Feet, Arms, Legs, Torso
     }
 
-    # Apply the mask over the original image
-    image_array = np.array(image)
-    masked_image = image_array.copy()
+    # Convert image to numpy array (RGBA)
+    image_array = np.array(image, dtype=np.uint8)
     
-    # Black out selected categories
+    # Create an empty transparent image
+    transparent_image = np.zeros_like(image_array, dtype=np.uint8)
+    
+    # Preserve only the selected mask regions, make everything else transparent
+    mask_combined = np.zeros_like(mask_data, dtype=bool)
     for category in categories_to_hide:
         for idx in grouped_mapping.get(category, []):
-            masked_image[mask_data == idx] = [0, 0, 0]
+            mask_combined |= (mask_data == idx)
+    
+    # Apply the mask (preserve only selected regions)
+    transparent_image[mask_combined] = image_array[mask_combined]
     
     # Convert back to PIL Image
-    result_image = Image.fromarray(masked_image)
+    result_image = Image.fromarray(transparent_image, mode="RGBA")
     
     return result_image
 
@@ -50,11 +56,11 @@ demo = gr.Interface(
         gr.File(label="Upload an Image"),
         gr.CheckboxGroup([
             "Background", "Clothes", "Face", "Hair", "Skin (Hands, Feet, Body)"
-        ], label="Select Categories to Hide")
+        ], label="Select Categories to Preserve")
     ],
-    outputs=gr.Image(label="Masked Image"),
+    outputs=gr.Image(label="Masked Image", type="pil"),
     title="Segmentation Mask Editor",
-    description="Upload an image, generate a segmentation mask, and select categories to black out."
+    description="Upload an image, generate a segmentation mask, and select categories to preserve while making the rest transparent."
 )
 
 if __name__ == "__main__":