BG Removal Updated

app.py

@@ -11,6 +11,7 @@ from io import BytesIO
 import base64
 import traceback
 from starlette.exceptions import HTTPException as StarletteHTTPException
+from bgrem import remove as bgrem_remove  # NEW: bgrem import
 
 # ===============================
 # SIMPLE DPT MODEL (DEPTH ESTIMATION)
@@ -85,17 +86,6 @@ def _process_saree_core(base_image: Image.Image, pattern_image: Image.Image):
     l_clahe = clahe.apply(l_channel)
     shading_map = l_clahe / 255.0
 
-    # GrabCut mask
-    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
-    grabcut_mask = np.zeros(img_bgr.shape[:2], np.uint8)
-    height, width = img_bgr.shape[:2]
-    margin = int(min(width, height) * 0.05)
-    rect = (margin, margin, width - 2 * margin, height - 2 * margin)
-    bgdModel = np.zeros((1, 65), np.float64)
-    fgdModel = np.zeros((1, 65), np.float64)
-    cv2.grabCut(img_bgr, grabcut_mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_RECT)
-    mask = np.where((grabcut_mask == cv2.GC_FGD) | (grabcut_mask == cv2.GC_PR_FGD), 255, 0).astype(np.uint8)
-
     # Tile pattern
     pattern_np = np.array(pattern_image.convert("RGB"))
     target_h, target_w = img_np.shape[:2]
@@ -119,13 +109,20 @@ def _process_saree_core(base_image: Image.Image, pattern_image: Image.Image):
     pattern_folded *= normal_boost
     pattern_folded = np.clip(pattern_folded, 0, 1)
 
-    # Clean mask and feather edges
-    mask_float = mask.astype(np.float32) / 255.0
-    kernel = np.ones((3, 3), np.uint8)
-    mask_clean = cv2.morphologyEx((mask_float * 255).astype(np.uint8), cv2.MORPH_OPEN, kernel)
-    mask_clean = cv2.morphologyEx(mask_clean, cv2.MORPH_CLOSE, kernel)
-    mask_clean = cv2.dilate(mask_clean, kernel, iterations=1)
-    mask_blurred = cv2.GaussianBlur(mask_clean, (15, 15), sigmaX=5, sigmaY=5)
+    # ==========================================================
+    # NEW: Background removal using bgrem (instead of GrabCut)
+    # ==========================================================
+    buf = BytesIO()
+    base_image.save(buf, format="PNG")
+    base_bytes = buf.getvalue()
+
+    # Get RGBA from bgrem
+    result_no_bg = bgrem_remove(base_bytes)
+    mask_img = Image.open(BytesIO(result_no_bg)).convert("RGBA")
+    mask_alpha = np.array(mask_img)[:, :, 3].astype(np.float32) / 255.0
+
+    # Feather edges for smoother blending
+    mask_blurred = cv2.GaussianBlur((mask_alpha * 255).astype(np.uint8), (15, 15), sigmaX=5, sigmaY=5)
     mask_blurred[mask_blurred < 25] = 0
     mask_blurred = mask_blurred.astype(np.float32) / 255.0