Update app.py

app.py

@@ -12,16 +12,12 @@ model = U2NET(3, 1)
 
 # Load the state dictionary
 state_dict = torch.load(model_path, map_location=torch.device('cpu'))
-
-# Remove the 'module.' prefix from the keys
-state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
-
-# Load the modified state dictionary into the model
+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 from the image using U²-Net."""
+    """Segment the dress from the image using U²-Net and refine the mask."""
     transform_pipeline = transforms.Compose([
         transforms.ToTensor(),
         transforms.Resize((320, 320))
@@ -32,8 +28,16 @@ def segment_dress(image_np):
     with torch.no_grad():
         output = model(input_tensor)[0][0].squeeze().cpu().numpy()
     
-    mask = (output > 0.5).astype(np.uint8) * 255  # Thresholding for binary mask
-    mask = cv2.resize(mask, (image_np.shape[1], image_np.shape[0]))  # Resize mask to original
+    mask = (output > 0.5).astype(np.uint8) * 255  # Binary mask
+    
+    # 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)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)  # Close small gaps
+    mask = cv2.dilate(mask, kernel, iterations=2)  # Expand the detected dress area
+    
     return mask
 
 def change_dress_color(image_path, color):
@@ -44,7 +48,7 @@ def change_dress_color(image_path, color):
     img = Image.open(image_path).convert("RGB")
     img_np = np.array(img)
     mask = segment_dress(img_np)
-    
+
     if mask is None:
         return img  # No dress detected
     
@@ -58,19 +62,17 @@ 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
     
-    # Convert image to HSV
-    img_hsv = cv2.cvtColor(img_np, cv2.COLOR_RGB2HSV)
-
-    # Convert new color to HSV
-    new_color_hsv = cv2.cvtColor(np.uint8([[new_color_bgr]]), cv2.COLOR_BGR2HSV)[0][0]
+    # 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]
     
-    # Apply the new color only to the masked dress area
-    img_hsv[..., 0] = np.where(mask == 255, new_color_hsv[0], img_hsv[..., 0])  # Hue
-    img_hsv[..., 1] = np.where(mask == 255, new_color_hsv[1], img_hsv[..., 1])  # Saturation
+    # Apply the new color while preserving texture
+    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_hsv, cv2.COLOR_HSV2RGB)
-    
+    img_recolored = cv2.cvtColor(img_lab, cv2.COLOR_LAB2RGB)
+
     return Image.fromarray(img_recolored)
 
 # Gradio Interface