Create app.py

app.py

@@ -0,0 +1,61 @@
+import torch
+import cv2
+import numpy as np
+import gradio as gr
+from PIL import Image
+from torchvision import transforms
+from cloth_segmentation.networks.u2net import U2NET  # Import U²-Net model
+
+# Load U²-Net model
+model_path = "u2net_model/u2net.pth"
+model = U2NET(3, 1)
+state_dict = torch.load(model_path, map_location=torch.device('cpu'))
+state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}  # Remove 'module.' prefix
+model.load_state_dict(state_dict)
+model.eval()
+
+def segment_dress(image_np):
+    """Detects dress using U²-Net and creates a binary mask."""
+    
+    # Convert image to tensor
+    transform_pipeline = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Resize((320, 320))
+    ])
+    
+    image = Image.fromarray(image_np).convert("RGB")
+    input_tensor = transform_pipeline(image).unsqueeze(0)
+
+    # U²-Net inference
+    with torch.no_grad():
+        output = model(input_tensor)[0][0].squeeze().cpu().numpy()
+
+    # Generate binary mask
+    dress_mask = (output > 0.5).astype(np.uint8) * 255
+    dress_mask = cv2.resize(dress_mask, (image_np.shape[1], image_np.shape[0]), interpolation=cv2.INTER_NEAREST)
+    
+    return dress_mask
+
+def remove_background(image_np):
+    """Removes background and replaces it with white while keeping the dress."""
+    
+    # Generate dress mask
+    mask = segment_dress(image_np)
+
+    # Make background white
+    white_bg = np.ones_like(image_np) * 255  # White background
+    segmented_dress = np.where(mask[..., None] > 128, image_np, white_bg)
+
+    return Image.fromarray(segmented_dress)
+
+# Gradio Interface
+demo = gr.Interface(
+    fn=remove_background,
+    inputs=gr.Image(type="numpy", label="Upload Dress Image"),
+    outputs=gr.Image(type="pil", label="Dress with White Background"),
+    title="Dress Segmentation & Background Removal",
+    description="Upload a dress image, and this AI model will detect the dress and replace the background with white."
+)
+
+if __name__ == "__main__":
+    demo.launch()