Update app.py

app.py

@@ -15,12 +15,11 @@ model.load_state_dict(state_dict)
 model.eval()
 
 def refine_mask(mask):
-    """Refines mask using morphological operations"""
-    kernel = np.ones((5, 5), np.uint8)
-    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
-    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
-    
-    return cv2.GaussianBlur(mask, (5, 5), 0)
+    """Refines mask using morphological closing followed by Gaussian blur"""
+    kernel = np.ones((7, 7), np.uint8)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)  # Close holes inside dress
+    mask = cv2.GaussianBlur(mask, (7, 7), 1.5)
+    return mask
 
 def segment_dress(image_np):
     """Segment dress using U²-Net"""
@@ -28,7 +27,7 @@ def segment_dress(image_np):
         transforms.ToTensor(),
         transforms.Resize((320, 320))
     ])
-    
+
     image = Image.fromarray(image_np).convert("RGB")
     input_tensor = transform_pipeline(image).unsqueeze(0)
 
@@ -49,13 +48,13 @@ def apply_grabcut(image_np, dress_mask):
     rect = (10, 10, image_np.shape[1] - 10, image_np.shape[0] - 10)
 
     cv2.grabCut(image_np, mask, rect, bgd_model, fgd_model, 5, cv2.GC_INIT_WITH_MASK)
-    
+
     refined_mask = np.where((mask == 2) | (mask == 0), 0, 255).astype("uint8")
     return refine_mask(refined_mask)
 
 def recolor_dress(image_np, dress_mask, target_color):
     """Changes dress color while keeping texture & lighting intact"""
-    
+
     # Convert target color to LAB
     target_color_lab = cv2.cvtColor(np.uint8([[target_color]]), cv2.COLOR_BGR2LAB)[0][0]
 
@@ -69,15 +68,17 @@ def recolor_dress(image_np, dress_mask, target_color):
 
     mean_L, mean_A, mean_B = np.mean(dress_pixels, axis=0)
 
-    # Apply histogram-based color transfer
-    img_lab[..., 1] = np.where(dress_mask > 128, img_lab[..., 1] - mean_A + target_color_lab[1], img_lab[..., 1])
-    img_lab[..., 2] = np.where(dress_mask > 128, img_lab[..., 2] - mean_B + target_color_lab[2], img_lab[..., 2])
+    # Apply histogram-based color transfer with topological adjustment
+    a_shift = target_color_lab[1] - mean_A
+    b_shift = target_color_lab[2] - mean_B
+    img_lab[..., 1] = np.clip(img_lab[..., 1] + (dress_mask / 255.0) * a_shift, 0, 255)
+    img_lab[..., 2] = np.clip(img_lab[..., 2] + (dress_mask / 255.0) * b_shift, 0, 255)
 
     # Convert back to RGB
-    img_recolored = cv2.cvtColor(img_lab, cv2.COLOR_LAB2RGB)
+    img_recolored = cv2.cvtColor(img_lab.astype(np.uint8), cv2.COLOR_LAB2RGB)
 
     # Create feathered mask for smooth blending
-    feathered_mask = cv2.GaussianBlur(dress_mask, (15, 15), 10)
+    feathered_mask = cv2.GaussianBlur(dress_mask, (15, 15), 5)
 
     # Blend the recolored dress with the original image
     img_final = (image_np * (1 - feathered_mask[..., None] / 255) + img_recolored * (feathered_mask[..., None] / 255)).astype(np.uint8)
@@ -94,7 +95,7 @@ def change_dress_color(image_path, color):
 
     # Get dress segmentation mask
     dress_mask = segment_dress(img_np)
-    
+
     if dress_mask is None:
         return img  # No dress detected