Create app.py

app.py

@@ -0,0 +1,36 @@
+import gradio as gr
+from transformers import pipeline
+from PIL import Image
+import numpy as np
+
+# Load U-2-Net segmentation pipeline
+pipe = pipeline("image-segmentation", model="BritishWerewolf/U-2-Net")
+
+# Segmentation function for Gradio
+def segment_dress(image: Image.Image):
+    # Run U-2-Net pipeline
+    segments = pipe(image)
+
+    # We'll assume the first segment is the foreground (person+clothes)
+    if not segments:
+        return image
+
+    # Load and apply mask
+    mask = Image.open(segments[0]["mask"]).convert("L").resize(image.size)
+    mask_np = np.array(mask) / 255.0
+    image_np = np.array(image).astype(np.uint8)
+
+    # Apply mask to image (keep only masked region)
+    segmented = (image_np * mask_np[..., None]).astype(np.uint8)
+    segmented_img = Image.fromarray(segmented)
+
+    return segmented_img
+
+# Gradio Interface
+gr.Interface(
+    fn=segment_dress,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=gr.Image(type="pil", label="Segmented Dress Region"),
+    title="Dress Segmentation using U-2-Net",
+    description="Upload an image. The U-2-Net model will segment the main foreground (usually a person wearing a dress)."
+).launch()