Update app.py

app.py

@@ -12,11 +12,12 @@ def get_segmentation_mask(image_url):
     result = client.predict(image=handle_file(image_url), model_name="1b", api_name="/process_image")
     return np.load(result[1])  # Result[1] contains the .npy mask
 
+
 def process_image(image, categories_to_hide):
     # Convert uploaded image to a PIL Image
     image = Image.open(image.name).convert("RGBA")
     
-    # Save temporarily and get the mask
+    # Save temporarily and get the segmentation mask
     image.save("temp_image.png")
     mask_data = get_segmentation_mask("temp_image.png")
     
@@ -24,7 +25,7 @@ def process_image(image, categories_to_hide):
     grouped_mapping = {
         "Background": [0],
         "Clothes": [1, 12, 22, 8, 9, 17, 18],  # Includes Shoes, Socks, Slippers
-        "Face": [2, 23, 24, 25, 26, 27],  # Face Neck, Lips, Teeth, Tongue
+        "Face": [2, 23, 24, 25, 26, 27],  # Face, Neck, Lips, Teeth, Tongue
         "Hair": [3],  # Hair
         "Skin": [4, 5, 6, 7, 10, 11, 13, 14, 15, 16, 19, 20, 21]  # Hands, Feet, Arms, Legs, Torso
     }
@@ -35,12 +36,28 @@ def process_image(image, categories_to_hide):
     # 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
+    # Create a binary mask for selected categories
     mask_combined = np.zeros_like(mask_data, dtype=bool)
+
     for category in categories_to_hide:
         for idx in grouped_mapping.get(category, []):
             mask_combined |= (mask_data == idx)
-    
+
+    # Expand clothing boundaries if clothes are in `categories_to_hide`
+    if "Clothes" in categories_to_hide:
+        clothing_mask = np.isin(mask_data, grouped_mapping["Clothes"]).astype(np.uint8)
+
+        # Determine kernel size (5% of the smaller image dimension)
+        height, width = clothing_mask.shape
+        kernel_size = max(1, int(0.05 * min(height, width)))  # Ensure at least 1 pixel
+        kernel = np.ones((kernel_size, kernel_size), np.uint8)
+
+        # Dilate the clothing mask
+        dilated_clothing_mask = cv2.dilate(clothing_mask, kernel, iterations=1)
+
+        # Update mask_combined with the expanded clothing mask
+        mask_combined |= (dilated_clothing_mask == 1)
+
     # Apply the mask (preserve only selected regions)
     transparent_image[mask_combined] = image_array[mask_combined]