Update app.py

app.py

@@ -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 the mask to reduce noise
+    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.dilate(mask, kernel, iterations=1)  # Expand the detected dress area
+    mask = cv2.medianBlur(mask, 5)  # Reduce noise and smooth edges
     
     return mask
 
@@ -66,17 +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]
     
-    # 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
-
+    # 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_lab, cv2.COLOR_LAB2RGB)
+    img_recolored = cv2.cvtColor(img_recolored, cv2.COLOR_LAB2RGB)
 
-    # 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-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)
 
-    return Image.fromarray(img_recolored)
+    return Image.fromarray(img_final)
 
 # Gradio Interface
 demo = gr.Interface(