Update app.py

app.py

@@ -14,6 +14,14 @@ state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
 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)
+
 def segment_dress(image_np):
     """Segment dress using U²-Net"""
     transform_pipeline = transforms.Compose([
@@ -30,17 +38,30 @@ def segment_dress(image_np):
     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_LINEAR)
 
-    return dress_mask
+    return refine_mask(dress_mask)
+
+def apply_grabcut(image_np, dress_mask):
+    """Refines the mask using GrabCut to avoid color bleeding"""
+    bgd_model = np.zeros((1, 65), np.float64)
+    fgd_model = np.zeros((1, 65), np.float64)
+
+    mask = np.where(dress_mask > 0, cv2.GC_FGD, cv2.GC_BGD).astype('uint8')
+    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):
-    """Change dress color naturally while keeping textures intact"""
+    """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]
 
-    # Convert image to LAB for better color control
+    # Convert image to LAB
     img_lab = cv2.cvtColor(image_np, cv2.COLOR_RGB2LAB)
-    
+
     # Compute mean LAB values of dress pixels
     dress_pixels = img_lab[dress_mask > 0]
     if len(dress_pixels) == 0:
@@ -48,21 +69,20 @@ def recolor_dress(image_np, dress_mask, target_color):
 
     mean_L, mean_A, mean_B = np.mean(dress_pixels, axis=0)
 
-    # Adjust LAB channels for the new color
+    # 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])
 
     # Convert back to RGB
     img_recolored = cv2.cvtColor(img_lab, cv2.COLOR_LAB2RGB)
 
-    # Convert mask to 3 channels for seamless blending
-    mask_bgr = cv2.cvtColor(dress_mask, cv2.COLOR_GRAY2BGR)
+    # Create feathered mask for smooth blending
+    feathered_mask = cv2.GaussianBlur(dress_mask, (15, 15), 10)
 
-    # Smooth edges for natural blending using seamlessClone
-    center = (image_np.shape[1] // 2, image_np.shape[0] // 2)
-    blended = cv2.seamlessClone(img_recolored, image_np, mask_bgr, center, cv2.NORMAL_CLONE)
+    # 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)
 
-    return blended
+    return img_final
 
 def change_dress_color(image_path, color):
     """Main function to change dress color naturally"""
@@ -78,6 +98,9 @@ def change_dress_color(image_path, color):
     if dress_mask is None:
         return img  # No dress detected
 
+    # Further refine mask with GrabCut
+    dress_mask = apply_grabcut(img_np, dress_mask)
+
     # Convert the selected color to BGR
     color_map = {
         "Red": (0, 0, 255), "Blue": (255, 0, 0), "Green": (0, 255, 0), "Yellow": (0, 255, 255),
@@ -99,7 +122,7 @@ demo = gr.Interface(
         gr.Radio(["Red", "Blue", "Green", "Yellow", "Purple", "Orange", "Cyan", "Magenta", "White", "Black"], label="Choose New Dress Color")
     ],
     outputs=gr.Image(type="pil", label="Color Changed Dress"),
-    title="Realistic Dress Color Changer",
+    title="AI-Powered Dress Color Changer",
     description="Upload an image of a dress and select a new color. The AI will change the dress color naturally while keeping the fabric texture."
 )