Replace ISNet with SAM3 for garment segmentation (remove hardcoded HF token)

app.py

@@ -79,19 +79,41 @@ _hrnet = models.pose_hrnet.get_pose_net(cfg, is_train=False)
 _hrnet.load_state_dict(torch.load(MODEL_FILE, map_location="cpu"), strict=False)
 _hrnet.eval()
 
-# Load rembg ISNet session (CPU-only; ONNX Runtime)
-import rembg
-_rembg_session = rembg.new_session(model_name="isnet-general-use")
+# Load SAM3 model (CPU — GPU assigned dynamically via @spaces.GPU)
+from huggingface_hub import hf_hub_download
+
+_sam3_model = None
+_sam3_processor = None
+
+def init_sam3():
+    global _sam3_model, _sam3_processor
+    try:
+        from sam3.model_builder import build_sam3_image_model
+        from sam3.model.sam3_image_processor import Sam3Processor
+        
+        token = os.getenv("HF_TOKEN")
+        if not token:
+            print("WARNING: HF_TOKEN environment variable not set. Downloading SAM3 may fail.")
+        print("Downloading/loading SAM3 model...", flush=True)
+        model_path = hf_hub_download(repo_id="facebook/sam3", filename="sam3.pt", token=token)
+        _sam3_model = build_sam3_image_model(checkpoint_path=model_path)
+        _sam3_model.eval()
+        _sam3_processor = Sam3Processor(_sam3_model)
+        print("SAM3 loaded successfully.", flush=True)
+    except Exception as e:
+        print(f"Failed to load SAM3: {e}")
+
+init_sam3()
 
 # ---------------------------------------------------------------------------
 # SINGLE @spaces.GPU function — runs both ISNet BG removal and HRNet
 # ---------------------------------------------------------------------------
 @spaces.GPU
-def run_pipeline_gpu(img_pil: Image.Image):
+def run_pipeline_gpu(img_pil: Image.Image, prompt_text="pants"):
     """
     Runs the full GPU pipeline:
-      1. ISNet background removal  (via rembg / ONNX Runtime — uses GPU if available)
-      2. HRNet landmark estimation (PyTorch — uses GPU)
+      1. SAM3 garment segmentation (uses GPU if available)
+      2. HRNet landmark estimation (uses GPU)
 
     Returns:
         img_np      : numpy HWC image array (for drawing)
@@ -99,35 +121,51 @@ def run_pipeline_gpu(img_pil: Image.Image):
         heatmaps    : numpy array of shape (num_joints, H, W)
         img_h, img_w: original dimensions used for landmark scaling
     """
-    global _hrnet, _rembg_session
+    global _hrnet, _sam3_model, _sam3_processor
 
     t_start = time.time()
 
-    # -- Step 1: ISNet background removal --
+    # -- Step 1: SAM3 garment segmentation --
     t1 = time.time()
-    REMBG_MAX_DIM = 800
-    rembg_input = img_pil
-    _rw, _rh = img_pil.size
-    if max(_rw, _rh) > REMBG_MAX_DIM:
-        _rs = REMBG_MAX_DIM / max(_rw, _rh)
-        rembg_input = img_pil.resize((int(_rw * _rs), int(_rh * _rs)), Image.BICUBIC)
-
-    removed_bg = rembg.remove(rembg_input, session=_rembg_session)
-
-    if removed_bg.mode == "RGBA":
-        white_bg = Image.new("RGBA", removed_bg.size, "WHITE")
-        white_bg.paste(removed_bg, mask=removed_bg)
-        img_rgb = white_bg.convert("RGB")
-        alpha = np.array(removed_bg.split()[-1])
-        mask_np = (alpha >= 128).astype(np.uint8) * 255
+    
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    if _sam3_model is not None and _sam3_processor is not None:
+        _sam3_model.to(device)
+        with torch.inference_mode():
+            inference_state = _sam3_processor.set_image(img_pil)
+            output = _sam3_processor.set_text_prompt(state=inference_state, prompt=prompt_text)
+            
+            masks = output["masks"]
+            scores = output["scores"]
+            
+            # Use the mask with the highest score
+            best_idx = 0
+            if hasattr(scores, "shape") and len(scores.shape) > 0 and scores.shape[0] > 0:
+                # Assuming scores is a tensor or numpy array
+                if torch.is_tensor(scores):
+                    best_idx = torch.argmax(scores).item()
+                else:
+                    best_idx = np.argmax(scores)
+            
+            best_mask = masks[best_idx]
+            if torch.is_tensor(best_mask):
+                best_mask = best_mask.cpu().numpy()
+            
+            mask_np = (best_mask > 0).astype(np.uint8) * 255
     else:
