Pipeline rewrite: SAM 3 boxes + rembg background removal + 4x Lanczos upscale

app.py

@@ -6,6 +6,7 @@ from PIL import Image
 import os
 import io
 import fitz  # PyMuPDF
+from rembg import remove as rembg_remove, new_session as rembg_new_session
 
 # ── UNCONDITIONAL BFloat16 → Float16 Patch for T4 Turing GPUs ────
 # CRITICAL: torch.cuda.is_bf16_supported() returns True on T4 because CUDA
@@ -139,31 +140,22 @@ ASSET_LIBRARY_DIR = os.path.join(tempfile.gettempdir(), "sam3_library")
 os.makedirs(ASSET_LIBRARY_DIR, exist_ok=True)
 asset_counter = 0
 
-def mask_iou(m1: np.ndarray, m2: np.ndarray) -> float:
-    # Threshold uint8 alpha arrays to binary for IoU comparison
-    b1 = m1 > 128 if m1.dtype == np.uint8 else m1
-    b2 = m2 > 128 if m2.dtype == np.uint8 else m2
-    inter = np.logical_and(b1, b2).sum()
-    union = np.logical_or(b1, b2).sum()
-    return float(inter) / float(union) if union > 0 else 0.0
+# Initialize rembg session (downloads model once, reuses for all requests)
+print("Loading background removal model...", flush=True)
+rembg_session = rembg_new_session(model_name="isnet-general-use")
+print("Background removal model loaded.", flush=True)
 
-def mask_to_crop(orig_rgb: np.ndarray, alpha: np.ndarray) -> np.ndarray:
-    """Crop an RGBA image using a binary alpha mask."""
-    h, w = orig_rgb.shape[:2]
-    
-    rgba = np.zeros((h, w, 4), dtype=np.uint8)
-    rgba[:, :, :3] = orig_rgb
-    rgba[:, :, 3] = alpha
-    
-    ys, xs = np.nonzero(alpha > 0)
-    if len(ys) == 0:
-        return rgba
-    y0, y1 = int(ys.min()), int(ys.max())
-    x0, x1 = int(xs.min()), int(xs.max())
-    pad = 4
-    y0, x0 = max(0, y0 - pad), max(0, x0 - pad)
-    y1, x1 = min(h - 1, y1 + pad), min(w - 1, x1 + pad)
-    return rgba[y0:y1+1, x0:x1+1]
+def box_iou(b1, b2):
+    """IoU between two boxes [x0, y0, x1, y1]."""
+    x0 = max(b1[0], b2[0])
+    y0 = max(b1[1], b2[1])
+    x1 = min(b1[2], b2[2])
+    y1 = min(b1[3], b2[3])
+    inter = max(0, x1 - x0) * max(0, y1 - y0)
+    a1 = (b1[2] - b1[0]) * (b1[3] - b1[1])
+    a2 = (b2[2] - b2[0]) * (b2[3] - b2[1])
+    union = a1 + a2 - inter
+    return inter / union if union > 0 else 0.0
 
 def upscale_4x(rgba: np.ndarray) -> np.ndarray:
     """4x Lanczos upscale with unsharp masking for crisp graphics."""
