Update app.py

app.py

@@ -28,7 +28,7 @@ def segment_dress(image_np):
         output = model(input_tensor)[0][0].squeeze().cpu().numpy()
 
     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_NEAREST)
+    dress_mask = cv2.resize(dress_mask, (image_np.shape[1], image_np.shape[0]), interpolation=cv2.INTER_LINEAR)
 
     return dress_mask
 
@@ -43,19 +43,26 @@ def recolor_dress(image_np, dress_mask, target_color):
     
     # Compute mean LAB values of dress pixels
     dress_pixels = img_lab[dress_mask > 0]
-    mean_L, mean_A, mean_B = dress_pixels[:, 0].mean(), dress_pixels[:, 1].mean(), dress_pixels[:, 2].mean()
+    if len(dress_pixels) == 0:
+        return image_np  # No dress detected
 
-    # Compute new color adjustment
+    mean_L, mean_A, mean_B = np.mean(dress_pixels, axis=0)
+
+    # Adjust LAB channels for the new color
     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)
 
-    # Smooth edges for natural blending
-    img_recolored = cv2.seamlessClone(img_recolored, image_np, dress_mask, (image_np.shape[1]//2, image_np.shape[0]//2), cv2.NORMAL_CLONE)
+    # Convert mask to 3 channels for seamless blending
+    mask_bgr = cv2.cvtColor(dress_mask, cv2.COLOR_GRAY2BGR)
+
+    # 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)
 
-    return img_recolored
+    return blended
 
 def change_dress_color(image_path, color):
     """Main function to change dress color naturally"""