Update app.py

app.py

@@ -32,7 +32,7 @@ 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)  # Thresholding for binary mask
+    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
     return mask
 
@@ -48,25 +48,27 @@ def change_dress_color(image_path, color):
     if mask is None:
         return img  # No dress detected
     
-    # Convert the selected color to HSV
+    # Convert the selected color to BGR
     color_map = {
-        "Red": (0, 255, 255),
-        "Blue": (120, 255, 255),
-        "Green": (60, 255, 255),
-        "Yellow": (30, 255, 255),
-        "Purple": (150, 255, 255)
+        "Red": (0, 0, 255),
+        "Blue": (255, 0, 0),
+        "Green": (0, 255, 0),
+        "Yellow": (0, 255, 255),
+        "Purple": (128, 0, 128)
     }
-    hsv_color = np.uint8([[color_map.get(color, (0, 255, 255))]])  # Default to Red
+    new_color_bgr = np.array(color_map.get(color, (0, 0, 255)), dtype=np.uint8)  # Default to Red
     
-    # Convert to BGR
-    new_color_bgr = cv2.cvtColor(hsv_color, cv2.COLOR_HSV2BGR)[0][0]
-
-    # Apply the color change
+    # Convert image to HSV
     img_hsv = cv2.cvtColor(img_np, cv2.COLOR_RGB2HSV)
-    img_hsv[..., 0] = mask * new_color_bgr[0] + (1 - mask) * img_hsv[..., 0]  # Adjust hue
-    img_hsv[..., 1] = mask * new_color_bgr[1] + (1 - mask) * img_hsv[..., 1]  # Adjust saturation
-    img_hsv[..., 2] = mask * new_color_bgr[2] + (1 - mask) * img_hsv[..., 2]  # Adjust value
 
+    # Convert new color to HSV
+    new_color_hsv = cv2.cvtColor(np.uint8([[new_color_bgr]]), cv2.COLOR_BGR2HSV)[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
+
+    # Convert back to RGB
     img_recolored = cv2.cvtColor(img_hsv, cv2.COLOR_HSV2RGB)
     
     return Image.fromarray(img_recolored)