-        img_rgb = removed_bg.convert("RGB")
-        gray = cv2.cvtColor(np.array(img_rgb), cv2.COLOR_RGB2GRAY)
-        _, mask_np = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY_INV)
-
+        print("SAM3 model not initialized, falling back to empty mask.")
+        mask_np = np.zeros((img_pil.size[1], img_pil.size[0]), dtype=np.uint8)
+        
     isnet_time = time.time() - t1
-    print(f"ISNet: {isnet_time:.2f}s", flush=True)
-
+    print(f"SAM3: {isnet_time:.2f}s", flush=True)
+
+    # Generate transparent cutout / white background
+    img_rgb = img_pil.convert("RGB")
+    white_bg = Image.new("RGB", img_pil.size, "WHITE")
+    mask_pil = Image.fromarray(mask_np)
+    white_bg.paste(img_rgb, mask=mask_pil)
+    
     # Portrait orientation helpers
     def _portrait(arr_img):
         h, w = arr_img.shape[:2]
@@ -135,7 +173,7 @@ def run_pipeline_gpu(img_pil: Image.Image):
             arr_img = cv2.rotate(arr_img, cv2.ROTATE_90_COUNTERCLOCKWISE)
         return arr_img
 
-    img_np = _portrait(np.array(img_rgb.convert("RGB")))
+    img_np = _portrait(np.array(white_bg))
     binary_mask = _portrait(mask_np)
 
     # -- Step 2: HRNet landmark estimation --
@@ -237,9 +275,9 @@ def verify_user_from_token(token):
 # ---------------------------------------------------------------------------
 def process_fit(name, waist_in, input_image, request: gr.Request):
     if not name:
-        return None, "Please enter your name.", "", ""
+        return None, None, "Please enter your name.", "", ""
     if input_image is None:
-        return None, "Please upload an image.", "", ""
+        return None, None, "Please upload an image.", "", ""
 
     try:
         # Auth
@@ -270,8 +308,8 @@ def process_fit(name, waist_in, input_image, request: gr.Request):
                 cv2.COLOR_BGR2RGB
             ))
 
-        # Run GPU pipeline (ISNet + HRNet in one @spaces.GPU call)
-        img_np, binary_mask, heatmaps, H, W = run_pipeline_gpu(img_pil)
+        # Run GPU pipeline (SAM3 + HRNet in one @spaces.GPU call)
+        img_np, binary_mask, heatmaps, H, W = run_pipeline_gpu(img_pil, prompt_text="pants")
 
         # Decode + refine landmarks (CPU)
         pts = decode_landmarks(heatmaps, H, W)
@@ -296,6 +334,7 @@ def process_fit(name, waist_in, input_image, request: gr.Request):
         inseam_val   = round_half(inseam_px  / waist_px * half_waist)
 
         result_pil = Image.fromarray(result_img)
+        binary_mask_pil = Image.fromarray(binary_mask)
 
         # Cloud save (if logged in)
         sync_status = "Not logged in — scan not saved to cloud profile."
@@ -345,11 +384,11 @@ def process_fit(name, waist_in, input_image, request: gr.Request):
             </div>
         </div>"""
 
-        return result_pil, results_html, post_msg_html, sync_status
+        return result_pil, binary_mask_pil, results_html, post_msg_html, sync_status
 
     except Exception as e:
         import traceback
-        return None, f"<div style='color:red;'>Error: {e}</div>", "", traceback.format_exc()
+        return None, None, f"<div style='color:red;'>Error: {e}</div>", "", traceback.format_exc()
 
 
 # ---------------------------------------------------------------------------
@@ -371,6 +410,7 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue="emerald", secondary_hue="ind
 
         with gr.Column(scale=1):
             image_output  = gr.Image(label="AI Detection Result")
+            mask_output   = gr.Image(label="Garment Mask", image_mode="L")
             results_panel = gr.HTML(label="Results")
             status_panel  = gr.Textbox(label="Sync Status", interactive=False)
             post_msg_panel = gr.HTML()
@@ -378,7 +418,7 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue="emerald", secondary_hue="ind
     analyze_btn.click(
         fn=process_fit,
         inputs=[name_input, waist_input, image_input],
-        outputs=[image_output, results_panel, post_msg_panel, status_panel]
+        outputs=[image_output, mask_output, results_panel, post_msg_panel, status_panel]
     )
 
 demo.launch()

requirements.txt

@@ -16,8 +16,7 @@ PyYAML
 scipy
 scikit-image
 certifi
-rembg>=2.0.57
-onnxruntime
+git+https://github.com/facebookresearch/sam3.git
 psutil
 fastapi
 uvicorn