Update app.py

app.py

@@ -52,37 +52,45 @@ def segment_dress(image_np):
     # 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)
-    
     return refined_mask
 
-def recolor_dress(image_np, mask, target_color):
-    """Change dress color while preserving texture and shadows."""
+def detect_design(image_np, dress_mask):
+    """Detect the design part of the dress and separate it from fabric."""
+    gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
+    edges = cv2.Canny(gray, 50, 150)
     
+    # Expand detected edges to mask the design area
+    kernel = np.ones((5, 5), np.uint8)
+    design_mask = cv2.dilate(edges, kernel, iterations=2)
+    
+    # Keep only the design within the dress area
+    design_mask = cv2.bitwise_and(design_mask, dress_mask)
+    return design_mask
+
+def recolor_dress(image_np, mask, design_mask, target_color):
+    """Change dress color while preserving texture and design."""
     img_lab = cv2.cvtColor(image_np, cv2.COLOR_RGB2LAB)
     target_color_lab = cv2.cvtColor(np.uint8([[target_color]]), cv2.COLOR_BGR2LAB)[0][0]
     
     # Preserve lightness (L) and change only chromatic channels (A & B)
     blend_factor = 0.7
-    img_lab[..., 1] = np.where(mask > 128, img_lab[..., 1] * (1 - blend_factor) + target_color_lab[1] * blend_factor, img_lab[..., 1])
-    img_lab[..., 2] = np.where(mask > 128, img_lab[..., 2] * (1 - blend_factor) + target_color_lab[2] * blend_factor, img_lab[..., 2])
+    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
 
 def change_dress_color(image_path, color):
-    """Change the dress color naturally while keeping textures."""
+    """Change the dress color naturally while keeping textures and design."""
     if image_path is None:
         return None
 
     img = Image.open(image_path).convert("RGB")
     img_np = np.array(img)
-    mask = segment_dress(img_np)
+    dress_mask = segment_dress(img_np)
+    design_mask = detect_design(img_np, dress_mask)
     
-    if mask is None:
+    if dress_mask is None:
         return img  # No dress detected
     
     # Convert the selected color to BGR
@@ -94,7 +102,7 @@ def change_dress_color(image_path, color):
     new_color_bgr = np.array(color_map.get(color, (0, 0, 255)), dtype=np.uint8)  # Default to Red
     
     # Recolor the dress naturally
-    img_recolored = recolor_dress(img_np, mask, new_color_bgr)
+    img_recolored = recolor_dress(img_np, dress_mask, design_mask, new_color_bgr)
 
     return Image.fromarray(img_recolored)
 
@@ -107,7 +115,7 @@ demo = gr.Interface(
     ],
     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."
+    description="Upload an image of a dress and select a new color to change its appearance naturally while preserving the design."
 )
 
 if __name__ == "__main__":