Update app.py

app.py

@@ -2,10 +2,23 @@ import gradio as gr
 import numpy as np
 import matplotlib.pyplot as plt
 import matplotlib.colors as mcolors
+from gradio_client import Client, handle_file
+from PIL import Image
+import requests
+from io import BytesIO
 
-def process_mask(file, category_to_hide):
-    # Load the .npy file
-    data = np.load(file.name)
+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[2])  # Result[2] contains the .npy mask
+
+def process_image(image, category_to_hide):
+    # Convert uploaded image to a PIL Image
+    image = Image.open(image.name).convert("RGB")
+    
+    # Save temporarily and get the mask
+    image.save("temp_image.png")
+    mask_data = get_segmentation_mask("temp_image.png")
     
     # Define grouped categories
     grouped_mapping = {
@@ -16,51 +29,31 @@ def process_mask(file, category_to_hide):
         "Skin (Hands, Feet, Body)": [4, 5, 6, 7, 10, 11, 13, 14, 15, 16, 19, 20, 21]  # Hands, Feet, Arms, Legs, Torso
     }
 
-    # Assign colors for the categories
-    group_colors = {
-        "Background": "black",
-        "Clothes": "magenta",
-        "Face": "orange",
-        "Hair": "brown",
-        "Skin (Hands, Feet, Body)": "cyan"
-    }
+    # Apply the mask over the original image
+    image_array = np.array(image)
+    masked_image = image_array.copy()
     
-    # Create a new mask with grouped categories
-    grouped_mask = np.zeros((*data.shape, 3), dtype=np.uint8)
-
-    for category, indices in grouped_mapping.items():
-        if category == category_to_hide:
-            continue  # Skip applying colors for the selected category to hide
-        for idx in indices:
-            mask = data == idx
-            rgb = mcolors.to_rgb(group_colors[category])  # Convert color to RGB
-            grouped_mask[mask] = [int(c * 255) for c in rgb]
-
-    # Save the mask image
-    fig, ax = plt.subplots(figsize=(6, 6))
-    ax.imshow(grouped_mask)
-    ax.axis("off")
-    plt.tight_layout()
+    # Black out selected category
+    for idx in grouped_mapping[category_to_hide]:
+        masked_image[mask_data == idx] = [0, 0, 0]
     
-    # Save to file for Gradio output
-    output_path = "output_mask.png"
-    plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
-    plt.close()
+    # Convert back to PIL Image
+    result_image = Image.fromarray(masked_image)
     
-    return output_path
+    return result_image
 
 # Define Gradio Interface
 demo = gr.Interface(
-    fn=process_mask,
+    fn=process_image,
     inputs=[
-        gr.File(label="Upload .npy Segmentation File"),
+        gr.File(label="Upload an Image"),
         gr.Radio([
             "Background", "Clothes", "Face", "Hair", "Skin (Hands, Feet, Body)"
         ], label="Select Category to Hide")
     ],
-    outputs=gr.Image(label="Modified Segmentation Mask"),
+    outputs=gr.Image(label="Masked Image"),
     title="Segmentation Mask Editor",
-    description="Upload a .npy segmentation file and select a category to mask (hide with black)."
+    description="Upload an image, generate a segmentation mask, and select a category to black out."
 )
 
 if __name__ == "__main__":