@@ -173,7 +165,7 @@ def upscale_4x(rgba: np.ndarray) -> np.ndarray:
     # Upscale full RGBA together with Lanczos4
     upscaled = cv2.resize(rgba, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
     
-    # Unsharp mask on RGB only (preserve smooth alpha)
+    # Unsharp mask on RGB only (preserve alpha from rembg)
     rgb = upscaled[:, :, :3]
     blurred = cv2.GaussianBlur(rgb, (0, 0), sigmaX=1.0)
     rgb_sharp = cv2.addWeighted(rgb, 1.5, blurred, -0.5, 0)
@@ -198,7 +190,7 @@ def extract_assets(input_image):
         print(f">>> Image size: {w}x{h}, area: {img_area}", flush=True)
         pil_img = Image.fromarray(orig_rgb)
         
-        all_masks = []
+        all_boxes = []
         all_scores = []
         
         with torch.inference_mode():
@@ -210,95 +202,102 @@ def extract_assets(input_image):
                 print(f">>> Processing concept: '{concept}'...", flush=True)
                 out = processor.set_text_prompt(state=state, prompt=concept)
                 
-                # Log all output keys to find raw logits
-                print(f"  [{concept}] Output keys: {list(out.keys())}", flush=True)
-                
                 masks = out["masks"]
                 scores = out["scores"]
-                
-                # Check for raw logits (sub-pixel smooth edges)
-                raw_logits = None
-                for logit_key in ["masks_logits", "low_res_masks", "logits", "mask_logits"]:
-                    val = out.get(logit_key)
-                    if val is not None and (not torch.is_tensor(val) or val.numel() > 0):
-                        raw_logits = val
-                        break
-                if raw_logits is not None:
-                    print(f"  [{concept}] RAW LOGITS found: shape={raw_logits.shape}, dtype={raw_logits.dtype}", flush=True)
+                boxes = out.get("boxes")
                 
                 if masks is None or len(masks) == 0:
-                    print(f"  [{concept}] No masks returned", flush=True)
+                    print(f"  [{concept}] No detections", flush=True)
                     continue
-                    
+                
                 if torch.is_tensor(masks): masks = masks.float().cpu().numpy()
                 if torch.is_tensor(scores): scores = scores.float().cpu().numpy()
-                if raw_logits is not None and torch.is_tensor(raw_logits):
-                    raw_logits = raw_logits.float().cpu().numpy()
+                if boxes is not None and torch.is_tensor(boxes): boxes = boxes.float().cpu().numpy()
                 
-                print(f"  [{concept}] MASK shape: {masks.shape}, SCORES max: {scores.max():.4f}, min: {scores.min():.4f}", flush=True)
+                print(f"  [{concept}] Found {len(masks)} detections, boxes: {boxes.shape if boxes is not None else 'None'}", flush=True)
                 
                 for j in range(len(masks)):
-                    m = masks[j]
-                    while m.ndim > 2: m = m[0]
-                    m_bool = m.astype(bool)
-                    
                     score = float(scores[j]) if scores.ndim > 0 else float(scores)
-                    area = m_bool.sum()
                     
-                    if score < 0.1 or area < 500 or area > img_area * 0.90:
-                        print(f"    mask[{j}] SKIPPED: score={score:.4f}, area={area}", flush=True)
+                    # Get bounding box: prefer SAM 3's boxes, fallback to mask bbox
+                    if boxes is not None and j < len(boxes):
+                        box = boxes[j].flatten()
+                        # SAM 3 boxes format: get x0,y0,x1,y1
+                        if len(box) >= 4:
+                            x0, y0, x1, y1 = int(box[0]), int(box[1]), int(box[2]), int(box[3])
+                        else:
+                            continue
+                    else:
+                        # Fallback: derive box from mask
+                        m = masks[j]
+                        while m.ndim > 2: m = m[0]
+                        ys, xs = np.nonzero(m > 0.5)
+                        if len(ys) == 0: continue
+                        x0, y0, x1, y1 = int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())
+                    
+                    # Validate box
+                    box_w = x1 - x0
+                    box_h = y1 - y0
+                    box_area = box_w * box_h
+                    
+                    if score < 0.1 or box_area < 500 or box_area > img_area * 0.90:
+                        print(f"    det[{j}] SKIPPED: score={score:.4f}, area={box_area}", flush=True)
                         continue
                     
