Update app.py

app.py

@@ -33,11 +33,13 @@ 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)
     
-    # Refine the mask to reduce noise
-    kernel = np.ones((5, 5), np.uint8)
+    # 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=1)  # Expand the detected dress area
-    mask = cv2.medianBlur(mask, 5)  # Reduce noise and smooth edges
+    mask = cv2.dilate(mask, kernel, iterations=2)  # Expand the detected dress area
+    
+    # Apply Gaussian blur for smoother transitions
+    mask = cv2.GaussianBlur(mask, (15, 15), 5)
     
     return mask
 
@@ -49,17 +51,15 @@ 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
     
     # Convert the selected color to BGR
     color_map = {
-        "Red": (0, 0, 255),
-        "Blue": (255, 0, 0),
-        "Green": (0, 255, 0),
-        "Yellow": (0, 255, 255),
-        "Purple": (128, 0, 128)
+        "Red": (0, 0, 255), "Blue": (255, 0, 0), "Green": (0, 255, 0), "Yellow": (0, 255, 255),
+        "Purple": (128, 0, 128), "Orange": (0, 165, 255), "Cyan": (255, 255, 0), "Magenta": (255, 0, 255),
+        "White": (255, 255, 255)
     }
     new_color_bgr = np.array(color_map.get(color, (0, 0, 255)), dtype=np.uint8)  # Default to Red
     
@@ -67,29 +67,22 @@ def change_dress_color(image_path, color):
     img_lab = cv2.cvtColor(img_np, cv2.COLOR_RGB2LAB)
     new_color_lab = cv2.cvtColor(np.uint8([[new_color_bgr]]), cv2.COLOR_BGR2LAB)[0][0]
     
-    # Adaptive color blending - preserves details and avoids over-saturation
-    l_channel, a_channel, b_channel = cv2.split(img_lab)
-    
-    a_channel = np.where(mask == 255, new_color_lab[1], a_channel)  # Modify A-channel
-    b_channel = np.where(mask == 255, new_color_lab[2], b_channel)  # Modify B-channel
-    
-    img_recolored = cv2.merge([l_channel, a_channel, b_channel])
-    
-    # Convert back to RGB
-    img_recolored = cv2.cvtColor(img_recolored, cv2.COLOR_LAB2RGB)
-
-    # Apply Poisson blending for natural-looking color transition
-    center = (img_np.shape[1]//2, img_np.shape[0]//2)
-    img_final = cv2.seamlessClone(img_recolored, img_np, mask, center, cv2.NORMAL_CLONE)
+    # Preserve texture by modifying A & B channels with blend factor
+    blend_factor = 0.6  # Controls intensity of color change
+    img_lab[..., 1] = np.where(mask > 128, img_lab[..., 1] * (1 - blend_factor) + new_color_lab[1] * blend_factor, img_lab[..., 1])
+    img_lab[..., 2] = np.where(mask > 128, img_lab[..., 2] * (1 - blend_factor) + new_color_lab[2] * blend_factor, img_lab[..., 2])
 
-    return Image.fromarray(img_final)
+    # Convert back to RGB
+    img_recolored = cv2.cvtColor(img_lab, cv2.COLOR_LAB2RGB)
+    
+    return Image.fromarray(img_recolored)
 
 # Gradio Interface
 demo = gr.Interface(
     fn=change_dress_color,
     inputs=[
         gr.Image(type="filepath", label="Upload Dress Image"),
-        gr.Radio(["Red", "Blue", "Green", "Yellow", "Purple"], label="Choose New Dress Color")
+        gr.Radio(["Red", "Blue", "Green", "Yellow", "Purple", "Orange", "Cyan", "Magenta", "White"], label="Choose New Dress Color")
     ],
     outputs=gr.Image(type="pil", label="Color Changed Dress"),
     title="Dress Color Changer",