Update app.py

app.py

@@ -7,41 +7,61 @@ from torchvision import transforms
 from cloth_segmentation.networks.u2net import U2NET  # Import U²-Net
 
 # Load U²-Net model
-model_path = "cloth_segmentation/networks/u2net.pth"  # Ensure this path is correct
+model_path = "cloth_segmentation/networks/u2net.pth"
 model = U2NET(3, 1)
-
-# Load the state dictionary
 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):
-    """Segment the dress from the image using U²-Net and refine the mask."""
+    """Segment the dress using U²-Net & refine with Lab color space."""
+    
+    # Convert to Lab space
+    img_lab = cv2.cvtColor(image_np, cv2.COLOR_RGB2LAB)
+    L, A, B = cv2.split(img_lab)
+
+    # Use K-means clustering to detect dominant dress region
+    pixel_values = img_lab.reshape((-1, 3)).astype(np.float32)
+    k = 3  # Three clusters: background, skin, dress
+    _, labels, centers = cv2.kmeans(pixel_values, k, None, (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0), 10, cv2.KMEANS_RANDOM_CENTERS)
+    labels = labels.reshape(image_np.shape[:2])
+    
+    # Assume dress is the largest non-background cluster
+    unique_labels, counts = np.unique(labels, return_counts=True)
+    dress_label = unique_labels[np.argmax(counts[1:]) + 1]  # Avoid background
+    
+    # Create dress mask
+    mask = (labels == dress_label).astype(np.uint8) * 255
+
+    # Use U²-Net prediction to refine segmentation
     transform_pipeline = transforms.Compose([
         transforms.ToTensor(),
         transforms.Resize((320, 320))
     ])
+    
     image = Image.fromarray(image_np).convert("RGB")
     input_tensor = transform_pipeline(image).unsqueeze(0)
-    
+
     with torch.no_grad():
         output = model(input_tensor)[0][0].squeeze().cpu().numpy()
+
+    u2net_mask = (output > 0.5).astype(np.uint8) * 255
+    u2net_mask = cv2.resize(u2net_mask, (image_np.shape[1], image_np.shape[0]), interpolation=cv2.INTER_NEAREST)
     
-    mask = (output > 0.5).astype(np.uint8) * 255  # Binary mask
-    
-    # Resize mask to original image size
-    mask = cv2.resize(mask, (image_np.shape[1], image_np.shape[0]), interpolation=cv2.INTER_NEAREST)
-    
-    # Apply morphological operations for better segmentation
+    # Combine K-means and U²-Net masks
+    refined_mask = cv2.bitwise_and(mask, u2net_mask)
+
+    # Morphological operations for smoothness
     kernel = np.ones((5, 5), np.uint8)
-    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)  # Close small gaps
-    mask = cv2.GaussianBlur(mask, (15, 15), 5)  # Smooth edges for natural blending
+    refined_mask = cv2.morphologyEx(refined_mask, cv2.MORPH_CLOSE, kernel)
+    refined_mask = cv2.GaussianBlur(refined_mask, (15, 15), 5)
     
-    return mask
+    return refined_mask
 
 def recolor_dress(image_np, mask, target_color):
     """Change dress color while preserving texture and shadows."""
+    
     img_lab = cv2.cvtColor(image_np, cv2.COLOR_RGB2LAB)
     target_color_lab = cv2.cvtColor(np.uint8([[target_color]]), cv2.COLOR_BGR2LAB)[0][0]