Update app.py

app.py

@@ -7,103 +7,74 @@ 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"
+model_path = "cloth_segmentation/networks/u2net.pth"  # Ensure this path is correct
 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 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
+    """Segment the dress from the image using U²-Net and refine the mask."""
     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)
     
-    # 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)
-    refined_mask = cv2.morphologyEx(refined_mask, cv2.MORPH_CLOSE, kernel)
-    refined_mask = cv2.GaussianBlur(refined_mask, (15, 15), 5)
+    mask = (output > 0.5).astype(np.uint8) * 255  # Binary mask
     
-    return refined_mask
-
-def detect_design(image_np):
-    """Detects design patterns on the dress using edge detection."""
-    gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
-    edges = cv2.Canny(gray, 50, 150)
-    kernel = np.ones((3, 3), np.uint8)
-    design_mask = cv2.dilate(edges, kernel, iterations=2)
-    return design_mask
-
-def recolor_dress(image_np, mask, design_mask, target_color):
-    """Change dress color while preserving texture and design."""
+    # Resize mask to original image size
+    mask = cv2.resize(mask, (image_np.shape[1], image_np.shape[0]), interpolation=cv2.INTER_NEAREST)
     
-    img_lab = cv2.cvtColor(image_np, cv2.COLOR_RGB2LAB)
-    target_color_lab = cv2.cvtColor(np.uint8([[target_color]]), cv2.COLOR_BGR2LAB)[0][0]
+    # Apply morphological operations for better segmentation
+    kernel = np.ones((7, 7), np.uint8)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)  # Close small gaps
+    mask = cv2.dilate(mask, kernel, iterations=2)  # Expand the detected dress area
     
-    # Preserve lightness (L) and change only chromatic channels (A & B)
-    blend_factor = 0.7
-    img_lab[..., 1] = np.where((mask > 128) & (design_mask == 0), img_lab[..., 1] * (1 - blend_factor) + target_color_lab[1] * blend_factor, img_lab[..., 1])
-    img_lab[..., 2] = np.where((mask > 128) & (design_mask == 0), img_lab[..., 2] * (1 - blend_factor) + target_color_lab[2] * blend_factor, img_lab[..., 2])
-
-    img_recolored = cv2.cvtColor(img_lab, cv2.COLOR_LAB2RGB)
-    return img_recolored
+    return mask
 
 def change_dress_color(image_path, color):
-    """Change the dress color naturally while keeping textures and design."""
+    """Change the dress color naturally while keeping textures."""
     if image_path is None:
         return None
 
     img = Image.open(image_path).convert("RGB")
     img_np = np.array(img)
     mask = segment_dress(img_np)
-    design_mask = detect_design(img_np)
-    
+
     if mask is None:
         return img  # No dress detected
     
     # Convert the selected color to BGR
     color_map = {
-        "Red": (0, 0, 255), "Blue": (255, 0, 0), "Green": (0, 255, 0), "Yellow": (0, 255, 255),
-        "Purple": (128, 0, 128), "Orange": (0, 165, 255), "Cyan": (255, 255, 0), "Magenta": (255, 0, 255),
-        "White": (255, 255, 255), "Black": (0, 0, 0)
+        "Red": (0, 0, 255),
+        "Blue": (255, 0, 0),
+        "Green": (0, 255, 0),
+        "Yellow": (0, 255, 255),
+        "Purple": (128, 0, 128)
     }
     new_color_bgr = np.array(color_map.get(color, (0, 0, 255)), dtype=np.uint8)  # Default to Red
     
-    # Recolor the dress naturally while preserving design
-    img_recolored = recolor_dress(img_np, mask, design_mask, new_color_bgr)
+    # Convert image to LAB color space for better blending
+    img_lab = cv2.cvtColor(img_np, cv2.COLOR_RGB2LAB)
+    new_color_lab = cv2.cvtColor(np.uint8([[new_color_bgr]]), cv2.COLOR_BGR2LAB)[0][0]
+    
+    # Preserve texture by only modifying the A & B channels
+    img_lab[..., 1] = np.where(mask == 255, new_color_lab[1], img_lab[..., 1])  # Modify A-channel
+    img_lab[..., 2] = np.where(mask == 255, new_color_lab[2], img_lab[..., 2])  # Modify B-channel
+
+    # Convert back to RGB
+    img_recolored = cv2.cvtColor(img_lab, cv2.COLOR_LAB2RGB)
+
+    # Apply Poisson blending for realistic color application
+    img_recolored = cv2.seamlessClone(img_recolored, img_np, mask, (img_np.shape[1]//2, img_np.shape[0]//2), cv2.NORMAL_CLONE)
 
     return Image.fromarray(img_recolored)
 
@@ -112,11 +83,11 @@ demo = gr.Interface(
     fn=change_dress_color,
     inputs=[
         gr.Image(type="filepath", label="Upload Dress Image"),
-        gr.Radio(["Red", "Blue", "Green", "Yellow", "Purple", "Orange", "Cyan", "Magenta", "White", "Black"], label="Choose New Dress Color")
+        gr.Radio(["Red", "Blue", "Green", "Yellow", "Purple"], label="Choose New Dress Color")
     ],
     outputs=gr.Image(type="pil", label="Color Changed Dress"),
     title="Dress Color Changer",
-    description="Upload an image of a dress and select a new color to change its appearance naturally while preserving any design patterns."
+    description="Upload an image of a dress and select a new color to change its appearance naturally."
 )
 
 if __name__ == "__main__":