Update app.py

app.py

@@ -1,36 +1,62 @@
 import gradio as gr
-from transformers import pipeline
+import torch
+import torch.nn as nn
+from torchvision import transforms
 from PIL import Image
 import numpy as np
+import requests
+from io import BytesIO
 
-# Load U-2-Net segmentation pipeline
-pipe = pipeline("image-segmentation", model="BritishWerewolf/U-2-Net")
+# U-2-Net architecture (simplified, or import from a .py file if you've saved it)
+# You can get the U-2-Net code from https://github.com/xuebinqin/U-2-Net
 
-# Segmentation function for Gradio
-def segment_dress(image: Image.Image):
-    # Run U-2-Net pipeline
-    segments = pipe(image)
+# For demo, let's download the pre-trained model and use a wrapper instead
+from huggingface_hub import hf_hub_download
 
-    # We'll assume the first segment is the foreground (person+clothes)
-    if not segments:
-        return image
+# Download u2net.pth from HuggingFace Hub
+model_path = hf_hub_download(repo_id="BritishWerewolf/U-2-Net", filename="u2net.pth")
 
-    # 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)
+# Use a known U2NET implementation (e.g., from https://github.com/xuebinqin/U-2-Net/blob/master/u2net_test.py)
+from u2net import U2NET  # Assume you copied the model code as u2net.py
 
-    # Apply mask to image (keep only masked region)
-    segmented = (image_np * mask_np[..., None]).astype(np.uint8)
-    segmented_img = Image.fromarray(segmented)
+# Load model
+model = U2NET(3, 1)
+model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+model.eval()
 
-    return segmented_img
+# Preprocessing
+transform = transforms.Compose([
+    transforms.Resize((320, 320)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225])
+])
 
-# Gradio Interface
+def segment_dress(image):
+    original = image.convert("RGB")
+    input_tensor = transform(original).unsqueeze(0)
+    
+    with torch.no_grad():
+        d1, _, _, _, _, _, _ = model(input_tensor)
+        pred = d1[0][0]
+        pred = (pred - pred.min()) / (pred.max() - pred.min())
+        pred_np = pred.cpu().numpy()
+    
+    # Resize to original size
+    pred_resized = Image.fromarray((pred_np * 255).astype(np.uint8)).resize(original.size)
+    
+    # Apply mask
+    mask = np.array(pred_resized) / 255.0
+    image_np = np.array(original).astype(np.uint8)
+    segmented = (image_np * mask[..., None]).astype(np.uint8)
+
+    return Image.fromarray(segmented)
+
+# Launch Gradio app
 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)."
+    outputs=gr.Image(type="pil", label="Segmented Dress"),
+    title="Dress Segmentation with U-2-Net",
+    description="Segments the dress (or full foreground) using U-2-Net from Hugging Face"
 ).launch()