Update app.py

app.py

@@ -7,12 +7,12 @@ 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
+state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}  
 model.load_state_dict(state_dict)
 model.eval()
 
@@ -33,10 +33,11 @@ def segment_dress(image_np):
     # 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
-    kernel = np.ones((7, 7), np.uint8)
+    # Refine mask using morphological operations
+    kernel = np.ones((5, 5), 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
+    mask = cv2.GaussianBlur(mask, (5, 5), 0)  # Smooth edges
     
     return mask
 
@@ -52,29 +53,28 @@ def change_dress_color(image_path, color):
     if mask is None:
         return img  # No dress detected
     
-    # Convert the selected color to BGR
+    # Convert image to HSV for color modification
+    img_hsv = cv2.cvtColor(img_np, cv2.COLOR_RGB2HSV)
+
+    # Define new color in HSV (only modifying the Hue)
     color_map = {
-        "Red": (0, 0, 255),
-        "Blue": (255, 0, 0),
-        "Green": (0, 255, 0),
-        "Yellow": (0, 255, 255),
-        "Purple": (128, 0, 128)
+        "Red": 0,      # Hue value for Red
+        "Blue": 120,   # Hue value for Blue
+        "Green": 60,   # Hue value for Green
+        "Yellow": 30,  # Hue value for Yellow
+        "Purple": 150  # Hue value for Purple
     }
-    new_color_bgr = np.array(color_map.get(color, (0, 0, 255)), dtype=np.uint8)  # Default to Red
-    
-    # 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
+    new_hue = color_map.get(color, 0)
+
+    # Modify only the Hue channel where the mask is applied
+    img_hsv[..., 0] = np.where(mask > 128, new_hue, img_hsv[..., 0])
 
     # Convert back to RGB
-    img_recolored = cv2.cvtColor(img_lab, cv2.COLOR_LAB2RGB)
+    img_recolored = cv2.cvtColor(img_hsv, cv2.COLOR_HSV2RGB)
 
-    # 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)
+    # Apply Poisson blending for natural integration
+    center = (img_np.shape[1] // 2, img_np.shape[0] // 2)
+    img_recolored = cv2.seamlessClone(img_recolored, img_np, mask, center, cv2.MIXED_CLONE)
 
     return Image.fromarray(img_recolored)