-                    # Build alpha: use logits for smooth edges if available
-                    if raw_logits is not None and j < len(raw_logits):
-                        logit = raw_logits[j]
-                        while logit.ndim > 2: logit = logit[0]
-                        alpha_smooth = 1.0 / (1.0 + np.exp(-logit.astype(np.float32)))
-                        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
-                        dilated = cv2.dilate(m_bool.astype(np.uint8), kernel, iterations=2)
-                        alpha_smooth = alpha_smooth * dilated
-                        alpha_uint8 = (alpha_smooth * 255).clip(0, 255).astype(np.uint8)
-                        alpha_mask = np.where(m_bool, np.uint8(255), alpha_uint8)
-                    else:
-                        alpha_mask = m_bool.astype(np.uint8) * 255
+                    # Add padding (10% of box size)
+                    pad_x = max(10, int(box_w * 0.10))
+                    pad_y = max(10, int(box_h * 0.10))
+                    bx0 = max(0, x0 - pad_x)
+                    by0 = max(0, y0 - pad_y)
+                    bx1 = min(w, x1 + pad_x)
+                    by1 = min(h, y1 + pad_y)
                     
-                    all_masks.append(alpha_mask)
+                    all_boxes.append([bx0, by0, bx1, by1])
                     all_scores.append(score)
-                    print(f"    mask[{j}] KEPT: score={score:.4f}, area={area}", flush=True)
-
-        print(f">>> Total masks kept: {len(all_masks)}", flush=True)
+                    print(f"    det[{j}] KEPT: score={score:.4f}, box=[{bx0},{by0},{bx1},{by1}]", flush=True)
+        
+        print(f">>> Total detections kept: {len(all_boxes)}", flush=True)
         
-        if not all_masks:
+        if not all_boxes:
             gr.Info("No assets found in this image. Try a different slide with more visual elements.")
-            print(">>> No masks passed filters, returning []", flush=True)
+            print(">>> No detections passed filters, returning []", flush=True)
             return []
-            
-        # Deduplicate
-        order = sorted(range(len(all_masks)), key=lambda i: all_scores[i], reverse=True)
+        
+        # Deduplicate by box IoU
+        order = sorted(range(len(all_boxes)), key=lambda i: all_scores[i], reverse=True)
         keep = []
         for i in order:
             dup = False
             for ki in keep:
-                if mask_iou(all_masks[i], all_masks[ki]) > 0.5:
+                if box_iou(all_boxes[i], all_boxes[ki]) > 0.5:
                     dup = True
                     break
             if not dup:
                 keep.append(i)
-                
-        # Crop to PNG files and save to library
+        
+        # For each kept box: crop → rembg → upscale → save
         results = []
         global asset_counter
         for idx, ki in enumerate(keep):
-            crop = mask_to_crop(orig_rgb, all_masks[ki])
-            # 4x upscale with Lanczos + sharpening
-            crop = upscale_4x(crop)
-            # Save to persistent library
+            bx0, by0, bx1, by1 = all_boxes[ki]
+            crop_rgb = orig_rgb[by0:by1, bx0:bx1]
+            
+            # Background removal with rembg (clean alpha matte)
+            crop_pil = Image.fromarray(crop_rgb)
+            rgba_pil = rembg_remove(crop_pil, session=rembg_session)
+            rgba_np = np.array(rgba_pil)
+            print(f"    crop[{idx}] rembg done: {rgba_np.shape}", flush=True)
+            
+            # 4x upscale
+            rgba_np = upscale_4x(rgba_np)
+            
+            # Save to library
             asset_counter += 1
             lib_path = os.path.join(ASSET_LIBRARY_DIR, f"asset_{asset_counter:04d}.png")
-            Image.fromarray(crop, "RGBA").save(lib_path, format="PNG")
+            Image.fromarray(rgba_np, "RGBA").save(lib_path, format="PNG")
             results.append(lib_path)
         
-        print(f">>> Returning {len(results)} cropped PNG assets (library total: {asset_counter})", flush=True)
+        print(f">>> Returning {len(results)} assets (library total: {asset_counter})", flush=True)
         return results
         
     except Exception as e:

requirements.txt

@@ -11,3 +11,4 @@ gradio
 git+https://github.com/facebookresearch/sam3.git
 opencv-python-headless
 PyMuPDF
+rembg