Create app.py

app.py

@@ -0,0 +1,67 @@
+import gradio as gr
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
+
+def process_mask(file, category_to_hide):
+    # Load the .npy file
+    data = np.load(file.name)
+    
+    # Define grouped categories
+    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
+        "Hair": [3],  # Hair
+        "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"
+    }
+    
+    # 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()
+    
+    # Save to file for Gradio output
+    output_path = "output_mask.png"
+    plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
+    plt.close()
+    
+    return output_path
+
+# Define Gradio Interface
+demo = gr.Interface(
+    fn=process_mask,
+    inputs=[
+        gr.File(label="Upload .npy Segmentation File"),
+        gr.Radio([
+            "Background", "Clothes", "Face", "Hair", "Skin (Hands, Feet, Body)"
+        ], label="Select Category to Hide")
+    ],
+    outputs=gr.Image(label="Modified Segmentation Mask"),
+    title="Segmentation Mask Editor",
+    description="Upload a .npy segmentation file and select a category to mask (hide with black)."
+)
+
+if __name__ == "__main__":
+    demo.launch()