Update app.py

app.py

@@ -34,6 +34,7 @@ def segment_dress(image_np):
     with torch.no_grad():
         output = model(input_tensor)[0][0].squeeze().cpu().numpy()
 
+    output = (output - output.min()) / (output.max() - output.min() + 1e-8)  # Normalize to [0, 1]
     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_LINEAR)
 
@@ -44,12 +45,16 @@ def apply_grabcut(image_np, dress_mask):
     bgd_model = np.zeros((1, 65), np.float64)
     fgd_model = np.zeros((1, 65), np.float64)
 
-    mask = np.where(dress_mask > 0, cv2.GC_FGD, cv2.GC_BGD).astype('uint8')
-    rect = (10, 10, image_np.shape[1] - 10, image_np.shape[0] - 10)
+    mask = np.where(dress_mask > 0, cv2.GC_PR_FGD, cv2.GC_BGD).astype('uint8')
+
+    # Get bounding box of the mask
+    coords = cv2.findNonZero(dress_mask)
+    x, y, w, h = cv2.boundingRect(coords)
+    rect = (x, y, w, h)
 
     cv2.grabCut(image_np, mask, rect, bgd_model, fgd_model, 5, cv2.GC_INIT_WITH_MASK)
 
-    refined_mask = np.where((mask == 2) | (mask == 0), 0, 255).astype("uint8")
+    refined_mask = np.where((mask == cv2.GC_FGD) | (mask == cv2.GC_PR_FGD), 255, 0).astype("uint8")
     return refine_mask(refined_mask)
 
 def recolor_dress(image_np, dress_mask, target_color):